GradientLSS-rs

Distributional Gradient Boosting for Location, Scale, and Shape (LSS) in Rust.

A Rust implementation of probabilistic gradient boosting, supporting XGBoost and LightGBM backends.

Features

• Unified Distribution Interface: Common Distribution trait for all probability distributions
• Multiple Backends: Support for XGBoost and LightGBM through feature flags
• Response Functions: Identity, exp, softplus, sigmoid, relu, and more
• Loss Functions: Negative log-likelihood (NLL) and CRPS
• Gradient Stabilization: MAD and L2 stabilization methods
• Multivariate Distributions: Support for multivariate normal (MVN) and other multivariate distributions

Installation

Add to your Cargo.toml:

[dependencies]
gradientlss = "0.1"

# Enable backends as needed:
# gradientlss = { version = "0.1", features = ["xgboost"] }
# gradientlss = { version = "0.1", features = ["lightgbm"] }
# gradientlss = { version = "0.1", features = ["full"] }  # Both backends

System Requirements for Backends

Both XGBoost and LightGBM backends require native libraries to be compiled:

1. CMake (required for both):

   # macOS
   brew install cmake
   
   # Ubuntu/Debian
   sudo apt-get install cmake
   
   # Windows
   # Download from https://cmake.org/download/

2. Clang/LLVM (for bindgen):

   # macOS - usually pre-installed with Xcode
   xcode-select --install
   
   # Ubuntu/Debian
   sudo apt-get install llvm-dev libclang-dev clang

Breaking Changes

If you're updating from an earlier build, the following changes require action or change numerical results:

TrainConfig has a new required field: collect_train_metrics

Any code constructing TrainConfig as an exhaustive struct literal will fail to compile until the field is added:

let config = TrainConfig {
    num_boost_round: 100,
    early_stopping_rounds: Some(20),
    verbose: false,
    seed: 123,
    collect_train_metrics: false, // <-- new field
};

// or sidestep future additions entirely:
let config = TrainConfig {
    num_boost_round: 100,
    ..TrainConfig::default()
};

The field controls whether the per-round training-set metric is computed when nothing else needs it. It defaults to false because skipping the pass measured ~20% faster end-to-end training with a validation set (see examples/bench_collect_train_metrics.rs and OPTIMIZATION.md). Consequences:

• TrainingResult::train_history is now empty by default when a validation set drives early stopping. Set collect_train_metrics: true to get the per-round train loss curve back.
• Nothing else changes: the train metric is still force-computed whenever it is actually consumed — no validation set (train loss drives early stopping), verbose: true, or registered callbacks — so early stopping, callbacks, and the trained model itself are identical either way.

LightGBM saved-model format changed

LightGBM models now persist best_iteration (an 8-byte header before the model text) so that predictions after early stopping use the best round, matching LightGBM's Python behavior. Models saved with an older build will fail to load — retrain and re-save them.

Predictions may shift: best_iteration now tracks the true optimum

When early stopping is active, prediction truncates the ensemble to the best boosting iteration. Previously that iteration only advanced when the validation metric improved by the early-stopping min-delta (1e-4 relative), so it could lag the real optimum and predict a slightly under-fit model. It now records the true argmin of the validation curve (any improvement counts), matching XGBoost/LightGBM in Python.

• Effect: predictions from a model trained with a validation set + early stopping can change slightly — generally a touch more accurate. Models trained without early stopping are unchanged (all trees are used).
• Action: none required. If you pinned expected prediction values in tests, regenerate the baselines.

The early-stopping decision is unchanged: the patience counter still uses the min-delta, so training stops after the same number of rounds as before — only which iteration is selected for prediction moved.

Hyperparameter search now shuffles CV folds (results change)

hyper_opt with the default CvScheme::KFold now draws one seeded row permutation (from HyperOptConfig::seed) and reuses it for every trial, instead of using contiguous row-order folds. This makes seed actually control fold composition (as documented) and matches xgboost.cv's random partitioning.

• Effect: CV scores — and therefore which hyperparameters are selected — will differ from previous runs. This is a one-time shift toward more reliable estimates (contiguous folds were biased whenever rows were ordered by target).
• Reproducibility: a fixed seed (and hp_seed) now gives identical results run-to-run; a previous bug tied fold/RNG order to HashMap iteration order, so even seeded runs varied.

Time-series data is unaffected — and this is the mechanism to use. Shuffling is applied only to CvScheme::KFold. Set cv_scheme: CvScheme::TimeSeries and folds stay in row order with forward-chaining (expanding window): fold i trains only on rows strictly before its test segment, so no future row ever informs a past prediction. If your rows encode time, use TimeSeries and none of the shuffling above applies:

