Title: Wasserstein Gradient Boosting: A Framework for Distribution-Valued Supervised Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Gradient Boosting; Wasserstein Gradient Flow; Uncertainty Quantification
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, which 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. 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 competitive performance of the probabilistic prediction by Wasserstein gradient boosting in comparison with existing uncertainty quantification methods.

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

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, which 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. 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 competitive performance of the probabilistic prediction by Wasserstein gradient boosting in comparison with existing uncertainty quantification methods.

1 Introduction

Gradient boosting is a celebrated machine learning algorithm that has achieved considerable success with tabular data [1]. Gradient boosting has been extensively used for point forecasts and probabilistic classification, yet a relatively small number of studies have been concerned with the predictive uncertainty of gradient boosting. Predictive uncertainty of machine learning models plays a growing role in today's real-world production systems [2]. It is vital for safety-critical systems, such as medical diagnoses [3] and autonomous driving [4], to assess the potential risk of their actions that partially or entirely rely on predictions from their models. Gradient boosting has already been applied in a diverse range of real-world applications, including click prediction [5], ranking systems [6], scientific discovery [7], and data competition [8]. There is a pressing need for methodology to harness the power of gradient boosting to predictive uncertainty quantification.

In classification and regression, we typically specify a noise distribution p ( y | θ ) of a response variable y and use a model to return a point estimate θ ( x ) of the response parameter for each input x . In recent years, the importance of capturing uncertainty in the model output θ ( x ) has increasingly been emphasised [2]. A variety of approaches have been proposed to obtain a distributional estimate p ( θ | x ) of the response parameter for each input x [e.g. 9, 10, 11]. For example, Bayesian neural networks (BNNs) quantify uncertainty in network weights and propagate it to the space of network outputs. Marginalising the predictive distribution p ( y | θ ) over the distributional estimate p ( θ | x ) has been demonstrated to confer enhanced predictive accuracy and robustness against adversarial attacks [11]. Furthermore, the dispersion of the distributional estimate has been used as a powerful indicator for out-of-distribution (OOD) detection [12].

(a)

0

weak learner trained

(b)

15

weak learners trained

(c)

100

weak learners trained

2

2

2

1

1

1

0

0

0

-1

-1

1

2

2

-2

-3

-2

-1

0

1

2

3

E-

2

-1

0

1

2

3

E-

2

-1

0

1

2

3

In this context, a line of research based on the concept of evidential learning has recently gained significant attention [11, 13, 14, 15]. The idea can be broadly interpreted as making use of the 'individual-level' posterior p ( θ | y i ) of the response parameter θ conditional on each individual datum y i , which arises from the response-distribution likelihood p ( y i | θ ) and a user-specified prior p ( θ ) . If each individual-level posterior falls into a closed form characterised by some hyperparameter, neural networks can be trained by using the finite-dimensional hyperparameter as a target value for each input. Outstanding performance and computational efficiency of the existing approaches have been delivered in a wide spectrum of engineering and medical applications [16, 17, 18, 19]. However, the existing approaches are limited to neural networks and to the case where every individual-level posterior is in closed form so that the finite-dimensional hyperparameter can be predicted by proxy. In general, posterior distributions are known only up to their normalising constants and, therefore, require an approximation typically by particles [20].

Without closed-form expression, each individual-level posterior needs to be treated as an infinitedimensional output for each input. This challenge poses the following fundamental question:

Consider a supervised learning setting whose outputs are probability distributions. Given a training set of input values and output distributions { x i , µ i } D i =1 , can we build a model that receives an input x and returns a nonparametric prediction of the output distribution?

Motivated by this question, we propose a general framework of Wasserstein gradient boosting (WGBoost). WGBoost receives an input and returns a particle approximation of the output distribution. Figure 1 illustrates inputs and outputs of WGBoost. In this paper, we focus on application of WGBoost to evidential learning, where the individual-level posterior p ( θ | y i ) of the response parameter θ is used as the output distribution µ i for each input x i in the training set. Figure 2 compares the pipeline of evidential learning based on WGBoost with that of Bayesian learning.

Contributions Our contributions are summarised as follows:

Methodology of WGBoost : Section 2 establishes the general framework of WGBoost. It is a novel family of gradient boosting that returns a set of particles that approximates an output distribution assigned at each input. In contrast to standard gradient boosting that fits a weak learner to the gradient of a loss function, WGBoost fits a weak learner to the estimated Wasserstein gradient of a loss functional over probability distributions.

Application to Evidential Learning : Section 3 establishes tree-based evidential learning based on WGBoost, with the loss functional specified by the Kullback-Leibler (KL) divergence. Following modern gradient-boosting libraries [21, 22] that uses second-order gradient boosting (c.f. Section 2.2), we implement a concrete second-order WGBoost algorithm built on an approximate Wasserstein gradient and Hessian of the KL divergence.

Experiment on Real-world Data : Section 4 demonstrates the performance of probabilistic regression, and classification with OOD detection, on real-world tabular datasets. To the author's knowledge, WGBoost is the first framework that enables evidential learning for (i) boosted tree models and (ii) cases without closed form of individual-level posteriors.

θ =f(x,w)

p(y 1 0)

(0(x),...,0N(x))

p(y 1 0)

C

f(α,w)

y

C

WGBoost

A

y

propagate

training

w

～ p(w| {ci,yi}1)

[2i, p(0 | yi)}²=

2 Wasserstein Gradient Boosting

This section establishes the general formulation of WGBoost. Section 2.1 recaps the notion of Wasserstein gradient flows, a 'gradient' system of probability distributions that minimises an objective functional in the space of probability distributions. Section 2.2 recaps the notion of gradient boosting, a sequential ensemble method that fits a new weak learner to the 'gradient' of the remaining loss at each iteration. Section 2.3 combines the above two notions to establish WGBoost, a novel family of gradient boosting that enables to solve distribution-valued supervised learning.

Notation and Setting Let X and Y denote the space of inputs and responses in classification and regression. Suppose Θ = R d . Let P 2 be the 2-Wasserstein space i.e. a set of all probability distributions on Θ with finite second moment equipped with the Wasserstein metric [23]. We identify a probability distribution in P 2 with its density whenever it exits. Denote by ⊙ and ⊘ , respectively, elementwise multiplication and elementwise division of two vectors in R d . Let ∇ be the gradient operator. Let ∇ 2 d be a second-order gradient operator that takes the second derivative at each coordinate i.e. ∇ 2 d f ( θ ) = [ ∂ 2 f ( θ ) /∂θ 2 1 , . . . , ∂ 2 f ( θ ) /∂θ 2 d ] T ∈ R d .

2.1 Wasserstein Gradient Flow

In the Euclidean space, a gradient flow of a function f means a curve of points x t that solves a differential equation ( d/dt ) x t = -∇ f ( x t ) from some initial value x 0 . That is the continuous-time limit of gradient descent, which minimises the function f as t →∞ . A Wasserstein gradient flow means a curve of probability distributions µ t minimising a given functional F on the 2-Wasserstein space P 2 from some initial distribution µ 0 . The Wasserstein gradient flow µ t is characterised as a solution of a partial differential equation, known as the continuity equation :

where ∇ W F ( µ ) : Θ → Θ denotes the Wasserstein gradient of F at µ [24, 25]. Appendix A recaps the derivation of the Wasserstein gradient, presenting the examples for several functionals.

One of the elegant properties of the Wasserstein gradient flow is casting the infinite-dimensional optimisation of the functional F as a finite-dimensional particle update [23]. The continuity equation (1) can be reformulated as a dynamical system of a random variable θ t ∼ µ t , such that

in the sense that the law µ t of such a random variable θ t is a weak solution of the continuity equation. Consider the case where the initial measure µ 0 is set to the empirical distribution ˆ µ 0 of N particles { θ n 0 } N n =1 . Discretising the continuous-time system (2) by the Euler method with a small step size ν > 0 yields an iterative update scheme of N particles { θ n m } N n =1 from step m = 0 :

where ˆ µ m denotes the empirical distribution of the particles { θ n m } N n =1 at step m .

In practice, it is common that the Wasserstein gradient of a chosen functional F is not well-defined for empirical distributions. In such case, the particle update scheme (3) is not directly applicable because it depends on the Wasserstein gradient ∇ W F (ˆ µ m ) at the empirical distribution ˆ µ m . For example, the KL divergence F ( µ ) = KL( µ | π ) with a reference distribution π leads to the Wasserstein gradient [ ∇ W F ( µ )]( θ ) = -( ∇ log π ( θ ) -∇ log µ ( θ )) ill-defined if µ is an empirical distribution. Hence, the particle update scheme (3) is often performed with the estimated or approximated Wasserstein gradient well-defined for empirical distributions [e.g. 26, 27, 28, 29, 30]. A main application of WGBoost in Section 3 uses the 'smoothed' Wasserstein gradient of the KL divergence [26].

2.2 Gradient Boosting

Gradient boosting [31] is a sequential ensemble method of M weak learners f 1 , . . . , f M . It iteratively constructs an ensemble F m of m weak learners f 1 , . . . , f m from step m = 0 to M . Given the current ensemble F m at step m , it trains a new weak learner f m +1 to construct the next ensemble by

where ν is a shrinkage hyperparameter called a learning rate . The initial state of the ensemble F 0 ( x ) at step m = 0 is typically set to a constant that best fits the data. Any learning algorithm can be used as a weak learner in principle, although tree-based algorithms are most used [32].

The fundamental idea of gradient boosting is to train the new weak learner f m +1 to approximate the negative gradient of the remaining error of the current ensemble F m . Suppose that a loss function L measures the remaining error R i ( F m ( x i )) := L ( F m ( x i ) , y i ) for each output vector y i ∈ R d . The new weak learner f m +1 is fitted to the set { x i , g i } D i =1 whose target variable g i is each specified by

The target g i is often called a pseudo residual . For each data input x i , the boosting scheme (4) updates the output of the current ensemble F m ( x i ) in the steepest descent direction of the error R i ( F m ( x i )) . Although [31] originally suggested an additional line search to determine a scaling constant of each weak learner, the line search has been reported to have a negligible influence on performance [33].

