hessboost 0.2.0

//! Virtual ensembles and uncertainty decomposition for models trained with
//! SGLB posterior sampling (CatBoost's
//! `virtual_ensembles_predict`).
//!
//! A model trained with
//! [`posterior_sampling`](crate::config::TrainingParams::posterior_sampling)
//! (Stochastic Gradient Langevin Boosting with model shrinkage; Ustimenko and
//! Prokhorenkova, *SGLB: Stochastic Gradient Langevin Boosting*, ICML 2021)
//! draws its later iterates from the Bayesian posterior of the ensemble.
//! Malinin, Prokhorenkova and Ustimenko (*Uncertainty in Gradient Boosting
//! via Ensembles*, ICLR 2021) turn one such model into a *virtual ensemble*:
//! the models after several iterations of its second half act as posterior
//! samples, and the spread of their predictions measures knowledge
//! (epistemic) uncertainty, which grows away from the training data.
//!
//! [`BoostedModel::predict_virtual_ensembles`] returns the members'
//! predictions ([`VirtualEnsembles`]), [`BoostedModel::predict_uncertainty`]
//! their decomposition ([`Uncertainty`]).
//!
//! # Members
//!
//! As in CatBoost's `ApplyVirtualEnsembles`
//! (`catboost/private/libs/algo/apply.cpp`), `K` members of a model with `T`
//! (effective) iterations are the models after `b + p`, `b + 2p`, ..., `T`
//! iterations, with period `p = ⌊T / (2K)⌋` and `b = T − pK`; `p` must be
//! positive and `b` too (at least `2K` iterations, roughly). Because model
//! shrinkage rescales the ensemble every iteration, a member is not a prefix
//! of the trees: CatBoost multiplies the prefix by `(1 − rate · eta)^(k − T)`
//! in `f32`, which holds for the `constant` shrink mode only. hessboost
//! rebuilds each member from the stored shrinkage coefficients with the
//! arithmetic of training (see
//! [`predict_margin`](BoostedModel::predict_margin)), so member
//! `k` predicts bit for bit what the same run stopped after `k` rounds
//! predicts, in either shrink mode. One pass over the trees yields every
//! member: the margins are copied out at each member's iteration count. A model trained without shrinkage has
//! plain prefixes as members. EBMs are refused (their trees are ordered by
//! bag, stage, and term), and so is a `booster = boulevard` model: its
//! leaves carry the average over all of its rounds, so its prefixes are not
//! the models of fewer rounds (its variance is
//! [`BoulevardInference`](crate::inference::BoulevardInference)'s).
//!
//! # Decomposition
//!
//! Following CatBoost (`CalcRegressionUncertaitny`,
//! `CalcClassificationUncertainty`, `CalcMulticlassUncertainty` in
//! `catboost/libs/eval_result/eval_helpers.cpp`; the CatBoost *Uncertainty*
//! reference), with `K` members:
//!
//! - **Regression** (`reg:*`, `count:*`, `survival:*`): knowledge
//!   uncertainty is the variance (over members, divided by `K`) of the
//!   members' predictions, per output. There is no data uncertainty.
//! - **Distributional** (`dist:*`, the analogue of CatBoost's
//!   `RMSEWithUncertainty` for `dist:normal`): knowledge uncertainty is the
//!   variance of the members' predicted means, data uncertainty the mean of
//!   their predicted variances, and the total their sum (the law of total
//!   variance).
//! - **Classification** (`binary:logistic`, `binary:logitraw`,
//!   `reg:logistic`, per label column for multi-label models, and
//!   `multi:softprob` / `multi:softmax` over the classes): total uncertainty
//!   is the entropy (in nats) of the mean predicted distribution, data
//!   uncertainty the mean of the members' entropies, and knowledge
//!   uncertainty their difference (the mutual information between the
//!   prediction and the ensemble member), non-negative up to rounding.
//!
//! Other objectives (ranking, `binary:hinge`, custom objectives) are
//! refused.
//!
//! ```
//! use hessboost::prelude::*;
//!
//! # fn main() -> Result<()> {
//! let x: Vec<f32> = (0..200).map(|i| (i % 50) as f32 / 50.0).collect();
//! let y: Vec<f32> = x.iter().map(|v| 2.0 * v).collect();
//! let dtrain = DMatrix::from_dense(&x, 200, 1)?.with_labels(&y)?;
//! let params = TrainingParams::builder()
//!     .posterior_sampling(true)
//!     .eta(0.1)
//!     .max_depth(3)
//!     .build()?;
//! let model = train(&params, &dtrain, 100)?;
//!
//! let ensembles = model.predict_virtual_ensembles(&dtrain, 10)?;
//! assert_eq!(ensembles.iterations(), &[55, 60, 65, 70, 75, 80, 85, 90, 95, 100]);
//! let uncertainty = model.predict_uncertainty(&dtrain, 10)?;
//! assert_eq!(uncertainty.knowledge.n_rows(), 200);
//! assert_eq!(uncertainty.knowledge.width(), 1);
//! assert!(uncertainty.data.is_none()); // squared error has no data uncertainty
//! # Ok(())
//! # }
//! ```