let config = HyperOptConfig {
    cv_scheme: CvScheme::TimeSeries, // preserves order; no shuffle
    n_folds: 5,
    ..HyperOptConfig::default()      // KFold is the default
};

(The single 80/20 holdout, n_folds <= 1, is likewise never shuffled.)

LightGBM: user-set parameters now actually apply (results change)

LightGBMParams::set previously appended to LightGBM's parameter list, and LightGBM keeps the first occurrence of a duplicated key — so any override of a default (learning_rate, num_leaves, boosting, objective, verbose, plus num_class on repeated set_n_dist_params) was silently ignored, with the warning suppressed by verbose=-1. This included every hyper_opt-tuned value: LightGBM trials all trained with learning_rate=0.1, num_leaves=31 regardless of what the optimizer sampled. set now replaces in place.

• Effect: LightGBM models trained with non-default params (directly or via hyper_opt) will differ — they now honor your configuration. Re-evaluate any previously tuned "best" LightGBM hyperparameters: the scores they were selected by did not measure them.
• Relatedly, the binned Dataset is now constructed with the full training params (as lgb.train does in Python), so dataset-level params like max_bin, min_data_in_leaf (feature pre-filtering), and zero_as_missing take effect instead of being pinned to defaults.
• The [low, high] hyper_opt shorthand with integer bounds (e.g. num_leaves: [2, 64]) now samples integers (optuna suggest_int semantics); LightGBM rejects float-valued integer params, which the override fix exposed.

XGBoost: objective and base_score are now forced (matching Python)

Python XGBoostLSS's set_params_adj unconditionally overwrites objective=None, base_score=0, and disable_default_eval_metric=True; the Rust backend now does the same at train time. A user-carried objective (e.g. "count:poisson") used to silently corrupt training — every internal predict returned link-transformed values while boosting stayed on margin scale — and a nonzero base_score shifted all predictions. If you were setting either, remove them; the distribution's response functions are the link.

TrainConfig::seed is now wired into both backends (previously silently ignored). Runs using subsample/colsample_*/bagging_fraction may differ; an explicit seed set via backend params still wins.

Mixture distribution fixes (predictions change)

• predict(PredType::Samples | Quantiles) drew mixture components with the wrong probabilities: a second Gumbel perturbation plus temperature softmax was layered on top of the already-softmaxed mixing weights, biasing selection toward uniform (weights 0.9/0.1 sampled at roughly 0.82/0.18). Components are now drawn as an exact categorical, matching MixtureSameFamily.sample().
• predict(PredType::Parameters) now returns mixing probabilities (summing to 1) for the mix_prob_* columns, matching Python's predict_dist. Previously it returned raw unbounded logits.

NLL metric now skips NaN terms (torch.nansum parity)

The evaluation metric aggregates per-sample log-probs with NaN treated as 0, matching Python's -torch.nansum(log_prob). Previously one NaN log-prob (e.g. a 0·inf at a support boundary) made the metric NaN for the rest of training, which froze best_iteration and burned out the early-stopping patience where Python trains on.

CRPS loss now errors for non-reparameterizable distributions

LossFn::Crps computes gradients by finite differences through fixed-seed sampling, which is only meaningful when the sampler is smooth in its parameters (the rsample analogue). Training now fails fast with a clear error for distributions whose samplers use rejection sampling or discrete draws (Gamma, Beta, Dirichlet, StudentT, Poisson, NegativeBinomial, the ZI/ZA families, MVT, Mixture) instead of silently producing garbage gradients. Torch errors the same way — these have no rsample. Loc-scale and inverse-CDF families (Gaussian, LogNormal, Logistic, Gumbel, Laplace, Cauchy, Weibull, Expectile, MVN, MVNLoRa, SplineFlow) still support CRPS.

Smaller behavioral fixes

• XGBoost feature_importance now honors the requested type: Gain/Cover are parsed from the model dump (per-split averages, matching get_score); previously every type silently returned split counts.
• XGBoost saved models append n_features (trailing, optional — old files still load, reporting 0); GradientLSS::num_features() now returns the real count instead of a stub 0.
• A CallbackAction::Stop on a new-minimum round no longer excludes that round from best_iteration.
• num_boost_round: 0 now reports n_iterations: 0 instead of 1.
• Sigmoid response derivatives are 0 in the clamp region (|x| ≳ 6.9), matching torch autograd through torch.clamp and this crate's own numerical path.
• Softplus keeps its tail below x = −20 (exp(x) + ε) instead of flooring to ε, matching torch.