In modern gradient-boosting libraries, such as XGBoost [21] and LightGBM [22], the standard practice is to use the diagonal (coordinatewise) Newton direction of the remaining error R i ( F m ( x i )) in lieu of the negative gradient g i . The new base leaner f m +1 is instead fitted to the set { x i , g i ⊘ h i } n i =1 , where the negative gradient g i is divided elementwise by the Hessian diagonal h i given by

The target variable g i ⊘ h i is the diagonal Newton direction that minimises the second-order Taylor approximation of the remaining error R i ( F m ( x i )) for each coordinate independently. Combining the second-order gradient boosting framework with tree-based weak learners has demonstrated exceptional scalability and performance [34, 35]. Although it is possible to use the 'full' Newton direction as the target variable of each weak learner, the impracticality of the full Newton direction has been pointed out [e.g. 36, 37]. In addition, the coordinatewise computability of the diagonal Newton direction is suitable for popular gradient-boosting tree algorithms [36].

2.3 General Formulation of Wasserstein Gradient Boosting

Now we consider the setting of distribution-valued supervised learning, where we are given a training set of input vectors and output distributions { x i , µ i } D i =1 ⊂ X × P 2 . Our goal is to construct a model that receives an input and returns a set of N particles whose empirical distribution approximates the output distribution. We specify a loss functional D( · | · ) between two probability distributions-such as the KL divergence-to measure the remaining error F i ( · ) = D( · | µ i ) for each i -th training output distribution µ i . Our idea is to combine gradient boosting with Wasserstein gradient, where we iteratively construct a set of N boosting ensembles F 1 m , . . . , F N m from step m = 0 to M .

Here, the output F n m ( x ) of each n -th boosting ensemble represents the n -th output particle for an input x . Given the current set of N ensembles F 1 m , . . . , F N m at step m , WGBoost trains a set of N

new weak learners f 1 m +1 , . . . , f N m +1 and computes the next set of N ensembles by

where ν is a learning rate. Similarly to standard gradient boosting, we specify the initial state of N ensembles F 1 0 , . . . , F N 0 at step m = 0 by a set of constants. Throughout, denote by ˆ µ m,i the empirical distribution of the N output particles F 1 m ( x i ) , . . . , F N m ( x i ) for each i -th training input x i .

As discussed in Section 2.1, the Wasserstein gradient often needs to be estimated for empirical distributions. For better presentation, let G i ( µ ) denote an estimate of the Wasserstein gradient ∇ W F i ( µ ) of the i -th remaining error F i ( µ ) , which is well-defined for any distribution µ . If the Wasserstein gradient ∇ W F i ( µ ) is originally well-defined for any distribution µ , it is a trivial choice of the estimate, i.e., G i ( µ ) = ∇ W F i ( µ ) . Otherwise, any suitable estimate can be used as G i ( µ ) . The foundamental idea of WGBoost is to train the n -th new learner f n m +1 to approximate the estimated Wasserstein gradient -G i (ˆ µ m,i ) evaluated at the n -th boosting output F n m ( x i ) for each x i , so that,

For each data input x i , the boosting scheme (5) approximates the particle update scheme (3) for the output particles F 1 m ( x i ) , . . . , F N m ( x i ) under the estimated Wasserstein gradient. The output particles are updated in the direction to decrease the remianing error F i (ˆ µ m,i ) = D(ˆ µ m,i | µ i ) at each step m .

Algorithm 1 summarises the general procedure of WGBoost. See Figure 1 for illustration of WGBoost. In Section 3, we choose the KL divergence as a loss functional D and use a kernel smoothing estimate of the Wasserstein gradient. See Appendix A for the Wasserstein gradient of other divergences.

Remark 1 ( Stochastic WGBoost ) . Stochastic gradient boosting [38] uses only a randomly sampled subset of data to fit a new weak learner at each step m to reduce the computational cost. The same subsampling approach can be applied for WGBoost whenever the dataset is large.

Remark 2 ( Second-Order WGBoost ) . If any estimate of the Wasserstein 'Hessian' of the remaining error F i is available, the Newton direction of F i may also be computable [e.g. 39, 40]. Implemention of a second-order WGBoost algorithm is immediate by plugging such a Newton direction into G i ( µ ) in Algorithm 1. Our default WGBoost algorithm for tree-based evidential learning is built on a diagonal approximate Newton direction of the KL divergence, aligning with the standard practice in modern gradient-boosting libraries to use the diagonal Newton direction.

Algorithm 1: Wasserstein Gradient Boosting

Input: training set { x i , µ i } D i =1 of input x i ∈ X and output distribution µ i ∈ P 2

Parameter:

loss

N

Output:

D

, estimate

, iteration set of

(

F

1

0

(

·

)

, . . . , F

for

m

G

i

(

µ

)

M

, learning rate

N

ν

, weak learner

∇

W

f

boosting ensembles

N

0

0

(

·

))

(

F

(

ϑ

1

1

0

, . . . , ϑ

N

0

)

, . . . , M

-

1

do

M

, . . . , F

N

M

)

D(

µ

|

µ

i

, initial constants at final step

▷

)

, particle number

(

ϑ

1

0

, . . . , ϑ

N

0

)

M

set initial state of

N

boosting ensembles

3 Application to Evidential Learning

This section provides our default setting to implement a concrete WGBoost algorithm for evidential learning, which enables classification and regression with predictive uncertainty. The individual-level posterior p ( θ | y i ) of a response distribution p ( y | θ ) is used as the output distribution µ i of the training set { x i , µ i } D i =1 . Section 3.1 recaps derivation of the individual-level posterior p ( θ | y i ) , followed by Section 3.2 discussing the default choice of the prior. We choose the KL divergence as a loss functional of WGBoost. Section 3.3 recaps a widely-used estimate of the Wasserstein gradient of the KL divergence based on kernel smoothing [26]. A further advantage of the kernel smoothing estimate is that the approximate Wasserstein Hessian is available, with which Section 3.4 establishes a second-order WGBoost algorithm similarly to modern gradient-boosting libraries.

3.1 Derivation of Individual-Level Posteriors and Predictive Distribution

Suppose that a response distribution p ( y | θ ) of a response variable y is specified, as is typically done for probabilistic prediction. Suppose also that a prior distribution p i ( θ ) of the response parameter θ is specified for each individual data input x i . For each individual data pair ( x i , y i ) , the responsedistribution likelihood p ( y i | θ ) and the prior p i ( θ ) determine the individual-level posterior

by Bayes' theorem. This individual-level posterior is set to the output distribution µ i of the training set { x i , µ i } D i =1 of WGBoost. The framework of WGBoost then constructs a model that returns a particle approximation of the output distribution µ i ( · ) = p ( · | y i ) for each data input x i .

For a new input x , the constructed WGBoost model provides a set of particles ( θ 1 ( x ) , . . . , θ N ( x )) as a distributional prediction p ( θ | x ) of the response parameter θ . We can define a predictive distribution p ( y | x ) of the response y for the new input x via marginalisation of the output particles:

We can also define a point prediction ˆ y for the new input x via the individual-level Bayes action ˆ y = argmin y ∈Y ∫ Θ U ( y, θ ) p ( θ | x ) dθ , which minimises the average of some error U : Y × Θ → R . For example, the Bayes action ˆ y is simply the mean of the output particles if U ( y, θ ) = ( y -θ ) 2 .

In general, the explicit form of the individual-level posterior p ( θ | y i ) is known only up to the normalising constant. Our full algorithm in Section 3.4 requires no normalising constant of the individual-level posterior p ( θ | y i ) . Our algorithm depends only on the log-gradient of the individuallevel posterior ∇ log p ( θ | y i ) that cancels any constant term by the gradient. Hence, knowing the form of the response-distribution likelihood p ( y i | θ ) and the prior p i ( θ ) suffices.

Remark 3 ( Difference from Bayesian Learning ) . Given a response distribution p ( y | θ ) and a model θ = f ( x, w ) with the parameter w , Bayesian learning of the model f means the use of the posterior p ( w | { x i , y i } D i =1 ) over w conditional on all data. The predictive distribution p ( y | x ) of the response y is defined via marginalisation over w : ∫ Θ p ( y | θ = f ( x, w )) p ( w | { x i , y i } D i =1 ) dw . In contrast, WGBoost returns a distributional prediction p ( θ | x ) of the response parameter θ , circumventing the marginalisation over the model parameter w that can be ultra-high dimensional.

3.2 Choice of Individual-Level Priors

The prior distribution p i ( θ ) of the response parameter θ is specified at each individual data input x i . The approach to eliciting the prior may differ, depending on whether past data are available. If past data are available, they can be utilised to elicit a reasonable prior for unobserved data. If no past data are available, we recommend the use of a noninformative prior that have been developed as a sensible choice of prior in the absence of past data; see [e.g. 41] for the introduction. To avoid numerical errors, if a noninformative prior is improper (i.e. nonintegrable), we recommend the use of a proper probability distribution that approximates the noninformative prior sufficiently well.

Example 1 ( Normal Location-Scale Response ) . Consider a scalar-valued response variable y ∈ R for regression. A normal location-scale response distribution N ( y | m,σ ) has the mean and scale parameters m ∈ R and σ ∈ (0 , ∞ ) . A typical noninformative prior of m and σ are given by,

respectively, 1 and 1 /σ which are improper. At every data point ( x i , y i ) , we use a normal prior N ( m | 0 , σ 0 ) over m and an inverse gamma prior IG( σ | α 0 , β 0 ) over σ , with the hyperparameters σ 0 = 10 and α 0 = β 0 = 0 . 01 , which approximate the non-informative priors.

Example 2 ( Categorical Response ) . Consider a label response variable y ∈ { 1 , . . . , k } for k -class classification. A categorical response distribution C ( y | q ) has the class probability parameter q = ( q 1 , . . . , q k ) in the k -dimensional simplex ∆ k . If k = 2 , it corresponds to the Bernoulli distribution. A typical noninformative prior of q is given by 1 / ( q 1 ×··· × q k ) which are improper. At every data point ( x i , y i ) , we use the logistic normal prior-a multivariate generalisation of the logit normal distribution [42]-over q with the mean 0 and identity covariance matrix scaled by 10 .

Remark 4 ( Reparametrisation and Standardisation ) . Section 2 supposed Θ = R d for some dimension d without no loss of generality. Any parameter that lies in a subset of the Euclidean space (e.g. σ in Example 1) can be reparametrised as one in the Euclidean space (e.g. log σ ). Appendix D details the reparametrisation used for Examples 1 and 2. In addition, if one's dataset has scalar outputs of a low or high order of magnitude, we recommend standardising the outputs.

