Uncertainty-Aware Extrapolation in Bayesian Oblique Trees

TL;DR

This work extends variational oblique predictive clustering trees by attaching Gaussian process leaves, enabling uncertainty-aware extrapolation within a single, interpretable Bayesian tree. A gating mechanism based on leaf training support ensures GP-based extrapolation is activated only when inputs fall outside the leaf, preserving in-distribution accuracy. Theoretical results decompose predictive uncertainty into routing and functional components and establish that the model generalizes existing tree approaches while enabling nonlinear leaf behavior. Empirically, VSPYCT-GP achieves competitive IID performance and substantial improvements in extrapolation tasks, with tunable trade-offs via the support threshold $\tau$. The approach offers calibrated uncertainty under distribution shift and is practically beneficial for risk-sensitive regression applications.

Abstract

Decision trees are widely used due to their interpretability and efficiency, but they struggle in regression tasks that require reliable extrapolation and well-calibrated uncertainty. Piecewise-constant leaf predictions are bounded by the training targets and often become overconfident under distribution shift. We propose a single-tree Bayesian model that extends VSPYCT by equipping each leaf with a GP predictor. Bayesian oblique splits provide uncertainty-aware partitioning of the input space, while GP leaves model local functional behaviour and enable principled extrapolation beyond the observed target range. We present an efficient inference and prediction scheme that combines posterior sampling of split parameters with \gls{gp} posterior predictions, and a gating mechanism that activates GP-based extrapolation when inputs fall outside the training support of a leaf. Experiments on benchmark regression tasks show improvements in the predictive performance compared to standard variational oblique trees, and substantial performance gains in extrapolation scenarios.

Figures (5)

• Figure 1: Extrapolation behaviour on synthetic 1D regression data. The true function $f(x)$ (dashed gray) exhibits a smooth upward trend. Training data (blue dots) are observed only within a bounded region. In the training region, both VSPYCT (violet) and VSPYCT-GP (teal) produce similar smooth predictions due to Monte Carlo averaging over probabilistic split parameters. In extrapolation regions, VSPYCT predictions flatten---bounded by the range of leaf prototypes---while VSPYCT-GP continues the learned trend with confidence intervals that widen to reflect increasing uncertainty. This illustrates how GP leaf models enable calibrated extrapolation while preserving in-distribution performance.
• Figure 2: Model overview comparing VSPYCT and VSPYCT-GP. (a) Standard VSPYCT uses Bayesian oblique splits with constant leaf prototypes $\mu_\ell$, bounding predictions to the training target range. (b) VSPYCT-GP replaces prototypes with Gaussian process predictors that provide both a predictive mean and uncertainty estimate. The prediction pipeline (bottom) shows how split parameters are sampled from their variational posteriors, instances are routed to leaves, and a gating mechanism selects between prototype and GP predictions based on support detection.
• Figure 3: Extrapolation-aware gating mechanism. (a) The support region of a leaf is defined by the Mahalanobis distance from the training data centroid $\bar{x}_\ell$, using the empirical covariance $\Sigma_\ell$. Test points inside the support ($x_{\mathrm{in}}$) receive predictions dominated by the prototype, while points outside ($x_{\mathrm{out}}$) are governed by the GP predictor. (b) Prediction behavior: within support, predictions approach the constant prototype with minimal variance; outside support, the GP posterior dominates, providing both a predictive mean that can extrapolate and uncertainty that grows with distance. Soft sigmoid gating ensures a smooth transition at the support boundary.
• Figure 4: Effect of support threshold $\tau$ on VSPYCT-GP performance. (a) RMSE as a function of $\tau$ for three datasets; dashed lines indicate baseline VSPYCT performance. (b) Mean predictive uncertainty decreases as $\tau$ increases, reflecting fewer out-of-support classifications. Shaded regions show standard deviation across multiple random seeds.
• Figure 5: Controlled interpolation vs. extrapolation experiment on synthetic data. (a) RMSE comparison: both models perform similarly for interpolation, but VSPYCT-GP substantially outperforms VSPYCT in extrapolation regimes. (b) VSPYCT-GP uncertainty increases appropriately as test data moves further from the training distribution. Error bars show standard deviation across five random seeds.