use super::{BoostedModel, ModelObjective, Predictions, transform_margins_in_place};
use crate::data::DMatrix;
use crate::error::{HessboostError, Result};
use crate::objective::Objective;
use crate::objective::distributional::{Dist, DistFamily};

/// The predictions of a model's virtual ensemble
/// ([`BoostedModel::predict_virtual_ensembles`]): one model per member, the
/// model after [`iterations`](Self::iterations)`()[m]` boosting iterations.
/// Every member's margins and predictions are stored contiguously,
/// `[member][row][value]`.
#[derive(Debug, Clone)]
pub struct VirtualEnsembles {
    iterations: Vec<usize>,
    n_rows: usize,
    /// Values per row of a member's margins (the model's outputs).
    margin_width: usize,
    /// Values per row of a member's predictions.
    width: usize,
    margins: Vec<f32>,
    predictions: Vec<f32>,
}

impl VirtualEnsembles {
    /// Number of members.
    pub fn n_members(&self) -> usize {
        self.iterations.len()
    }

    /// The iteration count of every member's model, ascending; the last is
    /// the whole (effective) model.
    pub fn iterations(&self) -> &[usize] {
        &self.iterations
    }

    /// Number of predicted rows.
    pub fn n_rows(&self) -> usize {
        self.n_rows
    }

    /// Values per row of a member's predictions (the width of
    /// [`BoostedModel::predict`]'s layout).
    pub fn width(&self) -> usize {
        self.width
    }

    /// Values per row of a member's margins (the width of
    /// [`BoostedModel::predict_margin`]'s layout, the model's outputs).
    pub fn margin_width(&self) -> usize {
        self.margin_width
    }

    /// Member `m`'s predictions (`[row][value]`, [`Self::width`] wide), as
    /// [`BoostedModel::predict`] returns them for its model, or `None` past
    /// the last member.
    pub fn member_predictions(&self, m: usize) -> Option<&[f32]> {
        let len = self.n_rows * self.width;
        (m < self.n_members()).then(|| &self.predictions[m * len..(m + 1) * len])
    }

    /// Member `m`'s raw margins (`[row][output]`, [`Self::margin_width`]
    /// wide), as [`BoostedModel::predict_margin`] returns them for its
    /// model, or `None` past the last member.
    pub fn member_margins(&self, m: usize) -> Option<&[f32]> {
        let len = self.n_rows * self.margin_width;
        (m < self.n_members()).then(|| &self.margins[m * len..(m + 1) * len])
    }

    /// Every member's predictions, `[member][row][value]`
    /// ([`Self::width`] values per row).
    pub fn into_predictions(self) -> Vec<f32> {
        self.predictions
    }

    /// Every member's raw margins, `[member][row][output]`
    /// ([`Self::margin_width`] values per row).
    pub fn into_margins(self) -> Vec<f32> {
        self.margins
    }
}