3.3 Approximate Wasserstein Gradient of KL Divergence

We consider the KL divergence KL( µ | µ i ) as a loss functional of WGBoost. One challenge of the KL divergence is that the resulting Wasserstein gradient [ G KL i ( µ ) ] ( θ ) := -( ∇ log µ i ( θ ) -∇ log µ ( θ )) is not well-defined when µ is an empirical distribution. A particularly successful solution-which originates in [43] and has been applied in wide contexts [26, 44, 45]-is to smooth the Wasserstein gradient through a kernel integral operator ∫ Θ [ G KL i ( µ )]( θ ∗ ) k ( θ, θ ∗ ) dµ ( θ ∗ ) [46]. By integration-bypart (see [e.g. 43]), the smoothed Wasserstein gradient, denoted G ∗ i ( µ ) , falls into the following form that is well-defined for any distribution µ :

where ∇ k ( θ ∗ , θ ) denotes the gradient of k with respect to the first argument θ ∗ . An approximate Wasserstein gradient flow based on the smoothed Wasserstein gradient G ∗ i ( µ ) is called the Stein variational gradient descent [43] or kernelised Wasserstein gradient flow [47]. In most cases, the kernel k is set to the Gaussian kernel k ( θ, θ ∗ ) = exp( -∥ θ -θ ∗ ∥ 2 /h ) with the scale h > 0 . Appendix B discusses a choice of kernel. This work uses the Gaussian kernel with h = 0 . 1 throughout.

Another common approach to approximating the Wasserstein gradient flow of the KL divergence is the Langevin diffusion approach [48]. The discretised algorithm, called the unadjusted Langevin algorithm [49], is a stochastic particle update scheme that adds a Gaussian noise at every iteration. However, several known challenges, such as asymptotic bias and slow convergence, often necessitate an ad-hoc adjustment of the algorithm [48]. Appendix B discusses a variant of WGBoost built on the Langevin algorithm, although it is not considered the default implementation.

3.4 Second-Order Implementation of WGBoost

We use a diagonal (coordinatewise) approximate Wasserstein Newton direction of the KL divergence, following the standard practice in modern gradient-boosting libraries [21, 22] to use the diagonal Newton direction of a loss. Similarly to smoothed Wasserstein gradient G ∗ i ( µ ) , the approximate Wasserstein Hessian of the KL divergence KL( µ | µ i ) can be obtained through the kernel smoothing. The diagonal of the approximate Wasserstein Hessian, denoted H ∗ i ( µ ) , is defined by

The diagonal approximate Wasserstein Newton direction of the KL divergence is then defined by -[ G ∗ i ( µ )] ( · ) ⊘ [ H ∗ i ( µ )] ( · ) . Appendix C provides the derivation based on [39] who derived the Newton direction of the KL divergence in the context of nonparametric variational inference.

The second-order WGBoost algorithm is established by plugging it into G i ( µ ) in Algorithm 1, that is,

Algorithm 1 under the choice (9) is considered our default WGBoost algorithm for evidential learning. We refer this algorithm to as the Wasserstein-boosted evidential learning (WEvidential). The explicit pseudocode is provided in Algorithm 2 for full clarity.

Remark 5 ( Computation ) . The diagonal Newton direction (9) has the computational complexity O ( N × d ) same as that of the smoothed Wasserstein gradient. Hence, there is essentially no reason not to use the diagonal Newton direction (9) instead of the smoothed Wasserstein gradient. Although it is possible to use the full Newton direction with no diagonal approximation, the computation requires the inverse and product of ( N × d ) × ( N × d ) matrices that result in the complexity up to O ( N 3 × d 3 ) . Appendix D presents a simulation study to compare computational time and convergence speed of four WGBoost algorithms built on different estimates of the Wasserstein gradient.

4 Experiment on Real-world Tabular Data

We empirically demonstrate the performance of the WGBoost algorithm through three experiments using real-world tabular data. The first application illustrates the output of WGBoost through a simple conditional density estimation. The second application benchmarks the probabilistic regression performance. The third application demonstrates the classification and OOD detection performance. The source code is available in https://github.com/takuomatsubara/WGBoost .

Common Hyperparameters Throughout, we set the number of output particles N to 10 and set each weak learner f to the decision tree regressor [50] with maximum depth 1 for Section 4.1 and 3 for the rest. We set the learning rate ν to 0 . 1 for regression and 0 . 4 for classification, unless otherwise stated. Appendix E contains further details, including a choice of the initial constant { ϑ n 0 } N n =1 .

4.1 Illustrative Conditional Density Estimation

We illustrate the output of WEvidential by estimating a conditional density p ( y | x ) from onedimensional scalar inputs and outputs { x i , y i } D i =1 . The normal output distribution N ( y | m,σ ) and the prior p i ( m,σ ) in Example 1 were used to define the individual-level posterior p ( m,σ | y i ) , in which case the output of the WGBoost algorithm is a set of 10 particles { ( m n ( x ) , σ n ( x )) } 10 n =1 of the mean and scale parameters for each input x . We chose the number of weak learners M , drawing on an early-stopping approach used in [32], where we held out 20% of the training set as a validation set and chose the number 1 ≤ M ≤ 4000 achieving the least validation error. Once the number M was chosen, WEvidential was trained again using all the entire training set.

The conditional density is estimated using the predictive distribution (6) by WEvidential. We used two real-world datasets, bone mineral density [51] and old faithful geyser [52]. Figure 3 depicts the result for the former dataset, demonstrating that the WGBoost algorithm captures the heterogeneity of the conditional density on each input well. The result for the latter dataset is contained in Appendix E.1.

0.2

0.2

0.1

0.1

0.0

0.0

-0.1

-0.1

10

12

14

16

18

20

22

24

26

10

12

14

16

18

20

22

24

26

4.2 Probabilistic Regression Benchmark

We examine the regression performance of WEvidential using a standard benchmark protocol that originated in [53] and has been used in a number of subsequent works [10, 9, 32]. The benchmark protocol uses real-world tabular datasets from the UCI machine learning repository [54], each with one-dimensional scalar responses. As in Section 4.1, the normal response distribution N ( y | m,σ ) and the prior p i ( m,σ ) in Example 1 were used to define the individual-level posterior p ( m,σ | y i ) .

We followed the data splitting protocol in [53] and randomly held out 10% of each dataset as a test set. The negative log likelihood (NLL) is measured by using the predictive distribution (6). The root mean squared error (RMSE) is measured by using the point prediction by the mean value of the predictive distribution. We chose the number of weak learners M by the same approach as in Section 4.1. We repeated this procedure 20 times for each dataset, except the protein and year msd datasets for which we repeated five times and once. For the year msd dataset only, we subsampled 10% of data to fit each weak learner and used the learning rate 0.01 due to the large dataset size.

Table 1 compares the performance of WEvidential with five other methods: Monte Carlo Dropout (MCDropout) [9], Deep Ensemble (DEnsemble) [10], Concrete Dropout (CDropout) [55], Natural Gradient Boosting (NGBoost) [32], and Deep Evidential Regression (DEvidential) [13]. Appendix E provides further details on the experiment and a limited yet additional comparison. The WGBoost algorithm achieves the best score or a score sufficiently close to the best score most often.

4.3 Classification and Out-of-Distribution Detection

We examine the classification and anomaly OOD detection performance of WEvidential on two realworld tabular datasets, segment and sensorless , following the protocol used in [14]. The categorical response distribution C ( y | q ) and the prior p i ( q ) in Example 2 were used to define the individuallevel posterior p ( q | y i ) , in which case the output of the WGBoost algorithm is a set of 10 particles { q n } 10 n =1 of the class probability parameter q in the simplex ∆ k for each input x . We set the number of weak learners M to 4000 without early stopping to reduce the computational cost.

The segment and sensorless datasets have 7 and 11 classes in total. For the segment dataset, the data subset that belongs to the last class was kept as the OOD samples. For the sensorless dataset, the data subset that belongs to the last two classes was kept as the OOD samples. For each dataset, 20% of the non-OOD samples is held out as a test set to measure the classification accuracy. There exist several ways of defining a OOD score for each input [56]. For the WGBoost algorithm, the inverse of the maximum norm of the output-particle variance was used as the OOD score. We measured the OOD detection performance by the area under the precision recall curve (PR-AUC), viewing non-OOD test data as the positive class and OOD data as the negative class. We repeated this procedure five times.

Table 2 compares the performance of WEvidential with four other methods: MCDropout, DEnsemble, and Distributional Distillation (DDistillation) [57], and Posterior Network (PNetwork) [14]. Appendix E provides further details on the experiment. Figure 4 exemplifies how the dispersion of the output particles differ between OOD and non-OOD inputs. WEvidential demonstrates a high classification and OOD detection accuracy simultaneously. Although PNetwork has the best OOD detection performance for the sensorless dataset, the performance of the WGBoost algorithm also exceeds 80%, which is distinct from MCDropout, DEnsemble, and DDistillation.

(a) non-OOD input in class

4

(b) OOD input

1.0

0.8

probability

0.6

0.4

0.2

0.0

class 1

class 2

class3

class 4

class 5

class 6

1.0

0.8

probability

0.6

0.4

0.2

0.0

class 1

class 2

class 3

class 4

class 5

class 6

5 Discussion

This work established the general framework of WGBoost and developed the concrete algorithm WEvidential for evidential learning. The established framework of WGBoost offers exciting avenues for future research. Important directions for future study include (i) exploring alternative loss functionals to the KL divergence, (ii) investigating the convergence properties, and (iii) evaluating robustness of obtained predictive uncertainty in comparison to other methods. A particular limitation of WGBoost may arise when data are not tabular, as is the case of standard gradient boosting. These questions require careful examination and are critical for future study.

References

Ravid Shwartz-Ziv and Amitai Armon. Tabular data: Deep learning is not all you need. Information Fusion , 81:84-90, 2022.

Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U. Rajendra Acharya, Vladimir Makarenkov, and Saeid Nahavandi. A review of uncertainty quantification in deep learning: Techniques, applications and challenges. Information Fusion , 76:243-297, 2021.

Eric Topol. High-performance medicine: the convergence of human and artificial intelligence. Nature Medicine , 25:44-56, 2019.

