difftree

Differentiable decision trees in JAX with built-in uncertainty quantification.

difftree is a JAX implementation of soft, oblique-split, gradient-trainable decision trees with scikit-learn compatible estimators. The package is built around uncertainty quantification: regression leaves are heteroscedastic Gaussians (or multivariate Gaussians for multi-output targets), and the predictive distribution at each input is a mixture of Gaussians weighted by soft path probabilities. This structure delivers a calibrated predictive variance, a clean decomposition of uncertainty into within-leaf (aleatoric) and between-leaf (routing) components, and, for multivariate outputs, closed-form cross-conditioning and gradient-based inverse design as direct consequences.

For surrogate optimization, leaves can be affine (\(A_\ell x + b_\ell\), with optional Jacobian matching), the soft tree can be hardened into a crisp piecewise-affine model, and that model exports as leaf polytopes or as the solver-neutral JSON spec of jax-ldt for exact MILP embedding.

The implementation is closest in spirit to the end-to-end deterministic-tree formulation of Hehn et al. [2020], with two practical departures motivated by UQ. First, splits are kept finitely soft at test time rather than annealed to hard, so soft-routing uncertainty contributes to the predictive variance. Second, a small post-hoc calibration step, in the same family as temperature scaling for classifier logits [Guo et al., 2017], jointly fits per-depth steepness and a scalar variance correction on a held-out validation slice by minimizing NLL. The broader design space is reviewed by Costa and Pedreira [2023].

from difftree import DifferentiableTreeRegressor

model = DifferentiableTreeRegressor(depth=5, random_state=0).fit(X_train, y_train)
mean, std = model.predict(X_test, return_std=True)

User guide

• Installation
  • Development install
  • Building the documentation
• Quickstart
  • Univariate regression
  • Multivariate regression with cross-conditioning
  • Inverse design
  • Classification
  • Next steps
• Surrogates for MILP embedding
  • Affine leaves and Jacobian matching
  • Crisp mode, hardening, and the soft/hard gap
  • Export
  • Split boundaries

Background

• Why differentiable trees
• Model
  • Leaves
  • Predictive distribution
  • Variance decomposition
  • Training
  • Calibration
  • Cross-conditioning and inverse design
• Limitations and caveats

Examples

• Heteroscedastic UQ with DifferentiableTreeRegressor
• Multivariate UQ + cross-conditioning + inverse design
• Piecewise-linear surrogates for MILP embedding
• Scaling: what a full-depth soft tree costs
• Failure modes: seeds, local inverse design, missing epistemic uncertainty, readability

Reference

• API reference
  • Estimators
  • Export
• References