Variational Boosted Soft Trees

TL;DR

This paper introduces a differentiable Bayesian gradient boosting method using soft decision trees, enhancing uncertainty estimation while maintaining strong predictive performance on tabular data.

Contribution

It presents a novel variational inference approach with soft decision trees for Bayesian GBMs, enabling better uncertainty calibration and computational efficiency.

Findings

Higher test likelihoods on 7/10 regression datasets

Improved out-of-distribution detection in 5/10 datasets

Provides useful uncertainty estimates with maintained accuracy

Abstract

Gradient boosting machines (GBMs) based on decision trees consistently demonstrate state-of-the-art results on regression and classification tasks with tabular data, often outperforming deep neural networks. However, these models do not provide well-calibrated predictive uncertainties, which prevents their use for decision making in high-risk applications. The Bayesian treatment is known to improve predictive uncertainty calibration, but previously proposed Bayesian GBM methods are either computationally expensive, or resort to crude approximations. Variational inference is often used to implement Bayesian neural networks, but is difficult to apply to GBMs, because the decision trees used as weak learners are non-differentiable. In this paper, we propose to implement Bayesian GBMs using variational inference with soft decision trees, a fully differentiable alternative to standard decision trees introduced by Irsoy et al. Our experiments demonstrate that variational soft trees and variational soft GBMs provide useful uncertainty estimates, while retaining good predictive performance. The proposed models show higher test likelihoods when compared to the state-of-the-art Bayesian GBMs in 7/10 tabular regression datasets and improved out-of-distribution detection in 5/10 datasets.