Sorin Grigorescu, Bogdan Trasnea, Tiberiu Cocias, and Gigel Macesanu. A survey of deep learning techniques for autonomous driving. Journal of Field Robotics , 37(3):362-386, 2020.

Matthew Richardson, Ewa Dominowska, and Robert Ragno. Predicting clicks: estimating the click-through rate for new ads. In Proceedings of the 16th International Conference on World Wide Web , page 521-530, 2007.

Christopher Burges. From ranknet to lambdarank to lambdamart: An overview. Learning , 11, 2010.

Byron P. Roe, Hai-Jun Yang, Ji Zhu, Yong Liu, Ion Stancu, and Gordon McGregor. Boosted decision trees as an alternative to artificial neural networks for particle identification. Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment , 543(2):577-584, 2005.

James Bennett and Stan Lanning. The Netflix prize. In Proceedings of the KDD Cup Workshop 2007 , pages 3-6, 2007.

Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing model uncertainty in deep learning. In Proceedings of The 33rd International Conference on Machine Learning , volume 48 of Proceedings of Machine Learning Research , pages 1050-1059. PMLR, 2016.

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems , volume 30, 2017.

Murat Sensoy, Lance Kaplan, and Melih Kandemir. Evidential deep learning to quantify classification uncertainty. In Advances in Neural Information Processing Systems , volume 31, 2018.

Jakob Gawlikowski, Cedrique Rovile Njieutcheu Tassi, Mohsin Ali, Jongseok Lee, Matthias Humt, Jianxiang Feng, Anna Kruspe, Rudolph Triebel, Peter Jung, Ribana Roscher, Muhammad Shahzad, Wen Yang, Richard Bamler, and Xiao Xiang Zhu. A survey of uncertainty in deep neural networks. Artificial Intelligence Review , 56:1513-1589, 2023.

Alexander Amini, Wilko Schwarting, Ava Soleimany, and Daniela Rus. Deep evidential regression. In Advances in Neural Information Processing Systems , volume 33, pages 1492714937, 2020.

Bertrand Charpentier, Daniel Zügner, and Stephan Günnemann. Posterior network: Uncertainty estimation without OOD samples via density-based pseudo-counts. In Advances in Neural Information Processing Systems , volume 33, pages 1356-1367. Curran Associates, Inc., 2020.

Dennis Thomas Ulmer, Christian Hardmeier, and Jes Frellsen. Prior and posterior networks: A survey on evidential deep learning methods for uncertainty estimation. Transactions on Machine Learning Research , 2023.

Edouard Capellier, Franck Davoine, Veronique Cherfaoui, and You Li. Evidential deep learning for arbitrary lidar object classification in the context of autonomous driving. In 2019 IEEE Intelligent Vehicles Symposium (IV) , pages 1304-1311, 2019.

Patrick Hemmer, Niklas Kühl, and Jakob Schöffer. Deal: Deep evidential active learning for image classification. In 2020 19th IEEE International Conference on Machine Learning and Applications (ICMLA) , pages 865-870, 2020.

Ava P. Soleimany, Alexander Amini, Samuel Goldman, Daniela Rus, Sangeeta N. Bhatia, and Connor W. Coley. Evidential deep learning for guided molecular property prediction and discovery. ACS Central Science , 7(8):1356-1367, 2021.

Jakob Gawlikowski, Sudipan Saha, Anna Kruspe, and Xiao Xiang Zhu. An advanced Dirichlet prior network for out-of-distribution detection in remote sensing. IEEE Transactions on Geoscience and Remote Sensing , 60:1-19, 2022.

Andrew Gelman, John B. Carlin, Hal S. Stern, David B. Dunson, Aki Vehtari, and Donald B. Rubin. Bayesian Data Analysis . Chapman and Hall/CRC, 3rd ed. edition, 2013.

Tianqi Chen and Carlos Guestrin. XGBoost: A scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , KDD '16, page 785-794, 2016.

Guolin Ke, Qi Meng, Thomas Finley, Taifeng Wang, Wei Chen, Weidong Ma, Qiwei Ye, and Tie-Yan Liu. LightGBM: A highly efficient gradient boosting decision tree. In Advances in Neural Information Processing Systems , volume 30, 2017.

C'edric Villani. Topics in Optimal Transportation . Americal Mathematical Society, 2003.

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows In Metric Spaces and in the Space of Probability Measures . Birkhäuser Basel, 2005.

Filippo Santambrogio. {Euclidean, metric, and Wasserstein} gradient flows: an overview. Bulletin of Mathematical Sciences , 7:87-154, 2017.

Qiang Liu. Stein variational gradient descent as gradient flow. In Advances in Neural Information Processing Systems , volume 30, 2017.

Jos'e Antonio Carrillo, Katy Craig, and Francesco S. Patacchini. A blob method for diffusion. Calculus of Variations and Partial Differential Equations , 58(53), 2019.

Yifei Wang, Peng Chen, and Wuchen Li. Projected Wasserstein gradient descent for highdimensional Bayesian inference. SIAM/ASA Journal on Uncertainty Quantification , 10(4):15131532, 2022.

Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of Fokker-Planck equations through gradient-log-density estimation. Entropy (Basel) , 22(8):802, 2020.

Ye He, Krishnakumar Balasubramanian, Bharath K. Sriperumbudur, and Jianfeng Lu. Regularized Stein variational gradient flow. arXiv:2211.07861 , 2022.

Jerome H. Friedman. Greedy function approximation: A gradient boosting machine. The Annals of Statistics , 29(5):1189-1232, 2001.

Tony Duan, Avati Anand, Daisy Yi Ding, Khanh K. Thai, Sanjay Basu, Andrew Ng, and Alejandro Schuler. NGBoost: Natural gradient boosting for probabilistic prediction. In Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 2690-2700. PMLR, 2020.

Peter Bühlmann and Torsten Hothorn. Boosting Algorithms: Regularization, Prediction and Model Fitting. Statistical Science , 22(4):477 - 505, 2007.

Leo Grinsztajn, Edouard Oyallon, and Gael Varoquaux. Why do tree-based models still outperform deep learning on typical tabular data? In Advances in Neural Information Processing Systems , volume 35, pages 507-520, 2022.

Piotr Florek and Adam Zagda´ nski. Benchmarking state-of-the-art gradient boosting algorithms for classification. arXiv:2305.17094 , 2023.

Zhendong Zhang and Cheolkon Jung. GBDT-MO: Gradient-boosted decision trees for multiple outputs. IEEE Transactions on Neural Networks and Learning Systems , 32(7):3156-3167, 2021.

Tianqi Chen, Sameer Singh, Ben Taskar, and Carlos Guestrin. Efficient Second-Order Gradient Boosting for Conditional Random Fields. In Proceedings of the Eighteenth International Conference on Artificial Intelligence and Statistics , volume 38 of Proceedings of Machine Learning Research , pages 147-155. PMLR, 2015.

Jerome H. Friedman. Stochastic gradient boosting. Computational Statistics & Data Analysis , 38(4):367-378, 2002.

Gianluca Detommaso, Tiangang Cui, Youssef Marzouk, Alessio Spantini, and Robert Scheichl. A Stein variational Newton method. In Advances in Neural Information Processing Systems , volume 31, 2018.

Yifei Wang and Wuchen Li. Information Newton's flow: second-order optimization method in probability space. arXiv:2001.04341 , 2020.

Malay Ghosh. Objective priors: An introduction for frequentists. Statistical Science , 26(2):187202, 2011.

J. Aitchison and S. M. Shen. Logistic-normal distributions: Some properties and uses. Biometrika , 67(2):261-272, 1980.

Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. In Advances in Neural Information Processing Systems , volume 29, 2016.

Dilin Wang, Zhe Zeng, and Qiang Liu. Stein variational message passing for continuous graphical models. In Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 5219-5227. PMLR, 2018.

Alexander Lambert, Fabio Ramos, Byron Boots, Dieter Fox, and Adam Fishman. Stein variational model predictive control. In Proceedings of the 2020 Conference on Robot Learning , volume 155 of Proceedings of Machine Learning Research , pages 1278-1297. PMLR, 2021.

Anna Korba, Adil Salim, Michael Arbel, Giulia Luise, and Arthur Gretton. A non-asymptotic analysis for Stein variational gradient descent. In Advances in Neural Information Processing Systems , volume 33, pages 4672-4682, 2020.

Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, and Philippe Rigollet. SVGD as a kernelized Wasserstein gradient flow of the chi-squared divergence. In Advances in Neural Information Processing Systems , volume 33, pages 2098-2109, 2020.

Andre Wibisono. Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem. In Proceedings of the 31st Conference On Learning Theory , volume 75 of Proceedings of Machine Learning Research , pages 2093-3027. PMLR, 2018.

Gareth O. Roberts and Richard L. Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli , 2(4):341 - 363, 1996.

Leo Breiman, Jerome Friedman, R.A. Olshen, and Charles J. Stone. Classification and Regression Trees . Chapman and Hall/CRC, 1984.

Trevor Hastie, Robert Tibshirani, and Jerome Friedman. The Elements of Statistical Learning . Springer New York, 2009.

Sanford Weisberg. Applied Linear Regression . John Wiley & Sons, 1985.

Jose Miguel Hernandez-Lobato and Ryan Adams. Probabilistic backpropagation for scalable learning of Bayesian neural networks. In Proceedings of the 32nd International Conference on Machine Learning , volume 37 of Proceedings of Machine Learning Research , pages 1861-1869. PMLR, 2015.

Dheeru Dua and Casey Graff. UCI machine learning repository, 2017.

Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. In Advances in Neural Information Processing Systems , volume 30, 2017.

Jingkang Yang, Kaiyang Zhou, Yixuan Li, and Ziwei Liu. Generalized out-of-distribution detection: A survey. arXiv:2110.11334 , 2024.

Andrey Malinin, Bruno Mlodozeniec, and Mark Gales. Ensemble distribution distillation. In International Conference on Learning Representations , 2020.

Filippo Santambrogio. Optimal Transport for Applied Mathematicians . Birkhäuser Cham, 2015.

Mingxuan Yi and Song Liu. Bridging the gap between variational inference and Wasserstein gradient flows, 2023.

