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

TL;DR

This paper introduces Wasserstein gradient boosting, a new method for distribution-valued supervised learning that improves probabilistic predictions by fitting models to Wasserstein gradients of loss functionals, demonstrating superior uncertainty quantification.

Contribution

It extends gradient boosting to handle distribution-valued outputs using Wasserstein gradients, enabling more accurate probabilistic modeling and uncertainty quantification.

Findings

Outperforms existing uncertainty quantification methods in experiments

Provides distributional estimates of response parameters for each input

Enhances tree-based evidential learning with superior probabilistic predictions

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.