Skip to content
brndnmtthwsPublic

About

Fast, deterministic gradient boosting in Rust (with Python bindings)

Topics

Resources

Stars

3 stars

Watchers

0 watching

Forks

Repository files navigation

hessboost

crates.io docs.rs CI license

hessboost

Fast, deterministic gradient boosting in Rust (with Python bindings). hessboost provides multi-core tree building with runtime-detected NEON and AVX2 SIMD, strict parameter validation, reproducible models on any thread count, and stable model storage. It supports modern extensions like conformal prediction, explainable boosting machines (EBMs), distributional modeling, and tree-based diffusion, alongside bidirectional XGBoost JSON/UBJSON model interchange.

The name comes from the Hessian: hessboost fits each tree to the loss's gradients and second derivatives (Newton boosting).

Why hessboost

  • Fast. Multi-core training with runtime-detected NEON and AVX2 kernels; see docs/performance.md.
  • Lean on memory. A matrix, its copies with other labels, and every training of them share one compact binned index, built once even for concurrent trainings; row selections are views of their source, a read-only array (a memory-mapped .npy included) can be borrowed instead of copied, and a binned matrix trains without the feature values at all.
  • Deterministic. The same parameters, data, and seed produce the exact same model on any thread count.
  • Strict. Invalid parameters and unsupported combinations fail loudly; nothing is silently ignored.
  • Stable storage. Models saved natively are forwards-compatible across releases.
  • Modern modeling. Built-in support for conformal intervals, Boulevard confidence bands, explainable boosting machines (EBMs), distributional boosting, SGLB uncertainty, tree-based diffusion, in-place updates, and compact models.
  • XGBoost compatible. Accepts standard XGBoost parameter, objective, and metric names, with bidirectional JSON and UBJSON model interchange.

Getting started

cargo add hessboost

For users pinning a release series in a Cargo manifest:

hessboost = "0.4"

Needs Rust 1.93 or newer and a C compiler (to build libzstd).

use hessboost::prelude::*;

fn main() -> Result<()> {
    // 100 rows × 4 features, row-major, and one label per row.
    let (n_rows, n_cols) = (100, 4);
    let x: Vec<f32> = (0..n_rows * n_cols).map(|i| (i % 17) as f32 / 17.0).collect();
    let y: Vec<f32> = x.chunks(n_cols).map(|row| 2.0 * row[0] - row[1]).collect();

    let dtrain = DMatrix::from_dense(&x, n_rows, n_cols)?.with_labels(&y)?;

    let params = TrainingParams::builder()
        .objective(Objective::SquaredError(RegLoss::default()))
        .tree_method(TreeMethod::Hist)
        .max_depth(6)
        .eta(0.1)
        .subsample(0.9)
        .build()?;

    let model = train(&params, &dtrain, 200)?;
    let preds = model.predict(&dtrain, Iterations::Best)?;
    println!("first prediction: {}", preds.get(0, 0).unwrap());

    model.save("model.bin", ModelFormat::Binary)?;
    let reloaded = BoostedModel::load("model.bin", ModelFormat::Binary)?;
    assert_eq!(reloaded.predict(&dtrain, Iterations::Best)?, preds);
    Ok(())
}

For eval sets, early stopping, custom objectives, or continued training, use Trainer (the xgb.train keyword-argument equivalent):

use std::num::NonZeroUsize;

let result = Trainer::new(&params, &dtrain, 1000)
    .eval(&dvalid, "valid")
    .early_stopping_rounds(NonZeroUsize::new(20).unwrap())
    .train()?;
let model = result.model; // predicts with the best iteration

To ship a self-contained (e.g. static) binary, compile the model file into it; it decodes on first use:

use hessboost::model::EmbeddedModel;

static MODEL: EmbeddedModel = EmbeddedModel::new(include_bytes!("model.bin"), ModelFormat::Binary);

let preds = MODEL.get()?.predict(&dtest, Iterations::Best)?;

Every type and option is in the API docs; runnable programs live in examples/ (cargo run --release --example <name>):