Michael Arbel, Anna Korba, Adil Salim, and Arthur Gretton. Maximum mean discrepancy gradient flow. In Advances in Neural Information Processing Systems , volume 32, 2019.

Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the FokkerPlanck equation. SIAM Journal on Mathematical Analysis , 29(1):1-17, 1998.

Grigorios A. Pavliotis. Stochastic Processes and Applications . Springer New York, 2014.

Vern I. Paulsen and Mrinal Raghupathi. An Introduction to the Theory of Reproducing Kernel Hilbert Spaces . Cambridge University Press, 2016.

Alex Leviyev, Joshua Chen, Yifei Wang, Omar Ghattas, and Aaron Zimmerman. A stochastic Stein variational Newton method. arXiv:2204.09039 , 2022.

Alex Smola, Arthur Gretton, Le Song, and Bernhard Schölkopf. A Hilbert space embedding for distributions.

Shun ichi Amari. Information Geometry and Its Applications . Springer Tokyo, 2016.

Pavel Izmailov, Wesley J. Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Subspace inference for bayesian deep learning. In Proceedings of 35th Conference on Uncertainty in Artificial Intelligence , 2019.

Appendix

This appendix contains the technical and experiment details referred to in the main text. Appendix A recaps derivation of the Wasserstein gradient. Appendix B discusses the variant of WGBoost built on the unadjusted Langevin algorithm. Appendix C derives the diagonal approximate Wasserstein Newton direction used for WEvidential. Appendix D provides simulation studies on kernel choice of WEvidential and comparison of WGBoost algorithms built on different estimate of the Wasserstein gradient. Appendix E contains the additional details of the experiment presented in the main text.

A Derivation and Example of Wasserstein Gradient

This section recaps derivation of the Wasserstein gradient of a functional F , with examples of common divergences. The Wasserstein gradient depends on a function on Θ called the first variation [24]. The first variation δ F ( µ ) /δµ of the functional F at µ is a function on Θ that satisfies

for all signed measure ν s.t. µ + ϵν ∈ P 2 for all ϵ sufficiently small. The Wasserstein gradient ∇ W F ( µ ) of the functional F at µ is derived as the gradient of the first variation (see [e.g. 24]):

It is common to suppose that the functional F consists of three energies, which are determined by functions U : R → R , V : Θ → R , and W : Θ → R respectively, such that

For a functional F that falls into the above form, the Wasserstein gradient is derived as

where U ′ is the derivative of U : R → R [23]. The KL divergence F ( µ ) = KL( µ | π ) of a distribution π falls into the form with U ( x ) = x log x , V ( θ ) = -log π ( θ ) , and W ( θ ) = 0 , where

Table 3 presents examples of Wasserstein gradients of common divergences F ( µ ) = D( µ | π ) .

In the context of Bayesian inference, the KL divergence is particularly useful among many divergences. The Wasserstein gradient of the KL divergence requires no normalising constant of a posterior distribution π . This is because the Wasserstien gradient depends only on the log-gradient of the posterior ∇ log π ( θ ) = ∇ π ( θ ) /π ( θ ) of the target π , in which case the normalising constant of the target π is cancelled out by fraction. Hence, any posterior known only up to the normalising constant can be used as the target distribution π in the Wasserstein gradient of the KL divergence.

B Langevin Gradient Boosting for KL Divergence

If a chosen functional F on P 2 is the KL divergence F ( µ ) = KL( µ | π ) of a target distribution π , the continuity equation (1) admits an equivalent representation as the Fokker-Planck equation [61]:

where ∆ denotes the Laplacian operator. Recall that the original continuity equation (1) can be reformulated as the deterministic differential equation (2) of a random variable θ t ∼ µ t . In contrast, the Fokker-Planck equation (10) can be reformulated as a stochastic differential equation of a random variable θ t ∼ µ t , known as the overdamped Langevin dynamics [62]:

where B t denotes a standard Brownian motion. Note that the deterministic system (2) in the case of the KL divergence and the above stochastic system (11) are equivalent at population level, in a sense that the law of the random variable θ t in both the systems solves the two equivalent equations.

At the algorithmic level, however, discretisation of each system leads to different particle update schemes. Set the initial distribution µ 0 in (11) to the empirical distribution ˆ µ 0 of N initial particles { θ n 0 } N n =1 . Discretising the stochastic system (11) by the Euler-Maruyama method with a step size ν > 0 yields a stochastic update scheme of particles { θ n m } N n =1 from step m = 0 :

where each ξ n denotes a realisation from a standard normal distribution on R d . The above updating scheme of each n -th particle is known as the unadjusted Langevin algorithm [49]. We can define a variant of WGBoost by replacing the term G i ( µ ) in Algorithm 1 with ∇ log µ i ( · ) + √ 2 /ν ξ i where µ i is an output distribution at each x i and ξ i is a realisation from a standard normal distribution. The procedure is summarised in Algorithm 3, which we call Langevin gradient boosting (LGBoost).

Algorithm 3: Langevin Gradient Boosting

Input: training set { x i , µ i } D i =1 of input x i ∈ X and output distribution µ i ∈ P 2

Parameter: particle number N , iteration M , rate ν , weak learner f , initial constants

