Wasserstein Gradient Boosting: A Framework for Distribution-Valued Supervised Learning (2405.09536v2)

Abstract: Gradient boosting is a sequential ensemble method that fits a new weaker learner to pseudo residuals at each iteration. We propose Wasserstein gradient boosting, a novel extension of gradient boosting that fits a new weak learner to alternative pseudo residuals that are Wasserstein gradients of loss functionals of probability distributions assigned at each input. It solves distribution-valued supervised learning, where the output values of the training dataset are probability distributions for each input. In classification and regression, a model typically returns, for each input, a point estimate of a parameter of a noise distribution specified for a response variable, such as the class probability parameter of a categorical distribution specified for a response label. A main application of Wasserstein gradient boosting in this paper is tree-based evidential learning, which returns a distributional estimate of the response parameter for each input. We empirically demonstrate the superior performance of the probabilistic prediction by Wasserstein gradient boosting in comparison with existing uncertainty quantification methods.

• The paper introduces WGBoost, a novel extension of gradient boosting that uses Wasserstein gradients to predict full probability distributions.
• It demonstrates robust performance in conditional density estimation and out-of-distribution detection, excelling on metrics like NLL and RMSE.
• The approach enhances uncertainty quantification, offering more reliable predictions for applications in high-stakes domains such as medical diagnostics and autonomous driving.

Introducing Wasserstein Gradient Boosting: Enhancing Predictive Uncertainty in Gradient Boosting

Gradient Boosting is a popular machine learning method, especially useful with tabular data. However, traditional gradient boosting techniques often focus on point predictions or probabilistic classification, with less attention given to capturing predictive uncertainty. This is crucial for fields like medical diagnostics and autonomous driving where assessing risks and predictions' uncertainty can make a huge difference.

What is Wasserstein Gradient Boosting?

The paper presents a new technique called Wasserstein Gradient Boosting (WGBoost). This is an extension of gradient boosting that fits new base learners (typically decision trees) to the Wasserstein gradient of a loss function over probability distributions. Simply put, WGBoost aims to better approximate and predict entire probability distributions rather than just point estimates.

This approach is useful for "posterior regression," where the goal is to model the distribution of a parameter given past data.

Key Highlights

General Methodology

WGBoost builds on gradient boosting by:

1. Introducing a loss functional that measures the divergence between a predicted distribution and a target distribution.
2. Training base learners to approximate the steepest descent direction (Wasserstein gradient) of this functional.

The algorithm outputs a set of particles that approximate the target distribution at each input. This is particularly fitting for applications requiring high predictive uncertainty.

Numerical Results

The paper provides comprehensive empirical results:

1. Conditional Density Estimation: WGBoost effectively captures the variability in data, even with complex distribution shapes.
2. Probabilistic Regression Benchmarking: WGBoost often matches or exceeds the performance of other state-of-the-art methods across a variety of datasets, particularly in terms of negative log likelihood (NLL) and root mean square error (RMSE).
3. Classification and Out-of-Distribution (OOD) Detection: WGBoost demonstrates strong classification accuracy while also excelling in OOD detection, a critical capability for identifying when an input sample markedly deviates from the training data.

Practical and Theoretical Implications

Practical Implications

The key benefit of WGBoost is its ability to provide a distributional prediction rather than a single point estimate. This improvement offers:

• Enhanced robustness: Predictions take into account distributional data, offering more reliable outputs.
• Improved uncertainty estimates: Beneficial for fields where understanding the confidence of predictions is critical (e.g., medical applications).

Theoretical Implications

WGBoost extends gradient boosting by incorporating Wasserstein gradients, thereby opening new avenues to utilize strong mathematical frameworks from optimal transport in machine learning. This can inspire further research in:

• Advanced loss functionals: Tailoring them to specific applications.
• Cross-disciplinary applications: Using WGBoost in fields like computational finance, climate modeling, and more where uncertainty quantification is vital.

Future Developments

Moving forward, WGBoost could see enhancements such as:

• Hybrid models that integrate other machine learning paradigms.
• Further scalability improvements to handle larger datasets seamlessly.
• Expansions in automated machine learning (AutoML) frameworks to leverage WGBoost without manual tuning.

Overall, the paper presents a compelling case for the adoption of Wasserstein Gradient Boosting, highlighting its strengths and potential for future enhancements in predictive modeling. Whether in academia or industry, WGBoost offers a promising pathway to more reliable and interpretable machine learning models.