/// A virtual ensemble's uncertainty decomposition
/// ([`BoostedModel::predict_uncertainty`]; see the
/// [module docs](self#decomposition) for the per-objective definitions).
///
/// Every field is one row per predicted row, each with its own width.
/// `knowledge`, `data`, and `total` share theirs: one per row, except one
/// per output for multi-output regression and one per label column for
/// multi-label classification. `mean` is as wide as the mean prediction:
/// `num_class` probabilities for multiclass models (`multi:softmax` too),
/// otherwise the width of the uncertainties.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct Uncertainty {
    /// The members' mean prediction: predictions for regression, predicted
    /// means for `dist:*`, probabilities for classification
    /// (`multi:softmax` too: one per class).
    pub mean: Predictions<f64>,
    /// Knowledge (epistemic) uncertainty.
    pub knowledge: Predictions<f64>,
    /// Data (aleatoric) uncertainty; `None` for plain regression, whose
    /// models predict no spread.
    pub data: Option<Predictions<f64>>,
    /// Total uncertainty (`data + knowledge`); `None` when `data` is.
    pub total: Option<Predictions<f64>>,
}

/// How [`BoostedModel::predict_uncertainty`] reads an objective's members.
enum Decomposition {
    Regression,
    Distributional(DistFamily),
    Binary,
    Multiclass,
}

impl Decomposition {
    fn of(objective: &ModelObjective) -> Result<Self> {
        let refused = || {
            Err(HessboostError::incompatible_model(
                "objective",
                format!(
                    "uncertainty is defined for regression, `dist:*`, and probabilistic \
                     classification objectives, not `{}`",
                    objective.name()
                ),
            ))
        };
        let Some(built_in) = objective.built_in() else {
            return refused();
        };
        Ok(match built_in {
            Objective::Dist(_) => match built_in.dist_family() {
                Some(family) => Decomposition::Distributional(family),
                None => return refused(),
            },
            Objective::BinaryLogistic(_)
            | Objective::BinaryLogitRaw(_)
            | Objective::RegLogistic(_) => Decomposition::Binary,
            Objective::Softmax(_) | Objective::Softprob(_) => Decomposition::Multiclass,
            Objective::SquaredError(_)
            | Objective::SquaredLogError
            | Objective::PseudoHuber(_)
            | Objective::AbsoluteError
            | Objective::Quantile(_)
            | Objective::Expectile(_)
            | Objective::Poisson
            | Objective::Gamma(_)
            | Objective::Tweedie(_)
            | Objective::Cox
            | Objective::Aft(_) => Decomposition::Regression,
            // Hinge predicts no probabilities, ranking scores are only
            // ordered, and a custom loss's predictions are unknown.
            Objective::BinaryHinge
            | Objective::RankPairwise(_)
            | Objective::RankNdcg(_)
            | Objective::RankMap(_)
            | Objective::RankXendcg
            | Objective::Custom(_) => return refused(),
        })
    }
}

/// Binary entropy in nats, `0 ln 0 = 0`.
fn binary_entropy(p: f64) -> f64 {
    let term = |q: f64| if q > 0.0 { -q * q.ln() } else { 0.0 };
    term(p) + term(1.0 - p)
}

/// Mean and population variance of `values`.
fn mean_variance(values: impl Iterator<Item = f64> + Clone) -> (f64, f64) {
    let n = values.clone().count() as f64;
    let mean = values.clone().sum::<f64>() / n;
    let variance = values.map(|v| (v - mean) * (v - mean)).sum::<f64>() / n;
    (mean, variance)
}

impl BoostedModel {
    /// The iteration counts of a virtual ensemble of `count` members over
    /// the effective iterations (CatBoost's `ApplyVirtualEnsembles`).
    fn virtual_ensemble_iterations(&self, count: usize) -> Result<Vec<usize>> {
        if self.linear.is_some() {
            return Err(HessboostError::incompatible_model(
                "model",
                "gblinear models have no boosting iterations to form virtual ensembles from",
            ));
        }
        if self.ebm.is_some() {
            return Err(HessboostError::incompatible_model(
                "model",
                "an EBM's trees are ordered by bag, stage, and term, so its iteration prefixes \
                 are not the models of fewer rounds",
            ));
        }
        if self.boulevard.is_some() {
            return Err(HessboostError::incompatible_model(
                "model",
                "a Boulevard model's leaves carry its average over every round, so its \
                 iteration prefixes are not the models of fewer rounds; its variance comes from \
                 `inference::BoulevardInference`",
            ));
        }
        if count == 0 {
            return Err(HessboostError::invalid_param(
                "virtual_ensembles_count",
                "must be at least 1",
            ));
        }
        let end = self.effective_num_trees() / self.trees_per_iteration();
        let period = end / count.saturating_mul(2);
        if period == 0 || period * count >= end {
            return Err(HessboostError::incompatible_model(
                "virtual_ensembles_count",
                format!(
                    "{count} virtual ensembles need a model of at least {} iterations, this one \
                     has {end}",
                    count.saturating_mul(2)
                ),
            ));
        }
        let begin = end - period * count;
        Ok((1..=count).map(|m| begin + m * period).collect())
    }