(

Output: set of N boosting ensembles ( F 1 , . . . , F N ) at final step M

C Derivation of Approximate Wasserstein Newton Direction

This section derives the diagonal approximate Wasserstein Newton direction based on the kernel smoothing. The approximate Wasserstein Newton direction of the KL divergence was derived in [39] under a different terminology-simply, the Newton direction-from a viewpoint of nonparametric variational inference. We place their result in the context of approximate Wasserstein gradient flows. Appendix C.1 shows the derivation of the smoothed Wasserstein gradient and Hessian. Appendix C.2 defines the Newton direction built upon the smoothed Wasserstein gradient and Hessian, following the derivation in [39]. Appendix C.3 derives the diagonal approximation of the Newton direction.

C.1 Smoothed Wasserstein Gradient and Hessian

Consider the one-dimensional case Θ = R for simplicity. For a map T : R → R and a distribution µ ∈ P 2 , let µ t be the pushforward of µ under the transform θ ↦→ θ + tT ( θ ) defined with a timevariable t ∈ R . This means that µ t is a distribution obtained by change-of-variable applied for µ . The Wasserstein gradient of a functional F ( µ ) can be associated with the time derivative ( d/dt ) F ( µ t ) [23]. In what follows, we focus on the KL divergence F ( µ ) = KL( µ | π ) as a loss functional. Under a condition T ∈ L 2 ( µ ) , the time derivative at t = 0 satisfies the following equality

where G KL ( µ ) denotes the Wasserstein gradient of F ( µ ) = KL( µ | π ) with the target distribution π made implicit. It gives an interpretation of the Wasserstein gradient as the steepest-descent direction because the decay of the KL divergence at t = 0 is maximised when T = -G KL ( µ ) .

The 'smoothed' Wasserstein gradient can be derived by restricting the transform map T to a more regulated Hilbert space than L 2 ( µ ) . A reproducing kernel Hilbert space (RKHS) H associated with a kernel function k : R × R → R is the most common choice of such a Hilbert space [e.g. 26]. An important property of the RKHS H is that any function f ∈ H satisfies the reproducing property f ( θ ) = ⟨ f ( · ) , k ( θ, · ) ⟩ H under the associated kernel k and inner product ⟨· , ·⟩ H [63]. As discussed in [e.g. 46], applying the reproducing property in (12) under the condition T ∈ H and exchanging the integral order, the time derivative satisfies an alternative equality as follows:

where [ G ∗ ( µ )]( · ) := ∫ Θ [ G KL ( µ )]( θ ) k ( θ, · ) dµ ( θ ) corresponds to the smoothed Wasserstein gradient used in the main text. The decay of the KL divergence at t = 0 is maximised by T = -G ∗ ( µ ) .

Similarly, the Wasserstein Hessian of the functional F ( µ ) can be associated with the second time derivative ( d 2 /dt 2 ) F ( µ t ) [23]. As discussed in [e.g. 46], the Wasserstein Hessian of the KL divergence, denoted Hess( µ ) , is an operator over functions T ∈ L 2 ( µ ) that satisfies

See [46] for the explicit form of the Wasserstein Hessian. In the same manner as the smoothed Wasserstein gradient, applying the reproducing property in (14) under the condition T ∈ H and exchanging the integral order, the second time derivative satisfies an alternative equality as follows:

where [Hess ∗ ( µ )]( ⋆ 1 , ⋆ 2 ) := ⟨ k ( ⋆ 1 , · ) , Hess( µ ) k ( ⋆ 2 , · ) ⟩ L 2 ( µ ) is the smoothed Wasserstein Hessian and the symbols ⋆ 1 and ⋆ 2 denote the variables to which each of the two inner products is taken.

In the multidimensional case Θ = R d , the transport map T is a vector-valued function T : R d → R d , where a similar derivation can be repeated by replacing L 2 ( µ ) and H with the product space of d independent copies of L 2 ( µ ) and H . It follows from Proposition 1 and Theorem 1 in [39]-which derives the explicit form of (13) and (15) under their terminology, first and second variations-that the explicit form of the smoothed Wasserstein gradient and Hessian is given by

where ∇ 2 denotes an operator to take the Jacobian of the gradient-i.e., ∇ 2 f ( θ ) is the Hessian matrix of f at θ -and ⊗ denotes the outer product of two vectors. Note that both the smoothed Wasserstein gradient and Hessian are well-defined for any distribution µ including empirical distributions.

C.2 Approximate Wasserstein Newton Direction

In the Euclidean space, the Newton direction of an objective function is a direction s.t. the secondorder Taylor approximation of the function is minimised. Similarly, [39] characterised the Newton direction T ∗ : R d → R d of the KL divergence KL( µ | π ) as a solution of the following equation

Here Θ = R d and H is the product space of d independent copies of the RKHS of a kernel k . To obtain a closed-form solution, [39] supposed that the Newton direction T ∗ can be expressed in a form T ∗ ( · ) = ∑ n i =1 W n k ( · , θ n ) dependent on a set of each particle θ n ∈ Θ and associated vector-valued coefficient W n ∈ R d . Once the set of the particles is given, the set of the associated vector-valued coefficients is determined by solving the following simultaneous linear equation

These equations (16) can be rewritten in a matrix form [64]. Let K := N × d . Define a block matrix H ∈ R K × K and a block vector G ∈ R K by the following partitioning

with each block specified as H ij := [Hess ∗ ( µ )]( θ i , θ j ) ∈ R d × d and G i := [ G ∗ ( µ )]( θ i ) ∈ R d . Define a block matrix K ∈ R K × K and a block vector W ∈ R K by the following partitioning

with each block of K specified as K ij := I d × k ( θ i , θ j ) ∈ R d × d , where I d denotes the d × d identity matrix. Notice that W is a block vector that aligns the vector-valued coefficients { W n } N n =1 . Using these notations, the optimal coefficients that solve (16) is simply written as W = -H -1 G [64].

Given the optimal coefficients W = -H -1 G , the Newton direction T ∗ ( θ n ) evaluated at the given particle θ n for each n = 1 , . . . , N can be written in the following block vector form

To distinguish from the standard Newton direction in the Euclidean space, we call (17) the approximate Wasserstein Newton direction. The approximate Wasserstein Newton direction yields a second-order particle update scheme. Suppose we have particles { θ n m } N n =1 to be updated at each step m . At each step m , define the above matrices H and G with the empirical distribution µ = ˆ π m of the particles { θ n m } N n =1 . Replacing the Wasserstein gradient in the particle update scheme (3) by the approximate Wasserstein Newton direction (17) provides the second-order update scheme in [39].

C.3 Diagonal Approximate Wasserstein Newton Direction

We derive the diagonal approximation of the approximate Wasserstein Newton direction, which we used for our second-order WGBoost algorithm. A few approximations of the approximate Wasserstein Newton direction were discussed in [39] for better performance of their particle algorithm. We derive the diagonal approximation so that no matrix product and inversion will be involved. Specifically, we replace the matrices K and H in (17) by the diagonal approximations ˆ K and ˆ H , that is,

where K nn = I d × k ( θ n , θ n ) = I d for the Gaussian kernel k used in this work, and the matrix h nn ∈ R d × d denotes the diagonal approximation of the diagonal block H nn of H .

Recall that H nn = [Hess ∗ ( µ )]( θ n , θ n ) . Denote by Diag( A ) the diagonal of a square matrix A . The diagonal approximation h nn is a diagonal matrix whose diagonal is Diag( H nn ) . We plug the diagonal approximations ˆ K and ˆ H in (17). It follows from inverse and multiplication properties of diagonal matrices that the approximate Wasserstein Newton direction turns into a form

At an arbitrary particle location θ , denote by [ H ∗ ( µ )]( θ ) the diagonal of the smoothed Wasserstein Hessian [Hess ∗ ( µ )]( θ, θ ) . It is straightforward to see that the diagonal can be written as

Notice that Diag ( H nn ) = [ H ∗ ( µ )] ( θ n ) by definition. It therefore follows from the formula (18) with G n = [ G ∗ ( µ )]( θ n ) and Diag ( H nn ) = [ H ∗ ( µ )] ( θ n ) that the diagonal approximate Wasserstein Newton direction at an arbitrary particle location θ can be independently computed by

We used this direction in Section 3. In the main text, the diagonal approximate Wasserstein Newton direction is defined for each loss functional F i ( · ) = D( · | µ i ) , with π = µ i , using the smoothed Wasserstein gradient G ∗ i ( µ ) and the diagonal of the smoothed Wasserstein Hessian H ∗ i ( µ ) for each i-th output distribution µ i .

D Simulation Study for WEvidential

This section provides simulation studies on kernel choice of WEvidential and comparison of different estimate of the Wasserstein gradient to use in the WGBoost framework.

D.1 Choice of Kernel

We used the kernel smoothing estimate of the Wasserstein gradient of the KL divergence in order to built WEvidential. The smoothed Wasserstein gradient originates in an approximate Wasserstein gradient flow called Stein variational gradient descent (SVGD) [43]. It is fairly common in SVGD in practice to use the Gaussian kernel k ( θ, θ ∗ ) = exp( -∥ θ -θ ∗ ∥ 2 /h ) with the scale h > 0 . For WEvidential, the scale h may be viewed as a hyperparameter to choose. We recommend using the value of the scale s.t. 0 . 01 ≤ h ≤ 1 . 0 in general. This work uses the value h = 0 . 1 throughout the experiments in the main text. One may opt for performing more advanced tuning of the scale h .

We provide a simulation study to examine sensitivity to the scale value. We prepared a synthetic dataset { x i , µ i } D i =1 whose inputs are 200 gird points on the interval [ -3 . 5 , 3 . 5] and output distributions are normal distributions µ i ( θ ) = N ( θ | sin ( x i ) , 0 . 5) conditional on each x i . We used WEvidential with different kernel scales h = 0 . 001 , 0 . 01 , 0 . 1 , 1 . 0 , 10 , 100 and fitted it to the synthetic dataset { x i , µ i } D i =1 . The decision tree regressor with the maximum depth 3 was used for each weak learner. The learning rate and the number of weak learners were set to 0 . 1 and 100 . The initial constant { ϑ n } 10 n =1 of WEvidential was set to 10 grid points in the interval [ -10 , 10] .

We computed the output of WEvidential for 500 grid points in the interval [ -3 . 5 , 3 . 5] . We used the maximum mean discrepancy (MMD) [65] to measure the approximation error between the empirical distribution ˆ µ i of the output of WEvidential and the output distribution µ i :

where k D is a Gaussian kernel k D ( θ, θ ′ ) = exp( -( θ -θ ′ ) 2 /s ) with the scale s = 0 . 025 . The total approximation error was measured by the MMD averaged over all the inputs. Figure 5 shows the total approximation error of WEvidential for each scale value h in the common log scale, together with examples of the output of WEvidential for different values of h . It demonstrates that the total error is minimised by h = 0 . 1 and stays in a relatively small value range for 0 . 01 ≤ h ≤ 1 . 0 .

(a) The total error

(b)

h

= 0

.

10

(c)

h

= 0

.

01

(d)

h

= 100

-0.8

1.0

common-logMMD

-1.2

1.4

-1.6

-1.8

2.0

-3

2

1

0

1

2

common-log kernel bandwidth

1

-2

1

0

X

-1

-3

2

-1

0

1

2

3

X

1

0

1

2

3

X

D.2 Comparison of Different Wasserstein Gradient Estimates

We compare four WGBoost algorithms built on different estimates of the Wasserstein gradient of the KL divergence. The first three algorithms set the estimate G i ( µ ) in Algorithm 1 to, respectively,

the smoothed Wasserstein gradient in (7);

the diagonal approximate Wasserstein Newton direction in (8);

the full approximate Wasserstein Newton direction in (17).

The fourth algorithm is LGBoost in Appendix B which is rather a variant of WGBoost. The first and third algorithms are called the first-order WEvidential and the full-Newton WEvidential, respectively. The second algorithm is our default WEvidential presented in Section 3. The first-order WEvidential is implemented by removing h n i and replacing g n i ⊘ h n i with g n i in Algorithm 2. The full-Newton WEvidential is implemented by replacing g n i ⊘ h n i in Algorithm 2 with v n i computed by the following Algorithm 4, where ∇ 2 f ( θ ) denotes the Hessian matrix of a function f : Θ → R at θ .

We prepared the same synthetic dataset { x i , µ i } D i =1 as Appendix D.1 and fitted each algorithm to the dataset. We computed the output of each algorithm for 500 grid points in the interval [ -3 . 5 , 3 . 5] and measured the approximation error by the MMD in the same manner as Appendix D.1. The decision tree regressor with the maximum depth 3 was used for each weak learner, and the learning rate was set to 0 . 1 . We used an increasing number of weak leaners up to 100 weak learners, in order to observe the decay of the approximation error and the increase of the computational time. The initial constant { ϑ n } 10 n =1 for each algorithm was set to 10 grid points in the interval [ -10 , 10] , which sufficiently differs from the output distributions so that the decay of the approximation error is clear.

Figure 6 shows the approximation error and the computational time of each algorithm with respect to the increasing number of weak learners, together with the output of each algorithm with 100 weak learners trained. It demonstrates that WEvidential and full-Newton WEvidential reduce the approximation error most efficiently, while full-Newton WEvidential takes the longest computational time among others. As in Algorithm 4, the computation of the full approximate Wasserstein Newton direction requires the inverse and product of the ( N × d ) × ( N × d ) block matrices, where N denotes the particle number N and d denotes the particle dimension. The error decay of LGBoost is not only slow but also shows stochasticity due to the Gaussian noise used in the algorithm. We therefore recommend our default WEvidential for better performance and efficient computation.

Input:

input

x

i

, output distribution

µ

i

, outputs of

N

boostings

(

F

1

m

(

x

i

)

, . . . , F

N

m

(

x

i

))

Output:

(

v

1

i

, . . . , v

N

i

)

evaluated at

(

F

1

m

(

x

i

)

, . . . , F

N

m

(

x

i

))

for

n

1

, . . . , N

do

g

n

i

1

N

∑

N

j

=1

∇

log

µ

i

(

θ

j

i

)

k

(

θ

j

i

, θ

n

i

) +

∇

k

(

θ

j

i

, θ

n

i

)

for

k

1

, . . . , N

do

H

nk

i

1

N

∑

N

j

=1

-∇

2

log

µ

i

(

θ

j

i

)

k

(

θ

j

i

, θ

n

i

)

k

(

θ

j

i

, θ

k

i

) +

∇

k

(

θ

j

i

, θ

n

i

)

⊗∇

k

(

θ

j

i

, θ

k

i

)

K

nk

i

I

d

·

k

(

θ

n

i

, θ

k

i

)

end

end

(a) Approximation Error

(b) Computational Time

(c) First-order

(d) Diagonal-Newton

(e) Full-Newton

(f) LGBoost

discrepancy

0.20-

First-orderWEvidential

time

0.0

DefaultWEvidential

Full-NewtonWEvidential

omputation

0.15

LGBoost

-0.5

maximummean

-1.0

0.10

common-log

-1.5

0.05

-2.0

0.00

20

40

60

80

100

1.0

1.2

1.4

1.6

1.8

2.0

weaklearnernumber

common-logweaklearnernumber

1

-1

-1

E Additional Details of Main Experiments

This section describes additional details of the applications in Section 4. All the experiments were performed with x86-64 CPUs, where some of them were parallelised up to 10 CPUs and the rest uses 1 CPU. The scripts to reproduce all the experiments are provided in the source code. Appendices E.1 to E.3 describe additional details of the applications in Sections 4.1 to 4.3 respectively. Appendix E.4 describes a choice of initial constants { ϑ n 0 } N n =1 for the WGBoost algorithm used in Section 4.

E.1 Detail of Section 4.1

The normal response distribution N ( y | m,σ ) in Example 1 was used in Section 4.1. The normal response distribution has the scale parameter σ that lies only in the positive domain of the Euclidean space R . We reparametrised it as one in the Euclidean space R by the log transform σ ′ := log( σ ) , which is the standard practice in Bayesian computation [20]. The inverse of the log transform

is the exponential transform σ = exp( σ ′ ) . The change of variable formula tells the form of the individual-level posterior on ( m,σ ′ ) up to the normalising constant. Under the prior in Example 1, the individual-level posterior on ( m,σ ′ ) conditional on each individual response y i is given by

where the Jacobian determinant | dσ/dσ ′ | = exp( σ ′ ) is used for the change of variable. Figure 7 shows the output of WEvidential for the old faithful geyser dataset and the conditional density estimated through (6).

6

5

4

3

2

1

50

60

70

80

90

x

50

60

70

80

90

E.2 Detail of Section 4.2

Section 4.2 used the same reparametrisation of the normal response parameter ( m,σ ) as that of Appendix E.2. The NLL and RMSE of each algorithm on test data { x i , y i } D i =1 were computed by

where p ( y i | x i ) and ˆ y i denote the predictive distribution and the point prediction provided by the algorithm. For WEvidential, we standardised the responses in the training data, in which case standardisation of test data were performed as y ′ i = ( y i -y train mean ) /y train std with the mean y train mean and standard deviation y train std of the training responses. Hence WEvidential provided the predictive distribution p ( y ′ i | x i ) and the point prediction ˆ y ′ i for the standardised responses y ′ i . With no loss of generality, the NLL and RMSE for the original responses y i can be computed as follows:

where the equality of the NLL follows from the change of variable p ( y i | x i ) = p ( y ′ i | x i ) /y train std and the equality of the RMSE follows from rearranging the terms.

Weprovide a brief description of each algorithm used in Section 4.2. For all the algorithms, the normal response distribution p ( y | m,σ ) was specified as the response distribution, where the algorithm produces a point or distributional estimate of the response parameter ( m,σ ) at each input x .

MCDropout [9] trains a single neural network F while dropping out each network weight with some Bernoulli probability. It can be interpreted as a variational approximation of a Bayesian neural network. It generates multiple subnetworks { F n } N n =1 by subsampling the network weights by the dropout. The predictive distribution p ( y | x ) is given by the model averaging (1 /N ) ∑ N i =1 p ( y | ( m,σ ) = F n ( x )) for each input x .

DEnsemble [10] simply trains independent copies { F n } N n =1 of a neural network F in parallel. It is one of the mainstream approaches to uncertainty quantification based on deep learning. The predictive distribution is given by the model averaging as in MCDropout.

6

5

4

3

2

1

CDropout [55] consider a continuous relaxation of the Bernoulli random variable used in MCDropout to optimise the Bernoulli probability of the dropout. It generates multiple subnetworks { F n } N n =1 by subsampling the network weights by the dropout with the optimised probability. The predictive distribution is the same as MCDropout.

NGBoosting [32] is a family of gradient booting that use the natural gradient [66] of the response distribution as a target variable of each weak learner. In contrast to other methods, NGBoost outputs a single value F ( x ) to be plugged into the response-distribution parameter. The predictive distribution p ( y | x ) is given by p ( y | ( m,σ ) = F ( x )) for each input x .

DEvidential [13] extends deep evidential learning [11], originally proposed in classification settings, to regression settings. It considers the case where the individual-level posterior of the response distribution falls into a conjugate parametric form, and predicts the hyperparameter of the individual-level posterior by a neural network. The predictive distribution is also given in a conjugate closed-form.

The algorithms used in Section 4.2 are computationally efficient uncertainty quantification methods that are commonly-used in practice. In the following, we provide a limited yet additional comparison of WEvidential with other Bayesian neural networks (BNNs) and a kernel-based model, Gaussian process (GP). We used a subset of the datasets in Section 4, which were used in [67] who employed BNNs with one hidden layer of 50 units. We additionally used a large-scale dataset, keggd . For the keggd dataset only, we reported the NLL and RMSE of the normalised outputs without the adjustment (19) and (20), in line with [67]. WEvidential was compared with BNNs learned by variational inference with the reparametrisation trick (VI), deterministic variational inference (DVI), stochastic weight averaging Gaussian (SWAG), principal component analysis subspace inference (PCA+VI), and two-layer deep Gaussian process with 50 induced points trained via expectation propagation (DGP1-50). Table 1 summarises the result.

E.3 Detail of Section 4.3

We reparametrised the parameter of the categorical response distribution C ( y | q ) used in Section 4.3, similarly to the normal response distribution used in Sections 4.1 and 4.2. The categorical response distribution has a class probability parameter q = [ q 1 , . . . , q k ] in the simplex ∆ k . We applied the log-ratio transform q ′ := [log( q 1 /q k ) , . . . , log( q k -1 /q k )] ∈ R k -1 that maps from the simplex ∆ k to the Euclidean space R k -1 [42]. The inverse of the log-ratio transform is the logistic transform

The logistic normal distribution on the original parameter q corresponds to a normal distribution on the transformed parameter q ′ [42]. The change of variable formula tells the individual-level posterior on q ′ up to the normalising constant. Under the prior in Example 2, the individual-level posterior

p ( q ′ | y i ) conditional on each individual observed response y i is given by

where [ y i = j ] is 1 if y i is the j -th class label and 0 otherwise.

We provide a brief description of each algorithm used in Section 4.3. MCDropout and DEnsemble are described in Appendix E.2.

DDistillation [57] predicts the parameter of a Dirichlet distribution over the simplex ∆ k by a neural network using the output of DEnsemble. The output of multiple networks in DEnsemble is distilled into the Dirichlet distribution controlled by one single network.

PNetwork [14] considers the case where the individual-level posterior of the categorical response distribution falls into a Dirichlet distribution similarly to deep evidential learning [11]. It predicts the hyperparameter of the individual-level posterior given in the form of the Dirichlet distribution.

For the OOD detection, the OOD score we used was the inverse of the maximum norm of the variance of the WEvidential output. There are other quantities that are possible to use as the OOD score. For example, the entropy of the predictive distribution p ( y | x ) is a quantity computable for any probabilistic classification method. We compared the OOD detection performance of WEvidential with that of NGBoost and Random Forest (RForest) based on the entropy of their predictive distributions. For reference, we also computed the OOD detection performance of WEvidential based on the entropy of the predictive distribution. For NGBoost, the decision tree regressor was used for each weak learner and 4000 weak learners were trained with the learning rate 0.01. For RForest, the maximum depth 3 was used. Table 5 shows the result on the segment dataset, demonstrating that the performance of WEvidential is higher, to a large margin, than that of NGBoost and RForest based on the entropy.

To investigate on the effective learning rate of WEvidential for classification, we performed a simple simulation study of WEvidential using the synthetic madelon dataset. We created 600 data points of three-dimensional inputs and three-class labels. The data subset that belongs to the last class was kept as the OOD samples. We randomly held out 20% of non-OOD samples as a test set to measure the test classification accuracy. We trained WEvidential using different learning rates from 0.1 to 1.0 while fixing the number of weak learners to 4000. Table 6 shows that the result for the learning rate 0.4 demonstrates the best classification accuracy and the third best OOD detection performance. We also trained WEvidential using different numbers of weak learners from 1000 to 4000 while fixing the learning rate to 0.4. Table 7 shows that the result for the number of weak learners 4000 demonstrates the best classification accuracy and the second best OOD detection performance. While the most effective learning rate differs depending on each dataset, we drew on the insight obtained from the simulation study and set the learning rate to 0.4 for classification.

E.4 Choice of Initial State of WGBoost

In standard gradient boosting, the initial state at step m = 0 is specified by a constant that most fits given data. Similarly, we specified the initial state { ϑ n 0 } N n =1 of WEvidential by a set of constants that most fits the output distributions in average. We find such a set of constants by performing an approximate Wasserstein gradient flow averaged over all the output distributions. Specifically, given the term G i ( µ ) in Algorithm 1, we define ¯ G ( µ ) := (1 /D ) ∑ D i =1 G i ( µ ) and perform the update scheme of a set of N particles { ¯ ϑ n m } N n =1 :

with the learning rate ν 0 = 0 . 01 up to the maximum step number m = 5000 . The initial particle locations for this update scheme were sampled from a standard normal distribution. We specified the initial state { ϑ n 0 } N n =1 by the set of particles { ¯ ϑ n m } N n =1 obtained though this scheme at m = 5000 .

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claim made in the abstract and introduction were elaborated in the subsequent section and empirically demonstrated through the experiments.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: A concise discussion on the important remaining questions and the potential limitation were provided in the last section.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA] .

