# Wasserstein Gradient Boosting: A Framework for Distribution-Valued Supervised Learning

arXiv: https://arxiv.org/abs/2405.09536 · HTML: https://mlsanity.com/paper/2405.09536/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 68 references that are in the map.
- First posted on arXiv 15 May 2024.

## Builds on (most cited first, 30 of 34)

- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges](https://mlsanity.com/paper/2011.06225/) (2020)
- [Generalized Out-of-Distribution Detection: A Survey](https://mlsanity.com/paper/2110.11334/) (2021)
- [Evidential Deep Learning to Quantify Classification Uncertainty](https://mlsanity.com/paper/1806.01768/) (2018)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [Tabular Data: Deep Learning is Not All You Need](https://mlsanity.com/paper/2106.03253/) (2021)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [A Survey of Deep Learning Techniques for Autonomous Driving](https://mlsanity.com/paper/1910.07738/) (2019)
- [Deep Evidential Regression](https://mlsanity.com/paper/1910.02600/) (2019)
- [Probabilistic Backpropagation for Scalable Learning of Bayesian Neural Networks](https://mlsanity.com/paper/1502.05336/) (2015)
- [Concrete Dropout](https://mlsanity.com/paper/1705.07832/) (2017)
- [Posterior Network: Uncertainty Estimation without OOD Samples via Density-Based Pseudo-Counts](https://mlsanity.com/paper/2006.09239/) (2020)
- [Ensemble Distribution Distillation](https://mlsanity.com/paper/1905.00076/) (2019)
- [Stein Variational Gradient Descent as Gradient Flow](https://mlsanity.com/paper/1704.07520/) (2017)
- [Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/1906.04370/) (2019)
- [Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem](https://mlsanity.com/paper/1802.08089/) (2018)
- [Subspace Inference for Bayesian Deep Learning](https://mlsanity.com/paper/1907.07504/) (2019)
- [NGBoost: Natural Gradient Boosting for Probabilistic Prediction](https://mlsanity.com/paper/1910.03225/) (2019)
- [Interacting particle solutions of Fokker-Planck equations through gradient-log-density estimation](https://mlsanity.com/paper/2006.00702/) (2020)
- [Prior and Posterior Networks: A Survey on Evidential Deep Learning Methods For Uncertainty Estimation](https://mlsanity.com/paper/2110.03051/) (2021)
- [SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergence](https://mlsanity.com/paper/2006.02509/) (2020)
- [A Non-Asymptotic Analysis for Stein Variational Gradient Descent](https://mlsanity.com/paper/2006.09797/) (2020)
- [A Stein variational Newton method](https://mlsanity.com/paper/1806.03085/) (2018)
- [Stein Variational Model Predictive Control](https://mlsanity.com/paper/2011.07641/) (2020)
- [Information Newton's flow: second-order optimization method in probability space](https://mlsanity.com/paper/2001.04341/) (2020)
- [Projected Wasserstein gradient descent for high-dimensional Bayesian inference](https://mlsanity.com/paper/2102.06350/) (2021)
- [DEAL: Deep Evidential Active Learning for Image Classification](https://mlsanity.com/paper/2007.11344/) (2020)
- [Regularized Stein Variational Gradient Flow](https://mlsanity.com/paper/2211.07861/) (2022)

## Built on by (most cited first, 4 of 4)

- [DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows](https://mlsanity.com/paper/2503.01140/) (2025)
- [Multivariate Uncertainty Quantification with Tomographic Quantile Forests](https://mlsanity.com/paper/2512.16383/) (2025)
- [Fr'echet Geodesic Boosting](https://mlsanity.com/paper/2509.18013/) (2025)
- [Odds-Shift Slippage in One-vs-Rest Rankers: Diagnosing and Repairing Reweighting-Induced Top-K Errors](https://mlsanity.com/paper/2609.13810/) (2026)

## Abstract (arXiv)

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.