Usage

Basic Example

use gradientlss::prelude::*;
use gradientlss::distributions::Gaussian;

// Create a Gaussian distribution
let dist = Gaussian::new(
    Stabilization::None,
    ResponseFn::Exp,  // For scale parameter
    LossFn::Nll,
    false,  // Don't initialize with start values
);

// With XGBoost backend
#[cfg(feature = "xgboost")]
{
    use gradientlss::backend::XGBoostBackend;
    
    let mut model = GradientLSS::<XGBoostBackend>::new(dist);
    
    // Create dataset
    let mut train_data = XGBoostDataset::from_data(
        features.view(),
        labels.view(),
    )?;
    
    // Train
    let params = XGBoostBackend::create_params(2);  // 2 params: loc, scale
    let config = TrainConfig {
        num_boost_round: 100,
        early_stopping_rounds: Some(20),
        verbose: true,
        seed: 123,
        collect_train_metrics: false,  // set true to record per-round train loss
    };
    
    model.train(&mut train_data, None, params, config)?;
    
    // Predict
    let predictions = model.predict(
        &test_features.view(),
        PredType::Parameters,
        1000,  // n_samples (for sampling)
        &[0.1, 0.5, 0.9],  // quantiles
        123,  // seed
    )?;
}

Available Distributions

Univariate Distributions

• Gaussian - Normal distribution with loc (mean) and scale (std) parameters
• Gamma - Gamma distribution
• Beta - Beta distribution
• StudentT - Student's t distribution
• Poisson - Poisson distribution (discrete)
• NegativeBinomial - Negative binomial distribution (discrete)
• Weibull - Weibull distribution
• LogNormal - Log-normal distribution
• Cauchy - Cauchy distribution
• Laplace - Laplace distribution
• Gumbel - Gumbel distribution
• Logistic - Logistic distribution
• ZAGamma - Zero-adjusted Gamma distribution
• ZINB - Zero-inflated Negative Binomial distribution
• ZIPoisson - Zero-inflated Poisson distribution

Multivariate Distributions

• MVN - Multivariate Normal distribution with mean vector and Cholesky-decomposed covariance matrix
• MVT - Multivariate Student's T distribution with degrees of freedom, mean vector, and Cholesky-decomposed covariance matrix
• Dirichlet - Dirichlet distribution for compositional data (proportions that sum to 1)

More distributions can be added by implementing the Distribution trait.

Response Functions

Function	Description	Use Case
Identity	No transformation	Location parameters
Exp	Exponential	Strictly positive (scale)
Softplus	ln(1 + exp(x))	Smooth positive
Sigmoid	1/(1 + exp(-x))	Bounded (0, 1)
ExpDf	exp(x) + 2	Degrees of freedom

Multivariate Usage

For multivariate distributions, the target data must be provided in a flattened format:

MVN (Multivariate Normal)

use gradientlss::distributions::MVN;
use gradientlss::backend::lightgbm_backend::LightGBMBackend;
use gradientlss::model::GradientLSS;
use ndarray::array;

// Create a 2-target MVN distribution
let mvn = MVN::new(2, Stabilization::None, ResponseFn::Exp, LossFn::Nll, false);
let mut model = GradientLSS::<LightGBMBackend>::new(mvn);

// For 3 observations with 2 targets each, labels should be:
// [y1_obs1, y2_obs1, y1_obs2, y2_obs2, y1_obs3, y2_obs3]
let features = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
let labels = array![1.0, 2.0, 2.0, 3.0, 3.0, 4.0]; // Flattened multivariate targets

let mut train_data = LightGBMBackend::Dataset::from_data(features.view(), labels.view())?;

