Conditioning Tree-Based Diffusions and Flows for Probabilistic Tabular Regression

Technical Summary: Conditioning Tree-Based Diffusions and Flows for Probabilistic Tabular Regression

Problem Statement
Probabilistic regression on tabular data requires models that combine the predictive strength of gradient-boosted trees with full predictive distributions suitable for calibration, decision-making, and uncertainty quantification. While tree-based diffusion models (e.g., Treeffuser) have emerged as a solution by training LightGBM regressors to estimate the score of a noise-perturbed conditional density, they currently inherit design defaults from neural diffusion settings. These defaults—including the noising path, parameterization, training distribution, features, and sampler—are often suboptimal for histogram-based tree learners. The paper argues that these defaults act as a binding constraint: the specific supervised regression problem a gradient-boosted ensemble solves is determined by how these conditioning choices are configured.

Methodology: DiffGBM
The author introduces DiffGBM, an extension of Treeffuser that treats the conditioning choices of tree-based diffusion as explicit, tunable axes rather than fixed defaults. The methodology is developed along two primary axes:

1. Score-Side Modifications (The "Score-Flex" Space):
   The author proposes a jointly tunable recipe for score-based diffusion over a shared LightGBM surface. Key modifications include:

   • Conditional-Mean Residualization: A cross-validated LightGBM mean predictor $\hat{\mu}(x)$ is fitted first. Diffusion is applied to the residual $y - \hat{\mu}(x)$ rather than $y$, moving conditional location out of the score-regression target.
   • EDM Preconditioning: Following EDM principles, the model predicts a denoised target $D_\theta(y_t, t)$ scaled to unit variance, reconstructing the score algebraically. This stabilizes the regression target for trees, preventing the need to extrapolate noise-prediction magnitudes across noise levels.
   • Log-Sigma Time Sampling & Features: Instead of uniform time sampling, the training distribution draws $\log \sigma(t) \sim \mathcal{N}(-1.2, 1.2^2)$. Crucially, an explicit $\log \sigma(t)$ feature is added to the input vector, allowing the tree splits to condition directly on the noise regime.
   • Loss Weighting and Histogram Resolution: These are exposed as tunable dimensions to control where the model's capacity is spent.

2. Gaussian-Path Flow Matching (FM):
   The author implements a Gaussian-path flow-matching trainer for $p(y|x)$.

   • Velocity Field Learning: Instead of learning a score directly, the model learns a velocity field $v_\theta$ via gradient-boosted trees.
   • Algebraic Score Recovery: For any Gaussian path, the conditional score can be recovered algebraically from the velocity (Lemma 1), allowing the model to enter stochastic-interpolant samplers if desired.
   • Sampler Efficiency: This approach enables the use of deterministic, few-step ODE samplers (specifically a 5-step Heun ODE), significantly reducing inference latency compared to stochastic SDEs.
   • Path Choices: The author compares linear, trigonometric, and variance-preserving (VP) paths. They note that while VP paths keep inputs near unit scale, the optimal path depends on the dataset and the interplay with residualization.

Key Contributions

1. Flow Matching for Conditional Tree-Based Prediction: A Gaussian-path FM trainer that learns a velocity field directly with gradient-boosted trees, offering the cheapest sampler and best-calibrated operating point among DiffGBM variants.
2. A Jointly Tunable Score Recipe: The exposure of score-side modifications (residualization, EDM preconditioning, log-sigma sampling, explicit noise features, loss weighting, and histogram resolution) as tunable axes. This "score-flex" space represents the published recipe as a special case but allows for dataset-specific optimization.
3. Regression Conditioning as the Shared Bottleneck: The paper posits that the performance gains in both axes stem from better matching the induced supervised problem to the finite split resolution of histogram trees. The constraint is not the population-level properties of the objectives but the finite bin allocation required to serve all noise levels simultaneously.

Experimental Results
The method was evaluated on eleven standard tabular regression benchmarks (nine UCI, two scikit-learn) using a nested tuning protocol (fold 0 for tuning, folds 1–5 for evaluation) with a matched 40-trial budget.

• Accuracy (Score-Flex): The jointly tuned score-flex recipe achieved the best aggregate CRPS skill score (0.725 vs. 0.699 for the published baseline) and normalized CRPS (1.106 vs. 1.248). It beat the published baseline on every single dataset (paired Wilcoxon $p < 10^{-3}$).
• Calibration and Speed (Flow Matching): The DiffGBM-FM variant was the best-calibrated DiffGBM row (lowest coverage error at 50/90% levels and highest PIT pass rate) and the cheapest sampler. It was 5.2× faster than the published baseline (61.3s total vs. 196.3s) while maintaining competitive CRPS (0.707).
• Ablation Studies:
  • Switching from Euler SDE to Heun ODE for the score model without other changes reduced accuracy, indicating the sampler is not a free speedup.
  • EDM preconditioning provided the largest single jump in score-side performance.
  • Residualization was critical for Flow Matching; without it, FM performance degraded significantly.
  • Stochastic samplers ($\epsilon > 0$) did not Pareto-dominate deterministic samplers ($\epsilon = 0$) on aggregate CRPS or latency, though they could improve tail coverage in high-uncertainty regions.

Significance and Claims
The paper claims that the design defaults inherited from neural diffusion settings are the primary bottleneck for tree-based probabilistic regression. By treating the noising path, parameterization, and training distribution as tunable conditioning choices, DiffGBM unlocks significant improvements in both accuracy and calibration.

The author emphasizes that the "score-flex" and "flow-matching" approaches are complementary:

• Score-Flex is the accuracy choice, dominating as data grows and winning on the largest datasets, though it incurs higher latency and looser central coverage due to its stochastic sampler.
• Flow Matching is the fast, well-calibrated choice, winning on small-to-medium datasets where its residualized few-step ODE is strongest.

The paper concludes that no single recipe axis is universal; the optimal configuration depends on the dataset's characteristics (e.g., near-deterministic datasets like ct_slices preferred finer histogram bins and noise prediction over EDM). The work establishes that conditioning the induced supervised problem to the finite resolution of histogram splits is the key to unlocking the potential of tree-based diffusion and flow models.