    /// Predict `data` with a virtual ensemble of `count` members (CatBoost's
    /// `virtual_ensembles_predict`, default count 10): the models after
    /// several iterations of the second half of the effective iterations
    /// (see the [`uncertainty`](super::uncertainty) module docs). Meant for
    /// models trained with
    /// [`posterior_sampling`](crate::config::TrainingParams::posterior_sampling);
    /// any tree model is accepted.
    ///
    /// # Errors
    ///
    /// [`HessboostError::InvalidParameter`] (`virtual_ensembles_count`) for
    /// `count == 0`; [`HessboostError::IncompatibleModel`] for a model with
    /// fewer iterations than `count` members need
    /// (`virtual_ensembles_count`), or for a `gblinear` model, a
    /// `booster = boulevard` model or an EBM (whose prefixes are not the
    /// models of fewer rounds; `model`); plus the errors of
    /// [`Self::predict_margin`].
    pub fn predict_virtual_ensembles(
        &self,
        data: &DMatrix,
        count: usize,
    ) -> Result<VirtualEnsembles> {
        let iterations = self.virtual_ensemble_iterations(count)?;
        let margins = self.prefix_margins(data, &iterations)?;
        let (n_rows, k) = (data.n_rows(), self.n_outputs);
        // `predict`'s layout: one class index per row for `multi:softmax`.
        let width = if self
            .objective
            .built_in()
            .is_some_and(Objective::predicts_class_index)
        {
            1
        } else {
            k
        };
        let mut predictions = Vec::with_capacity(iterations.len() * n_rows * width);
        let len = n_rows * k;
        for m in 0..iterations.len() {
            let start = predictions.len();
            predictions.extend_from_slice(&margins[m * len..(m + 1) * len]);
            let values = &mut predictions[start..];
            transform_margins_in_place(
                &self.objective,
                self.max_delta_step,
                self.n_targets,
                values,
                k,
            );
            predictions.truncate(start + n_rows * width);
        }
        Ok(VirtualEnsembles {
            iterations,
            n_rows,
            margin_width: k,
            width,
            margins,
            predictions,
        })
    }

