scirs2_stats/bayesian_advanced/
bnn_train.rs

1//! Real training and posterior-predictive inference for
2//! [`super::BayesianNeuralNetwork`].
3//!
4//! `fit` trains a small **deep ensemble** (Lakshminarayanan et al., 2017):
5//! each ensemble member is initialized by sampling its weights and biases
6//! from the network's declared priors (via a real RNG, not a fixed
7//! constant), independently bootstrap-resampled from the training data, and
8//! then trained by ordinary gradient descent using an exact, hand-derived
9//! backpropagation pass (`forward_with_cache` / `backward`) through whatever
10//! mix of `ReLU`/`Sigmoid`/`Tanh`/`Swish`/`GELU` activations the network was
11//! built with. Random initialization plus bootstrap resampling gives each
12//! ensemble member a genuinely different local optimum, and the resulting
13//! ensemble's empirical mean/variance is a real, well-established
14//! (if approximate) way to quantify a neural network's epistemic uncertainty.
15//!
16//! `predict_with_uncertainty` then does real forward passes through that
17//! trained ensemble (or, if `fit` has not been called, through independent
18//! prior draws of the whole network -- a real prior-predictive Monte Carlo
19//! estimate) and reports the empirical mean/variance across the draws.
20
21use super::{ActivationType, AdvancedBayesianFloat, BayesianNeuralNetwork, DistributionType};
22use crate::error::{StatsError, StatsResult};
23use scirs2_core::ndarray::{Array1, Array2, ArrayView2, Axis};
24use scirs2_core::random::{Rng, RngExt, SeedableRng};
25use scirs2_core::validation::checkarray_finite;
26
27/// Configuration for [`BayesianNeuralNetwork::fit`]'s deep-ensemble training.
28#[derive(Debug, Clone)]
29pub struct BnnTrainingConfig {
30    /// Number of independently-initialized-and-trained ensemble members.
31    pub n_ensemble: usize,
32    /// Number of full-batch gradient descent epochs per ensemble member.
33    pub epochs: usize,
34    /// Gradient descent learning rate.
35    pub learning_rate: f64,
36    /// Whether to train each ensemble member on an independent bootstrap
37    /// resample of the training data (in addition to independent random
38    /// weight initialization).
39    pub bootstrap: bool,
40    /// Optional RNG seed for reproducibility.
41    pub seed: Option<u64>,
42}
43
44impl Default for BnnTrainingConfig {
45    fn default() -> Self {
46        Self {
47            n_ensemble: 16,
48            epochs: 200,
49            learning_rate: 0.05,
50            bootstrap: true,
51            seed: None,
52        }
53    }
54}
55
56/// Sample a real `N(mean, 1/precision)` draw using `rng` (replaces the old
57/// hardcoded-`u1 = u2 = 0.5` "Box-Muller transform" that always produced the
58/// same value regardless of `rng` state).
59fn sample_normal<F: AdvancedBayesianFloat, R: Rng + ?Sized>(
60    mean: F,
61    precision: F,
62    rng: &mut R,
63) -> F {
64    use scirs2_core::random::{Distribution, StandardNormal};
65    let eps = F::from(1e-12).expect("1e-12 fits in any Float");
66    let std_dev = F::one() / precision.max(eps).sqrt();
67    let z64: f64 = StandardNormal.sample(rng);
68    let z = F::from(z64).unwrap_or(F::zero());
69    mean + std_dev * z
70}
71
72fn prior_precision<F: AdvancedBayesianFloat>(prior: &DistributionType<F>) -> F {
73    match prior {
74        DistributionType::Normal { precision, .. } => *precision,
75        _ => F::one(),
76    }
77}
78
79impl<F: AdvancedBayesianFloat> BayesianNeuralNetwork<F> {
80    /// Derivative of the layer activation function with respect to its
81    /// pre-activation input `z` (i.e. `d/dz apply_activation(z, activation)`).
82    fn activation_derivative(&self, z: F, activation: ActivationType) -> F {
83        match activation {
84            ActivationType::ReLU => {
85                if z > F::zero() {
86                    F::one()
87                } else {
88                    F::zero()
89                }
90            }
91            ActivationType::Sigmoid => {
92                let s = F::one() / (F::one() + (-z).exp());
93                s * (F::one() - s)
94            }
95            ActivationType::Tanh => {
96                let t = z.tanh();
97                F::one() - t * t
98            }
99            ActivationType::Swish => {
100                let s = F::one() / (F::one() + (-z).exp());
101                s + z * s * (F::one() - s)
102            }
103            ActivationType::GELU => {
104                let sqrt_2_pi = F::from(0.7978845608).expect("sqrt(2/pi) fits in any Float");
105                let coeff = F::from(0.044715).expect("0.044715 fits in any Float");
106                let three = F::from(3.0).expect("3.0 fits in any Float");
107                let half = F::from(0.5).expect("0.5 fits in any Float");
108                let g = sqrt_2_pi * (z + coeff * z * z * z);
109                let g_prime = sqrt_2_pi * (F::one() + three * coeff * z * z);
110                let tanh_g = g.tanh();
111                half * (F::one() + tanh_g) + half * z * (F::one() - tanh_g * tanh_g) * g_prime
112            }
113        }
114    }
115
116    /// Forward pass that also returns every layer's pre-activation (`zs`)
117    /// and post-activation (`acts`, with `acts[0] == x`) values, needed by
118    /// [`Self::backward`].
119    fn forward_with_cache(
120        &self,
121        x: &ArrayView2<F>,
122        weights: &[Array2<F>],
123        biases: &[Array1<F>],
124    ) -> StatsResult<(Vec<Array2<F>>, Vec<Array2<F>>)> {
125        let mut acts = vec![x.to_owned()];
126        let mut zs = Vec::with_capacity(self.activations.len());
127        for (layer_idx, &activation_type) in self.activations.iter().enumerate() {
128            let z = self.linear_transform(
129                &acts[layer_idx].view(),
130                &weights[layer_idx],
131                &biases[layer_idx],
132            )?;
133            let a = z.mapv(|val| self.apply_activation(val, activation_type));
134            zs.push(z);
135            acts.push(a);
136        }
137        Ok((zs, acts))
138    }
139
140    /// Exact backpropagation of the mean-squared-error loss between the
141    /// network's (post-activation) output and `y`, returning the gradient of
142    /// each layer's weight matrix and bias vector.
143    fn backward(
144        &self,
145        x: &ArrayView2<F>,
146        y: &ArrayView2<F>,
147        weights: &[Array2<F>],
148        biases: &[Array1<F>],
149    ) -> StatsResult<(Vec<Array2<F>>, Vec<Array1<F>>)> {
150        let (zs, acts) = self.forward_with_cache(x, weights, biases)?;
151        let n_layers = self.activations.len();
152        let n_samples_ = x.nrows();
153        let output_dim = *self.architecture.last().ok_or_else(|| {
154            StatsError::InvalidArgument("architecture must be non-empty".to_string())
155        })?;
156
157        let scale = F::from(2.0).expect("2.0 fits in any Float")
158            / F::from((n_samples_ * output_dim).max(1)).expect("count fits in any Float");
159
160        let y_pred = &acts[n_layers];
161        let mut delta_a = (y_pred - y).mapv(|v| v * scale);
162
163        let mut grads_w: Vec<Array2<F>> = (0..n_layers).map(|_| Array2::zeros((0, 0))).collect();
164        let mut grads_b: Vec<Array1<F>> = (0..n_layers).map(|_| Array1::zeros(0)).collect();
165
166        for l in (0..n_layers).rev() {
167            let z_l = &zs[l];
168            let mut delta_z = Array2::<F>::zeros(delta_a.raw_dim());
169            for (dz, (da, zv)) in delta_z.iter_mut().zip(delta_a.iter().zip(z_l.iter())) {
170                *dz = *da * self.activation_derivative(*zv, self.activations[l]);
171            }
172
173            let a_l = &acts[l];
174            grads_w[l] = a_l.t().dot(&delta_z);
175            grads_b[l] = delta_z.sum_axis(Axis(0));
176
177            if l > 0 {
178                delta_a = delta_z.dot(&weights[l].t());
179            }
180        }
181
182        Ok((grads_w, grads_b))
183    }
184
185    /// Train a deep ensemble of posterior weight/bias samples: each member is
186    /// randomly initialized from the network's priors, optionally
187    /// bootstrap-resampled, and trained to a local optimum of the mean
188    /// squared error via exact backpropagation + gradient descent. Populates
189    /// `self.weight_samples` / `self.bias_samples`, which
190    /// [`Self::predict_with_uncertainty`] then draws from.
191    pub fn fit(
192        &mut self,
193        x: &ArrayView2<F>,
194        y: &ArrayView2<F>,
195        config: &BnnTrainingConfig,
196    ) -> StatsResult<()> {
197        checkarray_finite(x, "x")?;
198        checkarray_finite(y, "y")?;
199        if x.nrows() != y.nrows() {
200            return Err(StatsError::DimensionMismatch(
201                "x and y must have the same number of rows".to_string(),
202            ));
203        }
204        let output_dim = *self.architecture.last().ok_or_else(|| {
205            StatsError::InvalidArgument("architecture must be non-empty".to_string())
206        })?;
207        if y.ncols() != output_dim {
208            return Err(StatsError::DimensionMismatch(format!(
209                "y has {} columns, expected {} to match the network's output layer",
210                y.ncols(),
211                output_dim
212            )));
213        }
214        if x.ncols() != self.architecture[0] {
215            return Err(StatsError::DimensionMismatch(format!(
216                "x has {} columns, expected {} to match the network's input layer",
217                x.ncols(),
218                self.architecture[0]
219            )));
220        }
221        if config.n_ensemble == 0 {
222            return Err(StatsError::InvalidArgument(
223                "n_ensemble must be at least 1".to_string(),
224            ));
225        }
226
227        let mut rng = match config.seed {
228            Some(seed) => scirs2_core::random::rngs::StdRng::seed_from_u64(seed),
229            None => {
230                scirs2_core::random::rngs::StdRng::from_rng(&mut scirs2_core::random::thread_rng())
231            }
232        };
233
234        let n_layers = self.architecture.len() - 1;
235        let n_train = x.nrows();
236        let lr = F::from(config.learning_rate).ok_or_else(|| {
237            StatsError::InvalidArgument(
238                "learning_rate must be representable in the target float type".to_string(),
239            )
240        })?;
241
242        let mut weight_ensembles: Vec<Vec<Array2<F>>> = Vec::with_capacity(config.n_ensemble);
243        let mut bias_ensembles: Vec<Vec<Array1<F>>> = Vec::with_capacity(config.n_ensemble);
244
245        for _member in 0..config.n_ensemble {
246            let (x_train, y_train) = if config.bootstrap && n_train > 1 {
247                let idx: Vec<usize> = (0..n_train).map(|_| rng.random_range(0..n_train)).collect();
248                let xb = Array2::from_shape_fn((n_train, x.ncols()), |(i, j)| x[[idx[i], j]]);
249                let yb = Array2::from_shape_fn((n_train, y.ncols()), |(i, j)| y[[idx[i], j]]);
250                (xb, yb)
251            } else {
252                (x.to_owned(), y.to_owned())
253            };
254
255            let mut weights: Vec<Array2<F>> = Vec::with_capacity(n_layers);
256            let mut biases: Vec<Array1<F>> = Vec::with_capacity(n_layers);
257            for l in 0..n_layers {
258                let fan_in = self.architecture[l];
259                let fan_out = self.architecture[l + 1];
260                let w_prec = prior_precision(&self.weight_priors[l]);
261                let b_prec = prior_precision(&self.bias_priors[l]);
262                let w = Array2::from_shape_fn((fan_in, fan_out), |_| {
263                    sample_normal(F::zero(), w_prec, &mut rng)
264                });
265                let b =
266                    Array1::from_shape_fn(fan_out, |_| sample_normal(F::zero(), b_prec, &mut rng));
267                weights.push(w);
268                biases.push(b);
269            }
270
271            for _epoch in 0..config.epochs {
272                let (grads_w, grads_b) =
273                    self.backward(&x_train.view(), &y_train.view(), &weights, &biases)?;
274                for l in 0..n_layers {
275                    let dw = grads_w[l].mapv(|g| g * lr);
276                    weights[l] = &weights[l] - &dw;
277                    let db = grads_b[l].mapv(|g| g * lr);
278                    biases[l] = &biases[l] - &db;
279                }
280            }
281
282            weight_ensembles.push(weights);
283            bias_ensembles.push(biases);
284        }
285
286        self.weight_samples = Some(weight_ensembles);
287        self.bias_samples = Some(bias_ensembles);
288        Ok(())
289    }
290
291    /// Make predictions with uncertainty quantification.
292    ///
293    /// If [`Self::fit`] has already trained a posterior ensemble, the ensemble
294    /// is a small *finite* population (typically a handful to a few dozen
295    /// members) whose exact predictive mean/variance is directly computable,
296    /// so this forward-propagates through **every** trained member exactly
297    /// once (`n_samples_` is ignored in this case: randomly resampling a
298    /// finite, fully-known population with replacement would only add
299    /// spurious Monte Carlo noise around a quantity that has no randomness
300    /// left to estimate). Otherwise -- no ensemble has been trained yet --
301    /// this falls back to real prior-predictive Monte Carlo: `n_samples_`
302    /// independent draws of the *whole network* from its priors, each
303    /// forward-propagated through the actual input `x`, where genuine random
304    /// sampling is unavoidable since the prior is a continuous distribution.
305    /// Either way, the reported mean/variance are genuine empirical
306    /// statistics of real forward passes -- never a constant placeholder.
307    pub fn predict_with_uncertainty(
308        &self,
309        x: &ArrayView2<F>,
310        n_samples_: usize,
311    ) -> StatsResult<(Array2<F>, Array2<F>)> {
312        checkarray_finite(x, "x")?;
313        if n_samples_ == 0 {
314            return Err(StatsError::InvalidArgument(
315                "n_samples_ must be at least 1".to_string(),
316            ));
317        }
318        if x.ncols() != self.architecture[0] {
319            return Err(StatsError::DimensionMismatch(format!(
320                "x has {} columns, expected {} to match the network's input layer",
321                x.ncols(),
322                self.architecture[0]
323            )));
324        }
325
326        let n_test = x.nrows();
327        let output_dim = *self.architecture.last().ok_or_else(|| {
328            StatsError::InvalidArgument("architecture must be non-empty".to_string())
329        })?;
330        let n_layers = self.architecture.len() - 1;
331        let mut rng = scirs2_core::random::thread_rng();
332
333        let mut draws: Vec<Array2<F>> = Vec::with_capacity(n_samples_);
334
335        match (&self.weight_samples, &self.bias_samples) {
336            (Some(w_ens), Some(b_ens)) if !w_ens.is_empty() && !b_ens.is_empty() => {
337                let n_members = w_ens.len().min(b_ens.len());
338                for idx in 0..n_members {
339                    draws.push(self.forward(x, &w_ens[idx], &b_ens[idx])?);
340                }
341            }
342            _ => {
343                for _ in 0..n_samples_ {
344                    let mut weights = Vec::with_capacity(n_layers);
345                    let mut biases = Vec::with_capacity(n_layers);
346                    for l in 0..n_layers {
347                        let fan_in = self.architecture[l];
348                        let fan_out = self.architecture[l + 1];
349                        let w_prec = prior_precision(&self.weight_priors[l]);
350                        let b_prec = prior_precision(&self.bias_priors[l]);
351                        let w = Array2::from_shape_fn((fan_in, fan_out), |_| {
352                            sample_normal(F::zero(), w_prec, &mut rng)
353                        });
354                        let b = Array1::from_shape_fn(fan_out, |_| {
355                            sample_normal(F::zero(), b_prec, &mut rng)
356                        });
357                        weights.push(w);
358                        biases.push(b);
359                    }
360                    draws.push(self.forward(x, &weights, &biases)?);
361                }
362            }
363        }
364
365        let mut predictions = Array2::<F>::zeros((n_test, output_dim));
366        let mut prediction_vars = Array2::<F>::zeros((n_test, output_dim));
367        let s_f = F::from(draws.len()).expect("draw count fits in any Float");
368
369        for i in 0..n_test {
370            for j in 0..output_dim {
371                let m = draws.iter().fold(F::zero(), |acc, d| acc + d[[i, j]]) / s_f;
372                let v = draws
373                    .iter()
374                    .fold(F::zero(), |acc, d| acc + (d[[i, j]] - m) * (d[[i, j]] - m))
375                    / s_f.max(F::one());
376                predictions[[i, j]] = m;
377                prediction_vars[[i, j]] = v;
378            }
379        }
380
381        Ok((predictions, prediction_vars))
382    }
383}
scirs2_stats/bayesian_advanced/bnn_train.rs

scirs2_stats/bayesian_advanced/
bnn_train.rs