Justification: This work primarily focused on the methodological development of the proposed algorithm, empirically demonstrating the performance on the various benchmarks.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The computational procedure of the proposed algorithm was clarified in the main text and appendix, with the setting to reproduce the experiments detailed.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The source code is available online with the instruction, and also submitted in the supplementary file. For the blind review process, the source code URL was masked in the manuscript. The source code URL will be unmasked after the review process.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All the experimental setting was clarified in the main text and appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All the experiments to measure the predictive performance of the proposed algorithm were repeated several times to estimate the standard deviation of the score.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The computer resource was clarified in the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The conducted research conforms the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This work proposed a generic algorithm that is not tied to specific applications.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This work proposed a generic algorithm.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The credit and license were appropriately included in the text and source code.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes] .

Justification: The documentation and license were included in the source code.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This work does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This work does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
This work proposes to perform boosting with base learner which are fitted to the Wasserstein gradient of a loss function on the space of probability distributions, which can be useful in particular to capture uncertainty of the models.

Several variants of the algorithm are discussed (by adding a diagonal Hessian preconditioner and with different approximation of Wasserstein gradients for functionals not differentiable on discrete measures). The method is demonstrated on posterior regression tasks where the functionals are KL divergences with respect to some prior, and applied on different real datasets benchmarks.

STRENGTHS:
This paper is well written and proposes a new interesting boosting method to minimize loss over probability distributions.