MVT (Multivariate Student's T)

use gradientlss::distributions::{MVT, ResponseFn};

// Create a 2-target MVT distribution with separate response functions
let mvt = MVT::new(
    2,
    Stabilization::None,
    ResponseFn::Exp,      // For scale parameters
    ResponseFn::ExpDf,    // For degrees of freedom (ensures df > 2)
    LossFn::Nll,
    false
);

Dirichlet (Compositional Data)

use gradientlss::distributions::Dirichlet;

// Create a 3-target Dirichlet distribution for compositional data
let dirichlet = Dirichlet::new(3, Stabilization::None, ResponseFn::Exp, LossFn::Nll, false);

// For Dirichlet, targets must sum to 1 for each observation
// [p1_obs1, p2_obs1, p3_obs1, p1_obs2, p2_obs2, p3_obs2, ...]
let features = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
let labels = array![0.3, 0.4, 0.3, 0.2, 0.5, 0.3, 0.1, 0.6, 0.3];

The model will automatically handle the reshaping and parameter estimation for each multivariate distribution type.

Model Interpretability

GradientLSS provides multiple approaches for model interpretation, matching the functionality of XGBoostLSS/LightGBMLSS.

Built-in Interpretability Features

// Feature importance
let importance = model.feature_importance(
    FeatureImportanceType::Gain,
    Some(vec!["age".to_string(), "income".to_string()]),
)?;

// Partial dependence plots
let pdp = model.partial_dependence(
    &features.view(),
    0,  // feature index
    0,  // parameter index (e.g., 0 for 'loc')
    Some(50),  // grid size
    None,
    None,
)?;

// ICE (Individual Conditional Expectation) curves
let ice = model.ice_curves(
    &features.view(),
    0,  // feature index
    0,  // parameter index
    Some(50),
    None,
    None,
)?;

Plotting (with plotting feature)

// Feature importance plot
model.plot_feature_importance(
    "importance.png",
    FeatureImportanceType::Gain,
    None,  // aggregated across parameters
    Some(feature_names),
    None,
)?;

// Partial dependence plot
model.plot_partial_dependence(
    &features.view(),
    0,  // feature index
    0,  // parameter index
    "pdp.png",
    None,
)?;

// Expectile plot (for Expectile distribution models)
model.expectile_plot(
    &features.view(),
    0,  // feature index
    "expectiles.png",
    Some(50),  // grid size
    Some("Age".to_string()),
    None,
    true,  // show confidence bands
)?;

SHAP Integration

For advanced SHAP-based visualizations (beeswarm plots, dependency plots), GradientLSS provides a data export workflow compatible with Python's shap library:

// Export data for SHAP analysis
let shap_data = model.export_for_shap(
    &features.view(),
    Some(vec!["feature1".to_string(), "feature2".to_string()]),
)?;

// Save to JSON for Python consumption
shap_data.to_json_file("shap_export.json")?;

Then in Python:

import json
import numpy as np
import shap
import xgboost as xgb

# Load the exported data
with open('shap_export.json', 'r') as f:
    data = json.load(f)

X = np.array(data['features'])
predictions = np.array(data['predictions'])
feature_names = data['feature_names']
param_names = data['param_names']

# Use feature importance directly from GradientLSS
if data['feature_importance']:
    for param, importance in data['feature_importance'].items():
        print(f"{param}: {importance}")

# Or train a surrogate model for full SHAP values
# (useful for beeswarm plots and SHAP dependency plots)
model_loc = xgb.XGBRegressor()
model_loc.fit(X, predictions[:, 0])  # Train on 'loc' predictions

explainer = shap.TreeExplainer(model_loc)
shap_values = explainer.shap_values(X)
shap.summary_plot(shap_values, X, feature_names=feature_names)

See examples/shap_integration.rs for a complete example.

Backend Differences

XGBoost (--features xgboost)

• Full support for custom objective functions
• Uses update_custom() for distributional gradient updates
• Row-major (C-order) gradient layout

LightGBM (--features lightgbm)

• Limited: Current Rust bindings don't expose custom objective API
• Uses built-in objectives only
• Column-major (Fortran-order) gradient layout
• Consider using XGBoost for full distributional regression support

Architecture

gradientlss/
├── distributions/     # Distribution implementations
│   ├── base.rs       # Distribution trait
│   └── gaussian.rs   # Gaussian distribution
├── backend/          # Gradient boosting backends
│   ├── traits.rs     # Backend trait definitions
│   ├── xgboost_backend.rs
│   └── lightgbm_backend.rs
├── utils.rs          # Response functions
├── model.rs          # GradientLSS model wrapper
└── error.rs          # Error types

Adding New Distributions

1. Create a new file in src/distributions/
2. Implement the Distribution trait:

impl Distribution for MyDistribution {
    fn n_params(&self) -> usize { /* ... */ }
    fn params(&self) -> &[DistributionParam] { /* ... */ }
    fn log_prob(&self, params: &[f64], target: f64) -> f64 { /* ... */ }
    fn nll(&self, params: &ArrayView2<f64>, target: &ArrayView1<f64>) -> f64 { /* ... */ }
    fn sample(&self, params: &ArrayView2<f64>, n_samples: usize, seed: u64) -> Array2<f64> { /* ... */ }
    // ... other required methods
}

3. Export from src/distributions/mod.rs

License

MIT