    /// The knowledge, data, and total uncertainty of `data` under a virtual
    /// ensemble of `count` members ([`Self::predict_virtual_ensembles`];
    /// decomposition per objective in the
    /// [`uncertainty`](super::uncertainty#decomposition) module docs).
    ///
    /// # Errors
    ///
    /// [`HessboostError::IncompatibleModel`] (`objective`) for objectives
    /// without a decomposition (ranking, `binary:hinge`, custom objectives),
    /// plus the errors of [`Self::predict_virtual_ensembles`].
    pub fn predict_uncertainty(&self, data: &DMatrix, count: usize) -> Result<Uncertainty> {
        let decomposition = Decomposition::of(&self.objective)?;
        let ensembles = self.predict_virtual_ensembles(data, count)?;
        let n = ensembles.n_rows;
        let k = self.n_outputs;
        let members = || 0..ensembles.n_members();
        let margin = |m: usize, row: usize, out: usize| {
            f64::from(ensembles.margins[(m * n + row) * k + out])
        };
        Ok(match decomposition {
            Decomposition::Regression => {
                let width = ensembles.width;
                let mut mean = Vec::with_capacity(n * width);
                let mut knowledge = Vec::with_capacity(n * width);
                for cell in 0..n * width {
                    let values =
                        members().map(|m| f64::from(ensembles.predictions[m * n * width + cell]));
                    let (mu, variance) = mean_variance(values);
                    mean.push(mu);
                    knowledge.push(variance);
                }
                Uncertainty {
                    mean: Predictions::new(mean, n, width),
                    knowledge: Predictions::new(knowledge, n, width),
                    data: None,
                    total: None,
                }
            }
            Decomposition::Distributional(family) => {
                let mut mean = Vec::with_capacity(n);
                let mut knowledge = Vec::with_capacity(n);
                let mut aleatoric = Vec::with_capacity(n);
                for row in 0..n {
                    let dists: Vec<_> = members()
                        .map(|m| {
                            let eta: Vec<f64> = (0..k).map(|out| margin(m, row, out)).collect();
                            family.dist_from_margins(&eta)
                        })
                        .collect();
                    let (mu, variance) = mean_variance(dists.iter().map(Dist::mean));
                    mean.push(mu);
                    knowledge.push(variance);
                    aleatoric.push(dists.iter().map(Dist::variance).sum::<f64>() / count as f64);
                }
                let total = knowledge
                    .iter()
                    .zip(&aleatoric)
                    .map(|(a, b)| a + b)
                    .collect();
                Uncertainty {
                    mean: Predictions::new(mean, n, 1),
                    knowledge: Predictions::new(knowledge, n, 1),
                    data: Some(Predictions::new(aleatoric, n, 1)),
                    total: Some(Predictions::new(total, n, 1)),
                }
            }
            Decomposition::Binary => {
                let mut mean = Vec::with_capacity(n * k);
                let mut aleatoric = Vec::with_capacity(n * k);
                let mut total = Vec::with_capacity(n * k);
                for row in 0..n {
                    for out in 0..k {
                        let probs = members().map(|m| 1.0 / (1.0 + (-margin(m, row, out)).exp()));
                        let p = probs.clone().sum::<f64>() / count as f64;
                        mean.push(p);
                        aleatoric.push(probs.map(binary_entropy).sum::<f64>() / count as f64);
                        total.push(binary_entropy(p));
                    }
                }
                let knowledge = total.iter().zip(&aleatoric).map(|(t, d)| t - d).collect();
                Uncertainty {
                    mean: Predictions::new(mean, n, k),
                    knowledge: Predictions::new(knowledge, n, k),
                    data: Some(Predictions::new(aleatoric, n, k)),
                    total: Some(Predictions::new(total, n, k)),
                }
            }
            Decomposition::Multiclass => {
                let mut mean = vec![0.0; n * k];
                let mut aleatoric = Vec::with_capacity(n);
                let mut total = Vec::with_capacity(n);
                let mut probs = vec![0.0; k];
                for row in 0..n {
                    let row_mean = &mut mean[row * k..(row + 1) * k];
                    let mut entropy_sum = 0.0;
                    for m in members() {
                        let max = (0..k)
                            .map(|c| margin(m, row, c))
                            .fold(f64::NEG_INFINITY, f64::max);
                        for (c, p) in probs.iter_mut().enumerate() {
                            *p = (margin(m, row, c) - max).exp();
                        }
                        let sum: f64 = probs.iter().sum();
                        for (p, acc) in probs.iter_mut().zip(row_mean.iter_mut()) {
                            *p /= sum;
                            *acc += *p;
                            if *p > 0.0 {
                                entropy_sum -= *p * p.ln();
                            }
                        }
                    }
                    let mut entropy_of_mean = 0.0;
                    for p in row_mean.iter_mut() {
                        *p /= count as f64;
                        if *p > 0.0 {
                            entropy_of_mean -= *p * p.ln();
                        }
                    }
                    aleatoric.push(entropy_sum / count as f64);
                    total.push(entropy_of_mean);
                }
                let knowledge = total.iter().zip(&aleatoric).map(|(t, d)| t - d).collect();
                Uncertainty {
                    mean: Predictions::new(mean, n, k),
                    knowledge: Predictions::new(knowledge, n, 1),
                    data: Some(Predictions::new(aleatoric, n, 1)),
                    total: Some(Predictions::new(total, n, 1)),
                }
            }
        })
    }
}