- The paper is well written
- A new boosting algorithm guided by functionals on probability distributions
- Application on real datasets outperforming baseline methods

WEAKNESSES:
The paper is good overall in my opinion, but still has some weaknesses.

- Experiments focus on KL divergence functional.
- No theoretical analysis

QUESTIONS:
Did you compare between the different algorithms (i.e. with and without Wasserstein Hessian preconditioner, and Kernel vs Langevin approximation) on the benchmark of Section 4.2?

Are there other possible functionals which could be interesting besides the KL or other divregences between probability distributions?

Could the method be used with input distributions, e.g. to do regression on probability distributions?

Typos:
- Line 145: "gradint"
- Line 520: "Wasserstien"

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed a new ensemble algorithm, called Wasserstein Gradient Boosting (WGBoost), which is a novel gradient boosting framework that leverages the wasserstein gradient for probabilistic prediction. Specifically, WGBoost fits a new base learner to the Wasserstein gradient of a loss functional on the space of probability distributions. Since we cannot access to the probability distribution at each iteration and we also cannot evaluate the wasserstein gradient at each iteration, the author proposed to implement it using the particle method and approximate the functional gradient by the kernel method. The algorithm returns a set of particles that approximate a target distribution for each input. The main application demonstrated is posterior regression, where WGBoost provides a distributional estimate of output-distribution parameters.\n\nSTRENGTHS:\n- WGBoost's ability to approximate target distributions with particles offers a robust approach to posterior regression,  capturing predictive uncertainty effectively. \n- The proposed method shows superior performance in empirical evaluations on real-world tabular datasets, both for regression and out-of-distribution detection tasks.\n- The implementation is seems easy by utilizing the particle method combined with kernel approximation.\n\nWEAKNESSES:\nThe approach seems interesting, but the paper lacks the comparison with existing work. Moreover, there is no analysis or discussion when this algorithm is useful as I will explain in the below. \n- It seems that the proposed method seems almost identical to the stein variational gradient descent (SVGD), but there is no qualitative and quantitative comparison with that. Please explain what is the fundamental difference compared to SVGD. \n- Since the proposed method is very similar to SVGD, I think the comparison with SVGD and its extended methods including [3, 5, 6, 7] (I think there are other many variants of extention in SVGD).\n\n- It has been known that the posterior approximation quality strongly depends on the choice of kernels in SVGD [1, 2, 3, 4]. Since the proposed algorithm is almost identical to SVGD, the quality of approximation by the WGBoost is strongly affected by the choice of the kernel function. However, there is no discussion about this point.\n\n- The numerical experiments are only conducted with respect to the final performance on benchmark dataset and I cannot understand when and what kind of problems the proposed algorithm is suitable to  approximate the posterior distribution. It is known that SVGD suffers from collapse phenomena.\n- In addition to the above point, there is no discussion about the computational cost of the proposed algorithm. When I say the computational cost, I point about the computational cost at each iteration and convergence speed. How large computational cost is compared to existing method regarding the number of particles and training dataset size ? I think the proposed algorithm is based on the boosting method, so it suffers from large computational cost with respect to the training dataset size.\n- As for the convergence speed, it has been known that the convergence of the SVGD is slow, which does not show the linear convergence [4], and I suspect that the proposed method suffers similar problem. However, no discussion is present about the convergence speed or no numerical comparison exists with existing method.\n\n[1] Stein Points \n\n[2] Measuring Sample Quality with Kernels\n\n[3] Kernel Stein Discrepancy Descent\n\n[4] On the geometry of Stein variational gradient descent\n\n[5] FUNCTION SPACE PARTICLE OPTIMIZATION FOR BAYESIAN NEURAL NETWORKS\n\n[6] Feature Space Particle Inference for Neural Network Ensembles\n\n[7] Repulsive Deep Ensembles are Bayesian\n\nQUESTIONS:\nI wrote questions in the Weakness.\n\nLIMITATIONS:\nThe limitation is unclear. As far as I read no formal description is presented.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a probabilistic boosting tree algorithm called Wasserstein boosting that uses a smoothed particle gradient to provide probabilistic predictions. Experiments are performed using UCI tabular regression, and out of distribution classification.\n\nSTRENGTHS:\nOriginality: \n\n-\tI like the application of Wasserstein gradients and gradient flow to gradient boosting trees, don’t think I’ve really seen any paper like this before.\n\nQuality:\n\n-\tI didn’t check especially carefully but the machinery for the algorithm seems to be well explained and correct.\n\nClarity:\n\n-\tThe relevant machinery of Wasserstein gradients and gradient boosting is mostly well explained.\n\nSignificance:\n\n-\tMaking trees more probabilistic with limited modifications to their tabular capabilities would be a quite nice advance.\n\nWEAKNESSES:\nOriginality: \n\n-\tNothing really noted. Being a straightforward application of some machinery is totally fine.\n\nQuality:\n\n-\tFor a paper on trees, I find the experiments quite small scale and limited. I would have expected NLL experiments on larger scale datasets, such as the xgboost and lightgbm papers. \n\n-\tComparisons to (approximate) Bayesian neural networks are quite weak – many other methods such as sgmcmc, etc. tend to outperform on datasets of these scales. references: https://arxiv.org/abs/1902.03932, https://arxiv.org/abs/1907.07504, https://arxiv.org/abs/2002.03704, amongst others\n\n-\tThe natural missing comparison here is to Gaussian processes and other kernel methods, which are naturally probabilistic and similarly nonparametric (like trees).\n\n-\tAnother missing set of experiments is comparison to quantile regression  / pinball loss using trees, which is implemented directly in lightgbm. Quantile regression itself is also naturally nonparametric in at least some sense.\n\nClarity:\n\n-\tMy understanding of Wasserstein particle flows is that the particles should interact in some manner during the gradient step. However, the writing of Algorithm 1 makes this quite unclear. I think that the kernel smoothing in the gradient step for the approximate flow is what makes the particles interact, but it’s overall quite unclear.\n\n         o\tThe code doesn’t seem to provide any clarity here.\n\n-\tOverall, it’s quite unclear which parameters in the loss are actually being estimated. If we’re only estimating uncertainty in regression parameters (or analogously classification), then there’s straightforward two stage approaches. \n\n      o\tFor example, one can easily fit a tree predicting the mean and then modify the loss function to predict its variance, or we can modify the loss in classification problems to do something analogous.\n\n\n\n-\tThe writing is extremely passive and non-specific. Suggestions below.\n\n       o\tL150: “procedure of exact or approximate” Please use the algorithm box to specifically write or point out which algorithm is used in the experiments. The current algorithm is so unspecific it’s very hard to follow.\n       o\tL113-115: “Although [32] … originally suggested…” rephrase to something like “Although Friedman [32]  originally proposed using a line search … , Buhlmann and Hothorn [34] recommend against the line search …”\n\nSignificance:\n\n-\tPart of the strength of trees in my experience is that they scale pretty well to large tabular datasets (e.g. n = 10 million). You also tend to need strong uncertainty quantification on these types of datasets, which is part of the reason why Bayesian neural nets became popular for a while. Yet, these large scale uncertainty quantification experiments are lacking from the paper.\n\n      o\tBayesian neural nets: https://proceedings.mlr.press/v115/izmailov20a.html, https://proceedings.mlr.press/v130/immer21a/immer21a.pdf, \n\n      o\tGaussian processes: https://arxiv.org/abs/1809.11165,\n\nQUESTIONS:\nAlgorithm 1: \n\n-\tIs the Wasserstein gradient where the particles manage to interact? Otherwise, I see nowhere else the various particles / learners would interact.\n\nTable 1: where do the error bars in the baseline methods come from?\n\n-\tWas the same preprocessing done / e.g. same train / test split for these?\n\n-\tIn general, I’m pretty sure that MC Dropout is a quite weak baseline for a Bayesian neural net, and many other BNN approaches are much stronger than it at this point (see references above)\n\n-\tthe clear tree baselines here would be quantile regression (implemented in lightgbm) and fitting a mean model with MSE loss and then a second tree to predict the variance by optimizing $\\max s = N(y | \\hat y, \\exp\\{s\\})$ which can be done in lightgbm with a manual loss function.\n\nSection 3.3: \n-\tso we’re just doing Wasserstein gradients on the mean and variance in a Gaussian regression problem? If so, then it seems like a mean / variance two stage fit model would be the natural “base” method to compare to (and could still plug in the “prior” as regularization”)\n\nSection 3.1:\n\nL190: “it depends only on the log gradient …”: yes, this is the requirement for most Bayesian (and probabilistic) inference, e.g. MCMC sampling, variational inference techniques.\n\nTable 1: what are the error bars?\n\nTable 2: I believe that MC Dropout is pretty well known for being poor at out of distribution detection [find reference], while the P Network is certainly a stronger baseline. What would entropies of a traditional multiclass classification tree perform like in terms of an out of distribution detector here?\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

comments: SUMMARY:
The paper introduces a novel gradient boosting framework called Wasserstein Gradient Boosting (WGBoost). Unlike traditional gradient boosting methods that fit base learners to the gradient of the loss function, WGBoost fits them to the Wasserstein gradient of a loss functional defined over probability distributions. This approach is particularly useful for probabilistic prediction, where the goal is to approximate the uncertainty in the model prediction. The authors provide a general formulation of WGBoost, its algorithmic implementation, and empirical evaluations on various benchmarks. 

Paper's major contributions include, the introduction of the Wasserstein gradient flow framework into gradient boosting; the development of an approximate algorithm for posterior regression using the KL divergence; and the demonstration of WGBoost's performance on regression, classification, and out-of-distribution (OOD) detection tasks. Authors also propose a second-order WGBoost algorithm built on the approximate Wasserstein gradient and Hessian of the KL divergence.

STRENGTHS:
S1: The application of Wasserstein gradient flows to gradient boosting is a novel and promising direction, providing a new perspective on ensemble learning algorithms.

S2: Solid theoretical concepts, with detailed derivations and explanations of the Wasserstein gradient and Hessian approximations.

WEAKNESSES:
Authors did not specify details about the hyperparameter selection for Conditional Density Estimation, Classification and OOD Detection tasks.

Misc:
There is a typo in the line no 284, root is written as room.

QUESTIONS:
I would be curious to know about the scalability and the computational complexity of the proposed algorithm.

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-15