Example Shows
train_regression end-to-end regression with feature importance
binary_classification a watched eval set, early stopping, AUC
balanced_bagging LightGBM class-stratified sampling for imbalanced binary classification
multiclass per-class probabilities and predicted classes
ranking / rank_xendcg LambdaMART and XE-NDCG with query bagging
constraints monotone and interaction constraints, categorical features
custom_objective a custom loss and eval metric
shap SHAP contributions and interaction values
model_io native and XGBoost JSON/UBJSON save and load
conformal calibrated prediction intervals
boulevard_inference confidence intervals for f(x) and prediction intervals
ebm an explainable boosting machine's shape functions and their confidence bands
distributional predictive distributions, intervals, and NLL
virtual_ensembles SGLB posterior sampling: knowledge uncertainty rising off the training data
tree_diffusion sampling multimodal and skewed p(y | x) with tree diffusion and flow matching
forest_flow synthetic tabular rows and imputation with ForestFlow / ForestDiffusion
ordered_target_stats encoding a high-cardinality categorical
compact_model reuse penalties and the compact model format
budget budget training against default and tuned training
online_update adding and deleting training rows in place, and exact unlearning
pfn_boost boosting from a pretrained model's logits
metal CPU vs GPU prediction (macOS, --features metal)
wgpu GPU training and prediction through wgpu, bit-identical to the CPU (--features wgpu; runs on Mesa's lavapipe without a GPU)

Python

python/ holds the Python package (pip install hessboost): DMatrix, train, cv, and Booster (taking XGBoost's parameter names), scikit-learn estimators, pandas and polars (1.x and 2.x; DataFrames and LazyFrames, with labels taken from the frame's columns by name) categorical input, and the conformal, distributional, tree-diffusion, ForestFlow, in-place update, ordered target statistics, budget training, and compact model extras, with GPU training and batch prediction (device="wgpu" on every platform, device="metal" on macOS, and device="cuda" on Linux with pip install 'hessboost[cuda]'; Booster.to_gpu()):

import hessboost

booster = hessboost.train(
    {"objective": "binary:logistic", "max_depth": 4}, hessboost.DMatrix(X, label=y), 100
)
probabilities = booster.predict(X_test)

See python/README.md.

Features

  • Core boosting: gbtree, dart, and gblinear boosters, boosted random forests, and exact, hist, and approx tree methods with native missing-value and categorical support.
  • Objectives & metrics: Regression (squared, log, Huber, quantile, expectile), binary/multiclass classification, ranking (LambdaMART, XE-NDCG), count, and survival (Cox, AFT), plus typed objective and metric APIs and custom loss hooks. Quantile and absolute-error models can set each leaf to the exact residual quantile of its rows (leaf_estimation_method = exact).
  • Validation & workflow: Cross-validation (including purged and forward time-series folds, optionally at caller-given block starts such as calendar-aligned walk-forward blocks, leakage-safe per-fold early stopping on a purged tail of each fold's training rows, out-of-fold margins and per-fold round counts, whole-query ranking folds, per-fold target statistics, continuing a model, and refitting on all rows at the chosen round count), early stopping, feature importance, SHAP values and interactions, model slicing and averaging (bagged models as one), and iteration ranges.
  • Prediction & inspection: Batch, in-place (borrowed rows), and allocation-free single-row prediction, bit-identical in every batch; read-only model metadata (tree weights, outputs, depths, and nodes, gblinear weights, model-shrinkage records).
  • Interchange: Native binary and JSON formats, XGBoost JSON/UBJSON import/export, LightGBM model import, and models embedded in the binary at compile time.
  • Modern modeling (opt-in):
    • Conformal intervals: Finite-sample coverage guarantees.
    • Boulevard inference: Asymptotic confidence and prediction intervals for f(x).
    • Explainable boosting machines: Interpretable cyclic GAM/GA²M models with shape-function confidence bands.
    • Distributional boosting: Full predictive distributions per row (dist:normal, dist:gamma, dist:student_t for heavy-tailed targets with learned degrees of freedom, etc.).
    • Uncertainty & virtual ensembles: SGLB posterior sampling and model shrinkage.
    • Generative tabular modeling: Tree-based conditional diffusion, flow matching, and ForestFlow synthetic data and imputation.
    • In-place updates: Fast incremental learning and exact or approximate unlearning.
    • Compact models: Bit-packed model format with identical margins.
    • Budget training: Training controlled by one budget value, based on PerpetualBooster.
    • Metal GPU: Apple Silicon GPU prediction and training (--features metal).
    • wgpu GPU: Vulkan, Metal, and DirectX 12 GPU prediction and training through wgpu (--features wgpu), bit-identical to the CPU.
    • CUDA GPU: NVIDIA GPU prediction and training on Linux (--features cuda; building it needs CUDA 13's headers and libclang), bit-identical to the CPU.

Caveats

  • Full technical details, invariants, and statistical assumptions are documented in the API reference.
  • Approximate in-place updates are designed for incremental shifts (under ~1% of rows); larger changes benefit from a retrain.
  • Asymptotic Boulevard inference and prediction intervals require specific noise and structure assumptions; see the inference docs for conditions and empirical coverage validation.

Not implemented

  • Distributed and external-memory training.
  • CLI and C bindings.
  • A few XGBoost options exist at one setting only, and a few metrics are missing; the API docs list them.

Contributing

AGENTS.md has the build, lint, and test commands and the project's invariants; scripts/README.md covers the XGBoost parity suite and benchmark harnesses.

License and attribution

Licensed under the Apache License, Version 2.0. Copyright 2026 Brenden Matthews.

hessboost is a fork of sequoia-boost (Copyright 2026 Patrick Garrett, Apache-2.0).

hessboost is not affiliated with or endorsed by the XGBoost project, and contains no XGBoost source code.

Budget training reimplements PerpetualBooster's algorithm (Copyright 2024 Perpetual ML, Apache-2.0); no Perpetual code is copied.

The error function used by the AFT normal distribution is ported from glibc 2.41's s_erf.c, derived from Sun Microsystems' fdlibm (Copyright (C) 1993 Sun Microsystems, Inc.); src/objective/distributional/special.rs carries its notice.

About

Fast, deterministic gradient boosting in Rust (with Python bindings)

Topics

Resources

Stars

3 stars

Watchers

0 watching

Forks

Releases

Contributors

Languages