Wasserstein Gradient Boosting: A Framework for Distribution-Valued Supervised Learning
Authors: Takuo Matsubara
Venue: NeurIPS 2024 poster (2024)
Keywords: Gradient Boosting; Wasserstein Gradient Flow; Uncertainty Quantification
Primary area: probabilistic_methods
TL;DR: We propose a novel family of gradient boosting, Wasserstein gradient boosting, which returns a set of particles that approximates a target probability distribution assigned at each input.

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- Valuue Supervised Learning
Takuo Matsubara
The University of Edinburgh Edinburgh, EH9 3JZ takuo mwtsubaraaedd ac. uk
Abstract
Gradient boosting is a sequential ensemble method that fits a new weaker learner pseudo residuals each iteration. We propose Wasserstein gradient boosting, novel extension of gradient boosting, which fits a new weak learner to alternative pseudo residuals that are Wasserstein gradients of loss functionals probability supervised learn-
distributions assigned at each input. It solves distribution-van ing, where the output values of the training dataset are probability distributions. In classification and regression, model typically returns, for each input, point estimate of parameter of noise distribution specified for a response variable, such as the class probability parameter of a categorical distribution specified for a response label. main application of Wasserstein gradient boosting in this paper is tree-based evidential learning, which returns 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.
Introduction
Gradient boosting is a celebrated machine learning algorithm that has achieved considerable success with tabular data [ Gradient boosting has been extensively used for point forecasts and probabilistic classification, yet relatively small number studies have heen concerned with the predictive uncertainty of gradient boosting. Predictive uncertainty of machine learning models plays growing role in today's 's real-world production systems 2 It is vital for safety-critical systems, such as medical diagnoses [3] and autonomous driving 4 assess the potential risk of their actions that partially or entirely rely on predictions from their models. Gradient boosting has already been applied in 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 model return point estimate f(x: of the response parameter for each input In recent years, the importance of capturing uncertainty in the model output ((r) has increasingly been emphasised (2] A variety of approaches have been proposed to obtain distributional estimate p(0) 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
the space of network outputs. Marginalising the predictive distribution p(y 0) over the distributional estimate p(0 x)) has been demonstrated confer enhanced predictive accuracy and robustness against adversarial attacks [1]]. Furthermore, the dispersion of the distributional estimate has been used as powerful indicator for -ut-of-dist (OOD) detection [2]
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

3
22 -1
(a) 0 weak leamner trained
(b) 15 weak learners trained
(c) 100 weak learners trained
Figure 1: Illustration of WGBoost trained on set { /1 whose inputs are 10 grid points in
[-3.5, 3.5] and each output distribution is a normal distribution (0: NN sin(z:). 0.5) over 0 E IR. The blue area indicates the 95 % high probability region of the conditional distribution N(O sin(x:), 0.5). WGBoost returns N
particles (red lines) to predict the output distribution
for each input æ. This illustration uses the Gaussian kernel regressor for every weaker learner.
In this context, a line of research based on the concept of evidential learning has recently gained significant attention [T
14 15]. The idea can be broadly interpreted as making use of the 13
iidiiiduullleeell posterior p(O yi) of the response parameter 6 conditional on each individual datum yes which arises from the response-distriee distrrbbutibiinnnn likelihood p(y: 0) and a user-specified prior p(0). each individual-level posterior falls into closed form characterised by some hyperaaraeeter, neural networks can be trained by using the finiteddimensional hyperparameter as a target value for each input. Outstanding performance and computational efficiency of the cxisting approaches have been delivered in a wide spectrum of engineering and medical applications [16 17 18 [9]. 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 finiee-dimensional hyperparameter can be predicted by proxy. In general, posterior distributions are known only up 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 infinite- dimensional output for each input. This challenge poses the following fundamental question:
Consider a supervised learning setting whose outputs are probability distributions. Given training set of input values and output distributions (0,, ) can we build model that receives an input and returns 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 TiaureDillssee Dillstraate inputs and outputs of WGBoost, In this paper, we focus on application of WGBoost to evidential learning, where the individual-level posterior p(0 yı) of the response parameter is used as the output distribution μ4 for each input 2; in the training set. Figure ecoomarees the pipeline of evidential learning based on WGBoost with that of Bayesian learning.
Contributions Our contributions are summarised as follows:
1. Methodology of WGBoost:
Section 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 standard gradient boosting that fits a weak learner to the gradient of loss function, WGBoost fits weak learner to the estimated Wasserstein gradient of loss functional over probability distributions.
2. Application to Evidential Learning: Section establishes tree- tree-based evidential learning based on WGBoost, with the loss functional specified by the Kullback -ulbbeck-Leibler (KL) diver- gence. 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 approximate Wasserstein gradient and Hessian of the KI divergence.
3. Experiment on Real-world Data: Section Jection[ddemone demonstrates the performance of probabilistic regression, and classification with OOD detection, 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.

---

= far w)
ply 8)
(0+ (2)..... (a))
ply 8)
E
f(z,, w)
y
WGBooss
θ
propagate
training
W
p twr
(aa.p(0
(a) Bayesian learning of a model (a, w)
(b) Evidential learning based on WGBoost
Figure 2: Comparison of the pipeline of (a) Bayesian learning and (b) evidential learning based on WGBoost. The former uses the (global-level) posterior p(w { 3/ 31-==)) 1) of the model parameter conditional on all data, and samples multiple models from it. The latter uses the individual-leve posterior p(0) yi) of the response parameter as the output distribution of the training set, and trains WGBoost that returns a particle-baser distributional estimate p(0) x)) of 0 for each input x.
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 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-van valuued supervised learning.
Notation and Setting Let X' and y denote the space of inputs and responses in classification and regression. Suppose θ = Rd. Let P2 be the 2-Wasserstein space i.e. a set of all probability distributions è with finite second moment equipped with the Wasserstein metric (23]. We identify probability distribution in P2 with its density whenever exits. Denote by O and Ø, respectively, elementwise multiplication and elementwise division of vectors in Rd. Let be the gradient operator. Let V be second-order gradient operator that takes the second derivative at each coordinate i.c. ( = [0e ((0)
f(0)/000 E Rd.
2.1 Wasserstein Gradient Flow
In the Euclidean space, gradient flow of a function f means a curve of points 2; that solves differential equation (d/di).r (x:) from some initial value zo. That the continuous-time limit of gradient descent, which minimises the function f as f. - co. A Wasserstein gradient flow means a curve of probability distributions f4 minimising a given functional on the 2-ssserrn space Pg from some initial distribution μoo The Wasserstein gradient flow μl is characterised solution of a partial differential equation, known as the continuity equation:
dt 14 V Vw given
140 E P2:
where W θ θ 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 infinite-dimenss dimensional
optimisation of the functional F finite-dimensional particle update (23 The continuity equation (1) can be reformulated as a dynamical system random variable 0g ~ 14, such that
d - 0:
W F(u)) (0c) given Oo mo:
(2)
dt
in the sense that the law μr of such a random variable 0, is a weak solution of the continuity equation. Consider the case where the initial measure 14 is set the empirical distribution fl of N particles system (2) by the Euler method with small step size
{0) i 1- Discretising the continuous yields an iterative update scheme of N particles {@ 1 from step m
0
00
[V w: Fo ((mm))
(3)
W F(Am))))
mm+1
m
3

---

where imm denotes the empirical distribution of the particles { at step m.
1.W
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 the Wasserstein gradient V F( the empirical distribution P.n. For example, the KL divergence F(u)= (L π) with reference distribution и leads to the Wasserstein gradient [Vw]F(p][0) (V log (0) V log n()) ill-defined if µ is an empirical distribution. Hence, the particle update scheme ( is often performed with the estimated or approximated Wasserstein
gradient well-defined for empirical distributions [e.g. 26 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 sequential ensemble method of M weak learners f1, .....,. sas. It iteratively constructs an ensemble F.n. of m weak learners /ı...., J' from step m to M. Given the current ensemble F... at step m, it trains new weak learner fm+1 to construct the next ensemble by
F (x) = (x) + f 1 (x),
where v is a shrinkage hyperparameter called a learning rate. The initial state of the ensemble F (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
The fundamental idea of gradient boosting is to train the new weak learner fm+11 to approximate the negative gradient of the remaining error of the current ensemble Fm. Suppose that loss function L measures the remaining error R2 (m (xx²) T.(F... (ara), y1) for each output vector yi E Rd. The new weak learner fm+1 is fitted to the set {2i,, 9: ) whose target variable 9i is each specified by
9 V (F,.(( E
The target 9 is often called a pseudo residual. For each data input zi the boosting scheme the output of the current ensemble F... (I;) in the steepest descent direction of the error R.(F (I;)
(4 updates
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 performance (33)
In modern gradient-boostinn practice is to use the diagonal (coordinatewise) Newton direction of the remaining error R (F (2:)) in lieu of the negative gradient 9r- The new base leaner fm+1 is instead fitted the set {x,, g. 1 where the negative gradient g, is divided elementwise by the Hessian diagonal h2 given by
libraries, such as XGBoost 21 and LightGBM
22 the standard
h2 = (Fmm)))) E Rd
The target variable @ h4 is the diagonal Newton direction that minimises the second-order Taylor approximation of the remaining error (I' (x:)) for each coordinate independently. Combining the second-order gradient boosting framework with rre-based weak learners has demonstrated exceptional scalability and performance [34,] 35). Although 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 Icarning, where are given a training set of input vectors and output distributions {æ,, 44 ) C. X x Pz. Our goal is to construct a model that receives an input and returns set of N particles whose empirical distribution approximates the output distribution. We specify a loss functional D I.) between two probability distributions such as the KL divergence output distribution μ4; Our idea is to combine gradient boosting with Wasserstein gradient, where we iteratively construct set of N boosting ensembles Fi,...,, F.N from step 0 to M.
measure the remaining error F.(.) D(- 14) for each i-th training
Here, the output (x)) of each n-th boosting ensemble represents the n-th output partiele for an .., FN step m. WGBoost trains set of N input Given the current set of N ensembles Fl .......

---

new weak learners fin +1 1..................... Jm+1 and computes the next set of N ensembles by
F (x)
F (x)
fm+1(x)
+
(5)
F (x)
Fm (a)
(
where // is a learning rate, Similarly to standard gradient boosting, we specify the initial state of N ensembles F) , Fg step m by set of constants. Throughout, denote by Pm, the empirical distribution of the N output particles (z.). ..... F (x;) for each i-th training input z;-
As discussed in Section 2.T the Wasserstein gradient often needs be estimated for empirical distributions. For better presentation, let G((n) denote an estimate of the Wasserstein gradient (w) of the i-th remaining error (1),, which is well-defined for any distribution jt. the Wasserstein gradient Fw (u) is originally well-defined for any distribution it is a trivial choice of the estimate, i.e., (1())) = V F. Otherwise, any suitable estimate can be used as Gi()). The foundamental idea of WGBoost is to train the n-th new learner fm+1 to approximate the estimated Wasserstein gradient =)((m,,)) evaluated at the n-th boosting output Fm (x;) for each 2i, so that,
f (2:4)
GG (imm.)) (Fi (x:))
f (a:))
[9s (Pm, (FN (z:))
For each data input xi, the boosting scheme (5) approximates the particle update scheme (3) for the output particles F1 (x:)..... F (I;) under the estimated Wasserstein gradient. The output particles updated in the direction to decrease the remianing error F, mmm.)) D( (1,) at each step m.
Algorithm I |summarises the general procedure of WGBoost. See Figure for illustration of WGBoost. In Section 3 we choose the KL divergence a loss functional D and use kernel smoothing estimate of the Wasserstein gradient. See Appendix A for the Wasserstein gradient of other divergences.
38 uses only 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 large.
Remark 1 (Stochastic WGBoost). Stochastic gradient boosting
Remark 2 (Second-Order WGBoost). If any estimate of the Wasserstein 'Hessian' of the remaining еrror F. available, the Newton direction of F; may also be computable [e.g. 33 40). Implemention of second-order WGBoost algorithm is immediate by plugging such a Newton direction into G, (u) in Algorithm Our default WGBoost algorithm for tree-based evidential learning built on diagonal approximate Newton direction of the KL divergence, aligning with the standard practice in modern gradient-boostig libraries to use the diagonal Newton direction.
Algorithm 1: Wasserstein Gradient Boosting
Input: training set {z,, p4 p of imput x, X' and output distribution μ4 P2 Parameter :loss D, estimate (1) of the Wasserstein gradient Vw D(p (4), particle number N, iteration M, learning rate v, weak learner f, initial constants (0)....,0
Output: set of N boosting ensembles (F) (F...... .......FN. FN at final step M
(F)( ())
FN (-)) + (0) (O.....,,N) (N)
D> set initial state of N boosting ensembles
for 0, ..... M do
for i +- L..... . D do
im, f- empirical distribution of output values (F)), (x:;), F. (a:)) for input oq
for n + I,...,N do
9
(P))) (FR (x:))
compute target value of n-th new weak learner
end
end
for n
L....,Ne N do
f +- fit( {xi g. }==1 F (-) t Fn + ff
D n-th new weak learner D set next state of n-th boosting ensemble
end
end
5

---

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-Iever posterior p(0 yi) of response distribution p(y 0) used as the output distribution μ4; of the training set { H4g 1.2 1--1- Section 3.1 recaps derivation of the individual-level posterior p(@ ys).
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 widely-used estimate of the Wasserstein gradient of the KI. divergence based on kernel smoothing [26]. further advantage of the kernel smoothing estimate is that the approximate Wasserstein Hessian is available, with which Section 3.4 estaabiiiss a second-order WGBoost algorithm similarly to modern gradient-boosting libraries.
3.1 Derivation of Individual-Leve Posteriors and Predictive Distribution
Suppose that a response distribution p(y 0) of a response variable y is specified, as is typically done for probabilistic prediction. Suppose also that a prior distribution (0) of the response parameter 0 is specified for each individual data input I₁. For each individual data pair (Iz, ys). the response- distribution likelihood p(ys 0) and the prior (0) determine the individual-level posterior
p(0 v:) x p(us (0)
by Bayes' theorem. This individual-level posterior set to the output distribution μ4; the training set {2i, _ of WGBoost. The framework of WGBoost then constructs model that returns particle approximation of the output distribution (n(+)) = p(- ys) for each data input s,,
For new input x, the constructed WGBoost model provides a set of particles (z). 0% (z)) as a distributional prediction p(0 x) of the response parameter 6. We can define predictive distribution p(y æ) of the response y for the new input a via marginalisation of the output particles:
N
/ ply 0)p(0) x)d0 = Ep(ulo'
1 N
ply 12) x))
We can also define a point prediction ÿ for the new input x via the individual- Bayes action ÿ argmingey So Uly ())(() z))8, which minimises the average of some error U 3 x Q + R. For example, the Bayes action ý is simply the mean of the output particles if U(y 0) (y 0)2.
In general, the explicit form of the lndividual-level posterior p(0 y4) 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(0 y:). Our algorithm depends only on the log-gradient of the individual- level posterior log p(0 y:) that cancels any constant term by the gradient. Hence, knowing the form of the
likelihood p(y, ()) and the prior (0) suffices.
Remark (Difference from Bayesian Learning). Given response distribution ply [0) and a model 0 f'xx, w) with the parameter w, Bayesian learning of the model ʃ means the use of the posterior p(w {z,, 3: } over w conditional on all data. The predictive distribution p(y x) of the response y defined via marginalisation over w: Je p(y
[ar w)) {æz, y1 )dw.. In contrast, WGBoost returns distributional prediction p(0 x) of the response parameter
circumventing the marginalisation over the model parameter w that can be ultra-high dimensional.
3.2 Choice of Individual- Priors
The prior distribution (f) of the response parameter 0 is specified at each individual data input ay The approach to eliciting the prior may differ, depending on whether past data are available. If past data available, they can be utilised to elicit 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. [1] for the introduction. To avoid numerical errors, 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-Scal. Response). Consider scalar-valued response variable y E IR for regression. A normal location-scall response distribution Ny m, a)) has the mean and scale parameters IR and E (0, oo). A typical noninformative prior of m and are given by,

---

respectively, 1 and 1/ which are improper. At every data point (2),3t), we use a normal prior N(m 0. 0o) over m and an inverse gamma prior [G ag. 56) over o, with the hyperparameters Øo = 10 and Og = 30 = 0.01, which approximate the non-informative priors.
Example 2 (Categorical Response). Consider a label response variable y E {1...., }} for k-class classification. A categorical response distribution C(y q) has the class probability parameter q == (q..... qs) ) in the k-dimensional simplex 4 If = 2, it corresponds to the Bernoulli distribution. A typical noninformative prior of q is given by 1/(11 x... x qs) which are improper. At every data point (#i, ys), we use the logistic normal prior-a multivariate generalisation of the logit normal distribution 42 -vvrr q with the mean 0 and identity covariance matrix scaled by 10.
Remark 4 (Reparametrisation and Standardisation). Section supposed Q = R for some dimension without no loss of generality. Any parameter that lies subset of the Euclidean space (e.g. a in Example I can be reparametrised as one in the Euclidean space (e.g. log Appendix [ details the reparametrisation used for Examples [ 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 KI. divergence KT,(1/ (4:) as loss functional of WGBoost. One challenge of the KL
divergence is that the resulting Wasserstein gradient (c) (0)
(V log (0 V log μ(0) not well-defined when ju is an empirical distribution. A particularly successful solution-
originates in [43] and has been applied in wide contexts (26] 44 45 is to smooth the Wasserstein
gradient through a kernel integral operator Je [gK] (1*) (e, 0 (dp) [46]. By integration-by
part (see [c.g. (33). the smoothed Wasserstein gradient, denoted 9 (n), falls into the following form that is well-defined for any distribution J.C.
[ (10) (0) == -Eg-~
V log pu (0")(((00 0) + Vk(O", 0] E R.
where Vk(0", 0) denotes the gradient of with respect the first argument 6", An approximate Wasserstein gradient flow based on the smoothed Wasserstein gradient G (p) is called the Stein variational gradient descent 43 or kernelised Wasserstein gradient flow 47 In most cases, the kernel l is set to the Gaussian kernel l:(0, 0*) exp( xxp(-110 1|0 0" I2²/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 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 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 (f),, the approximate Wasserstein Hessian of the KI. divergence KT.( 14: can be obtained through the kernel smoothing. The diagonal of the approximate Wasserstein Hessian, denoted Hj (14), is defined by
(n)) (0) = Egg
V (og Vk 0*) @ Vk(0, or)] E RR..
(8)
The diagonal approximate Wasserstein Newton direction of the KL divergence is then defined by provides the derivation based on [39 who derived the Newton direction of the KL. divergence in the context of nonparametric variational inference.
- [9] (1)) (-) o [?)] (44) Appendix C
The second-order WGBoost algorithm is established by plugging it into Ge(rc) in Algorithm that is,
[Ga(p)) = [9] (A) @ [MC (-)
(9)
Algorithm| I Junder the choice is considered our default WGBoost algorithm for evidential learning. We refer this algorithm to as the besseretein-boosteed evidential learning (WEvidential). The explicit pseudocode is provided in Algorithm for full clarity.
7

---

Algorithm 2: Wasserstein- Boosted Evidential Learning
Input: dataset (I,, Yi R of input z; and response yi of classification or regression
Parameter inniividual-level posterior p(0 yy)) of response distribution ply 0), particle number
N, iteration M, learning rate v, weak learner f. initial constants (v0 J
Output: set of N boosting ensembles (I : FM) at final step M
(F ... P6Y (-))
(0)----. ,0)
set initial state of N boostings
for m - 0, M 1 do
for <-1 ,.... D do
(0) ..... 0 f (F (z.)
PN (a:))
> get output particles for i-th data input
for n +- 1, ,,...,, do
g. t N Σ>
V p(0 D:)a(( 03) + Vk(0,, (?)
1=
h ΣN
k=1 logpp(0 11)k(e,, 07) VA(0,, 03) @ VK(0,
end
end
for n2 f 1,...,I1 .N do
.....,,N..
f 1m+1 t fit( {x,. Ø h m+ (.) < Pm (.) + vs (1)
>> fit n-th new weak learner >> set next state of n-th boosting
end
end
Remark 5 (Computation). The diagonal Newton direction ( has the computational complexity ONN x d) same as that of the smoothed Wasserstein gradient. Hence, there is essentially no reason not 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 x d) matrices that result in the complexity up O x d)). Appendix presents a simulation study to compare computational time and convergence speed of four WGBoost algorithms built on different estimates of the Wasserstein gradient.
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 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 /tttts:://githu/t com/ takkuonatsubarr//WBBooooo
Common Hyperparameters Throughout, we set the number of output particles N to 10 and set each weak learner to the decision tree regressor 50) with maximum depth 1 for Section 4. and 3 for the rest. We set the learning rate v to 0.1 for regression and 0.4 for classification, unless otherwise stated. Appendix F contains further details, including a choice of the initial constant {00 J==1-
4.1 Illustrative Conditional Density Estimation
We illustrate the output of WEvidential by estimating conditional density ply x) from one- dimensional scalar inputs and outputs {s;, ys ) 'he normal output distribution N N vn, or) and the prior (,,,, or) in Example I were used to define the individual-level posterior p(m, or y:). in which case the output of the WGBoost algorithm is set of 10 particles { {(m"(() o" (x)) of the mean and scale parameters for each input æ. We chose the number of weak learners M, drawing on an carly-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 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.
8

---

0.1
0.1
-0.1
0.1
Figure 3: Conditional density estimation for the bone mineral density dataset (grey dots) by WEvi- dential, where the normal response distribution N(y m, m. a) is used for the response variable y. Left: distributional estimate (10 particles) of the location parameter (m"( ) for each input. Right: es- timated conditional density (6) through marginalisation of the output particles {(m" (x), (0))o" (=))) n=1
4.2 Probabilistic Regression Benchmark
We examine the regression performance of WEvidential using standard benchmark protocol that originated in [53] and has been used in a number of subsequent works [10] [ 32]. The benchmark
protocol uses real-world tabular datasets from the UCI machine learning repository [54), each with one-dimensional scalar responses. As Section 4.1 the normal response distribution N m. and the prior Ps (m,,or 0r) in Example 1 were used to define the individual-level posterior p(m, y4).
followed the data splitting protocol in [53] and randomly held out 10% each dataset as a test set. The negative log likelihood (NLL) is measured by using the predictive distribution (6). The root mean squared eтror (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 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, subsampled 10% of data to fit each weak learner and used the learning rate 0.01 due to the large dataset size.
Table compares the performance of WEvidential with five other methods: Monte Carlo Dropout (MCDropout) [9], Deep Ensemble (DEnsemble) [[O], Concrete Dropout (CDropout) (55). Natural Gradient Boosting (NGBoost) [32), and Deep Evidential Regression (DEvidential) [I3] Appendix
B
provides further details on the experiment and a limited yet additional comparison. The WGBoost algorithm achieves the best score score sufficiently close to the best score most often.
Table 1: The NLLs and RMSEs for each dataset, where the best score is underlined and the scores
whose standard deviation ranges include the best score are in bold. Results of MCDropout, DEnsem-
bles, CDropout, NGBoost, and DEvidential were reported in 9 10,, (55). 32 and [33 respectively.
Dataset Criteria
WEvidential
MCDropout
DEnsemble
CDropout
NGBoost
DEvidential
boston
2.47 + 0.16
2,46 0.06
2.41 #
2.72
2.43 0.15
2.35 0.06
concrete
2.83
3,04
3.06
3.00
3.11
3.01
energy
0.53 + -0.44
1.99
1.38
2.30
0.45
0.09
kin8nm naval
-0.95
-1.20
-0.65
-0.09
1.04
NLL
5.447
-3.00
-5.63
0.08
-..44
5.07
power
2.60
2.80
2.79
2.75
0.79
2.07
protein
2.70
2.89
2.83
2.00
2.00
2.63
wine
0.95 0.08
0.93
0.94
1.00
0.01
0.89 0.05
yacht
0.24
1.55
0.28
1.75 0.00
0.20
1.03 0.19
year msd
3.50 NA
NA
3.35 NA
NA# NA
3.43 NA
NA NA
boston
2,78 + 0,60
2.99
3,28 1.00
2.65 0.17
2.94
3,06 0.16
concrete
4.15 0.52
5.23
6.03 + 0.58
4,46 0.16
5.06
0.661
5.85 0.15
energy
0.42 0.07
1.66 0.04
2.09 + 0.29
0.46 0.02
0.46 0.06
2,06 0.100
kin8nm naval
0.15 + 0.00
0.00
0.09 + 0.00
0.07 0.00
0.00
0.09 0.00
RMSE
0.00 +: 0.00
0.01 0.00
0.00 + 0.00
0.00 0.00
0.00
0.00 0.00
power
3.19 0.25
4.02
0.11
3.70 0.04
3.798
4.23 0.09
protein
4.09 1: 0.02
4.06
0.06
3.85 0.02
4.33
4.64 0.03
wine
0.61 0.05
0.62 0.01
0.64 0.04
0.62 0.00
0.04
0.61 0.02
yacht
0.48 + 0.18
1.11 0.09
1.58 0.48
0.57 0.05
0.20
1.57 0.56
year msd
9.11 + NA
8.85 + NA
8.89 NA
NA NA
NA
NA NA

---

Table The classification accuracies and OOD detection PR-AUCs for each dataset, where the best score is underlined and in bold. The results other than WEvidential were reported in [14]
Dataset
Criteria
WEvidential
MCDropout
DEnsemble
DDistillation
PNetwork
Accuracy
96.57 0.6
95.25 0.1
97.27 + 0.1
96.21 = 0.1
96.92 + 0.1
segment
OOD
99.67 + 0.2
43.11 1
58.13 + 1.7
35.83 = 0.4
96.74 + 0.9
Accuracy
99.54 + 0.1
89.
99.37 = 0.0
93.66 - 1.5
99.52 = 0.0
sensorless
OOD
81.13
40.61 # 0.7
50.62 + 0.1
#1.1=002
88.65 + 0.4
4.3
Classification and Out-of-Distribution Detection
We examine the classification and anomaly OOD detection performance of WEvidential on two real- world tabular datasets, segmen! and sensorless, following the protocol used in [14]. 'hhe categorical response distribution C(y q) and the prior pz(4) in Example 2 were used to define the individual-
level posterior p(q yı), which case the output of the WGBoost algorithm is of 10 particles
{q"]10 {g" In=1 of the class probability parameter q in the simplex A & for each input We 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 the 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 test set to measure the classification accuracy. There exist several ways of defining OOD score for each input [56). For the WGBoost algorithm, the inverse of the maximum norm of the output-particnl 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, DEnsem ble, and Distributional Distillation (DDistillation) [57], and Posterior Network (PNetwork) [14]. Appendix EeaendixEpproveeeeeee provides further details on the experiment. Figure Figure@eexemliffes 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.
1.0
1.0
brobaiiity
probabiity
0.6
0.5
0.2
0.2
dlas 1
class
<loss 3
class
claso
clocs
dlass
C.SS
cl355 3
closs 4
closs
class
(a) non-OOD input class
(b) OOD input
Figure 4: Examples the output particles (red dot) of WEvidential on the segment dataset, where the coloured area indicate the kernel density estimation of the output particles for each class.
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, is the case of standard gradient boosting. These questions require careful examination and are critical for future study.
10

---

References
[1] Ravid Shwartz-Ziv and Amitai Armon. Tabular data: Deep learning is not all you need. Information Fusion, 81:84-90, 2022.
[2] 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.
[3] Eric Topol. Bigh-performance medicine: the convergence of human and artificial intelligence. Nature Medicine. 25:44-50 56, 2019.
[4] Sorin Grigorescu, Bogdan Trasnea, Tiberiu Cocias, and Gigel Macesanu. A survey of deep learning techniques for autonomous driving. Journal of Field Robotics, 37(3):36 2020.
[5] 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.
[6] Christopher Burges. From ranknet to lambdarank lambdamart: An overview. Learning, 11, 2010.
[7] Byron P. Roe, Hai-Jun Yang, Ji Zhu, Yong Liu, Ion Stancu, and Gordon McGreeor. 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.
[8] James Bennett and Stan Lanning. The Netflix prize. In Proceedings of the KDD Cup Workshop 2007, pages 3-6, 2007.
the
[9] Yarin Gal and Zoubin Ghahramani. Dropout as 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.
[10] Balaji Lakshninarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems, volume 30, 2017.
[11] Murat Sensoy, Lance Kaplan, and Melih Kandemir. Evidential deep learning to quantify classification uncertainty. In Advances in Neural Information Processing Systems, volume 31, 2018.
[12] 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.
[13] Alexander Amini, Wilko Schwarting, Ava Soleimany, and Daniela Rus. Deep evidential regression. In Advances in Neural Information Processing Systems, volume 33, pages 14927- 14937, 2020.
[14] Bertrand Charpentier, Daniel Zügner, and Stephan Günnemann. Posterior network: Uncertainty estimation without OOD samples via density-basee bbaed pseudo-counts. In Advances in Neural Information Processing Systems, volume 33, pages 1356-1367. Curran Associates, Inc., 2020.
[15] Dennis Thomas Ulmer, Christian Hardmeier, and Jes Frellsen. Prior and posterior networks: survey on evidential deep learning methods for uncertainty estimation. Transactions on Machine Learning Research, 2023.
[16] Edouard Capellier, Franck Davoine, Veronique Cherfaoui, and You Li. Evidential deep learning for arbitrary lidar object classification in the context of autonomous driving. 20/9 IEEE Intelligent Vehicles Symposium (IV), pages 1304-1311, 2019.
11

---

[17] Patrick Hemmer, Niklas Kühl, and Jakob Schöffer. Deal: Deep evidential active learning for image classification. 2020 19th IEEE International Conference on Machine Learning and Applications (ICMIA), pages 865-870, 2020.
[18] 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, (8):1336-1367, 2021.
[19] Jakob Gawlikowski, Sudipan Saha, Anna Kruspe, and Xiao Xiang Zhu. An advanced Dirichlet prior network for outto--disttibution detection in remote sensing. IEEE Transactions on Geoscience and Remote Sensing, 60:1-19, 2022.
[20] 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.
[21] Tianqi Chen and Carlos Guestrin. XGBoost: scalable tree boosting system. Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD 16, page 785-794, 2016.
[22] 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.
[23] C'edric Villani. Topies in Optimal Transportation. Americal Mathematical Society, 2003.
[24] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows In Metric Spaces and in the Space of Probability Measures. Birkhäuser Basel, 2005.
[25] Filippo Santambrogio. [Euclidean, metric, and Wasserstein] gradient flows: an overview. Bulletin of Mathematical Sciences, 7:87-154, 2017.
[26] Qiang Liu. Stein variational gradient descent as gradient flow. In Advances in Neural Information Processing Systems, volume 30, 2017.
[27] 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.
[28] Yifei Wang, Peng Chen, and Wuchen Li. Projected Wasserstein gradient descent for high- dimensional Bayesian inference. SIAM/ASA Journal Uncertainty Quanifiicatio, 10(4):1513- 1532, 2022.
[29] Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of Fokker-Planck equations through gradient-log-densit estimation. Entropy (Basel), 22(8):802, 2020.
[30] Ye He, Krishnakumar Balasubramanian, Bharath K. Sriperumbudur, and Jianfeng Lu. Regular- ized Stein variational gradient flow. arXiv:2211.0786) 2022.
[31] Jerome H. Friedman. Greedy function approximation: A gradient boosting machine. The Annals of Statistics, 29(5): 1 189-1232, 2001.
[32] Tony Duan, Avati Anand, Daisy Yi Ding, Khanh K. Thai, Sanjay Basu, Andrew Ng, and Alejan- dro 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.
[33] Peter Bühlmann and Torsten Hothorn. Boosting Algorithms: Regularization, Prediction and Model Fitting. Statistical Science, 22(4):477 - 505, 2007.
[34] Leo Grinsztajn, Edouard Oyallon, and Gael Varoquaux. Why do tree-based models still outperform deep learning typical tabular data? Advances in Neural Information Processing Systems, volume 35, pages 507-520, 2022.
[35] Piotr Florek and Adam Zagdański. Benchmarking state-of-the-ar gradient boosting algorithms for classification, arXivv2305.11009 2023.
12

---

[36] 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-316) 2021.
[37] 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.
[38] Jerome H. Friedman. Stochastic gradient boosting. Computational Statistics & Data Analysis, 38(4):367-378, 2002.
[39] 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.
[40] Yifei Wang and Wuchen Li. Information Newton's 's flow: second-order optimization method in probability space. w0Xiv:2001.04341 2020.
[41] Malay Ghosh. Objective priors: An introduction for frequentists. Statistical Science, 26(2):187- 202, 2011.
[42] J. Aitchison and S. M. Shen. Logistic-normal distributions: Some properties and uses. Biometrika, 6222):261-272, 1980.
[43] Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose Bayesian inference algorithm. In Advances in Neural Information Processing Systems, volume 29, 2016.
[44] 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
[45] 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.
[46] Anna Korba, Adil Salim, Michael Arbel, Giulia Luise, and Arthur Gretton. non-asymptotic analysis for Stein variational gradient descent. In Advances in Neural Information Processing Systems, volume 33, pages 4672 4682, 2020.
[47] 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.
[48] Andre Wibisono. Sampling as optimization in the space of measures: The Langevin dynamics 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.
[49] Gareth O. Roberts and Richard L. Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli, 2(4):341 - -33,, 1996.
[50] Leo Breiman, Jerome Friedman, R.A. Olshen, and Charles J. Stone. Classification and Regres- sion Trees. Chapman and Hal//CRC, 1984.
[51] Trevor Hastie, Robert Tibshirani, and Jerome Friedman. The Elements of Statistical Learning. Springer New York, 2009.
[52] Sanford Weisberg. Applied Linear Regression. John Wiley & Sons, 1985.
[53] Jose Miguel Hernandez-Lote 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.
13

---

[54] Dheeru Dua and Casey Graff. UCI machine learning repository, 2017.
[55] Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. In Advances in Neural Information Processing Systems, volume 30, 2017.
[56] Jingkang Yang, Kaiyang Zhou, Yixuan Li, and Ziwei Liu. Generalized out-of-distribution detection: A survey. 1rXiv:2/10.11334, 2024.
[57] Andrey Malinin, Bruno Mlodozeniec, and Mark Gales. Ensemble distribution distillation. In International Conference on Learning Representations, 2020.
[58] Filippo Santambrogio. Optimal Transport for Applied Mathematicians. Birkhäuser Cham, 2015.
[59] Mingxuan Yi and Song Liu. Bridging the gap between variational inference and Wasserstein gradient flows, 2023.
[60] Michael Arbel, Anna Korba, Adil Salim, and Arthur Gretton. Maximum mean discrepancy gradient flow. Advances in Neural Information Processing Systems, volume 32, 2019.
[61] Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker- Planck equation. SIAM Journal Mathematical Analysis, 29(1):1-17, 1998.
[62] Grigorios A. Pavliotis. Stochastic Processes and Applications. Springer New York. 2014.
[63] Vern I. Paulsen and Mrinal Raghupathi. Introduction to the Theory of Reproducing Kernel Hilbert Spaces. Cambridge University Press, 2016.
[64] Alex Leviyev, Joshua Chen, Yifci Wang, Omar Ghattas, and Aaron Zimmerman. A stochastic Stein variational Newton method. arXiv:2204.0000 00009,, 2022.
[65] Alex Smola, Arthur Gretton, Le Song, and Bernhard Schölkopf. Hilbert space embedding for distributions.
[66] Shun ichi Amari. Information Geometry and Its Applications. Springer Tokyo, 2016.
[67] Pavel Izmailov, Wesley Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Subspace inference for bayesian deep learning. In Proceedings of 35th Conference on Uncertainty Artificial Intelligence, 2019.
14

---

Appendix
This appendix contains the technical and experiment details referred to the main text. Appendix A recaps derivation of the Wasserstein gradient. Appendix discusses the variant of WGBoost built on the unadjusted Langevin algorithm. Appendix C derives the diagonal approximate Wasserstein Newton direction used for WEvidential. Appendix proovides simulation studies kernel choice of WEvidential and comparison of WGBoost algorithms built on different estimate of the Wasserstein gradient. Appendix 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 function on e called the first variation 24 The first variation (F(()) of the functional F at is a function on e that satisfies
F(u + EV) -
/ 6F(m) (0))(O)d0 8μ
lim ++
for all signed measure v s.t. lr + av E P2 for all sufficiently small. The Wasserstein gradient F(js of the functional F at is derived as the gradient of the first variation (see [e.g. (24))
V F()](()) = V 6F((H) (0).
8μ
is common to suppose that the functional F consists of three energies, which are determined by functions U R - R. V θ - R, and W : O - IR respectively, such that
1 F()) U()(0))dO ((0) + jxx W(O (0))(0))0(((())))))
internal energy
potential energy
interaction energy
For a functional F that falls into the above form, the Wasserstein gradient is derived as
FwJ) (0) VU (v())) VV + / VW(B ()))((')
where U' is the derivative of U R
The KL divergence F(u) KL(u T) of distribution æ falls into the form with Uxx) = x log H, V(0) = log (0), and ( O, where
R 23
KL T) / log log)(0))(()d0)
log - (0)
Table 3 presents examples of Wasserstein gradients of common divergences F(u) = D(H IT).
In the context 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 V log n (0) Vx (0) 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.
Table 3: Wasserstein gradients of four divergences: the KL divergence 58 the chi-squared diver-
gence [47], the alpha divergence [59], and the maximum mean discrepancy 600
Divergence F(u) D(p
Wasserstein gradient [w)) (0)
KL( T)
V log (0) - V log (a0))
Chi((u x)
(V(x(0)=(())
Alpha( T)
(n(0) (=00)=(0)=-
MMD(u
))
So (1(0,0')(()dd00 So vk(0,0")<(0)ll
15

---

B Langevin Gradient Boosting for KL Divergence
If chosen functional F on P2 is the KI. divergence F(a) (T. x) of target distribution x, the continuity equation (T) admits an equivalent represttatinn as the Fokker-Planck equation [61]
d dt
14 V (ee ( Vog + Ap4 given
No P2
(10)
where A denotes the Laplacian operator. Recall that the original continuity equation 1 can be reformulated as the deterministic differential equation 2 of a random variable 9 ~ f4g. In contrast, the Nokker-Planck equation (00) can be reformulated stochastic differential equation of random variable 0, ~ fl,, known as the overdamped Langevin dynamics [62]:
d0, v (lg + V2dB, given 8o ~ Poo
(11)
where B, denotes a standard Brownian motion. Note that the deterministic system 2 the case of the KI. divergence and the above stochastic system (1I are equivalent at population level, in a sense that the law of the random variable 0, 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 μl (TT) to the empirical distribution Ño of N initial particles {0h Jn=z- Discretising the stochastic system (TT) by the Euler-Maruyama method with step size > yields a stochastic update scheme of particles {6%h }.N from step m
01
01
V log x (6 + /2/vg'
m++1
AN
V log π (0%) + v 2/v N.
where each ed denotes a realisation from a standard normal distribution on IRd The above updating
scheme of each n-th particle is known as the unadjusted Langevin algorithm 49
We can define
variant of WGBoost by replacing the term G((p) in Algorithm ! with log p.(.) + 2/v / & where H4 is an output distribution at each and &a is a realisation from a standard normal distribution. The procedure is summarised in Algorithm which we call Langevin gradient boosting (LGBoost).
Algorithm 3: Langevin Gradient Boosting
Input: training set (Ii, p M of input x, E X' and output distribution μ4 E P2 Parameter: :particle number N, iteration M, rate v, weak learner f. initial constants (1) ..., ,.,)
Output: set of N boosting ensembles (F)r.... FN) at final step M
(P)()..... P/N (-)) + (0)....,0)) 0
for m + 0, ,
do
for 1, ,,... do
for - 1..... do
9.h t V log (f + 22/2 where er ((O,Ia))
end
m++1 +- fit (Tit oe )00 D
Fi=1
FE F f Fn (.) f
end
end
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 different terminology ciimiooogyy-iiplyy.. the Newton direction-frre a viewpoint of nonparametric variational inference. We place their result in the context of approximate Wasserstein gradient flows. Appendix C.I 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.
16

---

C.1 Smoothed Wasserstein Gradient and Hessian
Consider the one-dimensional case θ R for simplicity. For a map T: R and a distribution P2, let 14 be the pushforward of je under the transform 0 - 0 (T) defined with time- variable 1 E IR. This means that μl is distribution obtained by change-of-variable applied for μ.. The Wasserstein gradient of functional F(u) can be associated with the time derivative (d/di)F((w) [33. In what follows, we focus on the KL divergence F(u) = KL TT) as a loss functional. Under condition T' E (M),, the time derivative at t 0 satisfies the following equality
d KL(Me dt
(=)=
= -So T(Ø) [g (1)) (0) du(0) (T, (P(D))
(12)
where (2) denotes the Wasserstein gradient of F(P) = KL(H л) with the target distribution ' made implicit. It gives interpretation of the Wasserstein gradient the steepest-descent direction because the decay of the KL divergence at = 0 maximised when T -gki (w).
The 'smoothed' Wasserstein gradient can be derived by restricting the transform map T to a more regulated Hilbert space than L (p)). reproducing kernel Hilbert space (RKHS) H associated with a kernel function & IR x R + R is the most common choice of such Hilbert space [e.g.26]. An important property of the RKHS H is that any function / E H satisfies the reproducing property under the associated kernel & and inner product {-, 1) fr [63]. As discussed in H and exchanging the (e.g.46. applying the reproducing property in (12) under the condition
S(0) = (f(.). k(0,))) ))
integral order, the time derivative satisfies alternative equality follows:
d KL( | m) = So (T(), (4)] (0) du(0)
dt
= (I((), / [gKL(W) (0))(0)))
(T, (T)
(13)
where (9*(4)(( = -== SalgkL (W)) (ddp())))))))(()) corresponds to the smoothed Wasserstein gradient used in the main text. The decay the KL divergence t is maximised by T (p))
Similarly, the Wasserstein Hessian of the functional F(u) can be associated with the second time derivative (d² ((()) (2 As discussed in [e.g. 46 the Wasserstein Hessian of the KI. diver- gence, denoted Hess(u), is an operator over functions T E L (2(() that satisfies
d²
KT. T)
(T,
(14)
dta
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 E H and exchanging the integral order, the second time derivative satisfies an alternative equality as follows:
dc
KL T) == (T(+)))( [Hess" (4)] (+1, T(+2)) H)
(15)
diz
where [Hess" (m)) *2) = (k(*1,), :) Hess( 1) (P)))) is the smoothed Wasserstein Hessian and the symbols and *2 denote the variables to which each of the two inner products is taken.
In the multidimensional case θ Rdd the transport map T is vector-valued function T' Rd + gd where a similar derivation can be repeated by replacing (1²()) and H with the product space independent copies of L (2) and H. It follows from Proposition 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
(+)))))+) Eo-p
log n (0) vk(0,-) Rd
Hess" " (H1) x2) = Eonp
V logm(8)k((B.. +1 (0 x2) + Vh(O, +1) & vk(0,*2)) Radd
where 2 denotes operator to take the Jacobian of the gradient-i.e. i.e., ( is the Hessian matrix of f at 0-and 69 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.
17

---

C.2 Approximate Wasserstein Newton Direction
In the Euclidean space, the Newton direction of an objective function is direction s.t. the second- order Taylor approximation of the function is minimised. Similarly, [39] characterised the Newton direction T R I2d of the KL divergence KL(H I)) as solution of the following equation
([Hesss' (1)](+1, +2), T (+2)) -1 (10()))) ((+)) ()(())(())))())))
for all
H.
Here θ = _2 and H is the product space of d independent copies of the RKHS of kernel k. To obtain closed-form solution, [39 supposed that the Newton direction 7 can expressed in form (.) - (,, 0*) dependent on a set of each particle ( E è and associated vector-valued coefficient 4 E R 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
Σ == [Hess" (4)](0) 0*). W
[ (K)())
(16)
Heess" (M)](N, 0n) W
9 (1))(O))
These equations (16) can rewritten in a matrix form 64 H E R and block vector G E I&K by the following partitioning
Let K == N Define a block matrix
H1
HIN
and
G
HNI
HNN
GN
with each block specified as Hig == Define block matrix K E Ig x and a block vector W E IRK by the following partitioning
Hess" (H) 00) E Rdxd and G; = [S (())))) e R..
K11
KIN
W
K =
and W =
:
K
K N
W
with each block of K specified as K.j == Ia x k(0',
(0')0))) 01) E Rdxd, where L denotes the d x identity
matrix. Notice that W a block vector that aligns the vector-valued coefficients {1²m p=11 Using
N
these notations, the optimal coefficients that solve (16) is simply written W - G [64
Given the optimal coefficients W -- the Newton direction T (0n) evaluated the given particle 0 for each n I,..., N can be written in the following block vector form
(0')
K11
KIN
H
HIN
G1
:
: HNN
(17)
(0))
KN1
KNN
HN1
GN
To distinguish from the standard Newton direction in the Fuclidean space, we call (17) the ap- proximate Wasserstein Newton direction. The approximate Wasserstein Newton direction yields second-order particle update scheme. Suppose we have particles {0%) J.N to be updated at each step m. At each step m, define the above matrices H and with the empirical distribution μı = fm of the particles {0tm J Replacing the Wasserstein gradient in the particle update scheme (3) by the approximate Wasserstein Newton direction (17) provides the second-order update scheme in
C.3 Diagonal Approximate Wasserstein Newton Direction
We derive the diagonal approximation of the approximate Wasserstein Newton direction, which we used for 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 that no matrix product and inversion will involved. Specifically, we replace the matrices K and in I1 by the diagonal approximations K and H. that is,
I
0
h
0
K
and
i
=
I
hNN
18

---

where K = Ia x k(0", 8*) = Ia for the Gaussian kernel & used in this work. and the matrix han E RRdxd denotes the diagonal approximation of the diagonal block H of H.
Recall that Haz [Hess" (O) 0*). Denote by Diag(A) the diagonal of square matrix A. The diagonal approximation hnn is a diagonal matrix whose diagonal is Diag(Han) 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 form
T (0+)
b
0
G
G1 2 Diag (H1))
: (*N)N)
: N
(18)
NNN
G Ø Diag [H NNN)
At an arbitrary particte location 0, denote by [H* (O) the diagonal of the smoothed Wasserstein Hessian [Hess" (O)](0,0) 0). It is straightforward to see that the diagonal be written
[HL [v)] (.) = Eo-n V log π (0)A(0, ) Vk(0,-) o Vk(O,))
Notice that Diag (a))) [7L' (H)) (0") by definition. therefore follows from the formula (88)
with G, [ (p)] and Diag (Hmm) [HL- (4)] (#") that the diagonal approximate Wasserstein
Newton direction at an arbitrary particle location 0 can be independently computed by
- (x)]() Ø (H)) (0).
We used this direction in Section 3 In the main text, the diagonal approximate Wasserstein Newton direction is defined for each loss functional (-).) = D(- 14). with æ μ4, using the smoothed Wasserstein gradient 9 (p) and the diagonal of the smoothed Wasserstein Hessian H17 (n) for each i-th output distribution (4z-
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 KI. divergence in order to built WEvidential. The smoothed Wasserstein gradient originates in approximate Wasserstein gradient flow called Stein variational gradient descent (SVGD) 43]. is fairly common SVGD practice use the Gaussian kernel k:00, 0*) exp(-110 0 0-1 /2/ with the scale 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 simulation study to examine sensitivity to the scale value. We prepared a synthetic dataset {z,, 1. whose inputs are 200 gird points on the interval [-3.5, 3.5] and output distri- butions normal distributions (0) .N sin( 0.5) conditional
each Ii. We used
WEvidential with different kernel scales = 0.001, 0.01, 0.1, 1.0, 10, 100 and fitted it to the syn- 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 ( J==1 of WEvidential was set 10 grid points in the interval [-10, 10).
thetic dataset {a;, /4 ) The decision
We computed the output of WEvidential for 500 grid points the interval [-3.5, 3.5.. We used the maximum mean discrepancy (MMD) (65] to measure the approximation error between the empirical distribution f4; of the output of WEvidential and the output distribution J4z;
MMD2 (fs, (4: Bo- (0, 0') 2B p 0 (O, 0') + Ea-p,. (O, 0'))
where kp is a Gaussian kernel kp 8* exp( (0 0²)²/s) with the scale = 0.025. The total approximation error was measured by the MMD averaged over all the inputs. Figure shows the total approximation enror of WEvidential for each scale value h in the common log scale, together with examples of the output of WEvidential for different values h. demonstrates that the total
error is minimised by h = 0.1 and stays in relatively small value range for 0.01 < h < 1.0.
19

---

-08
Commons 409 MMO
-1.2
-1.4
-18
2
common log kemel banowisth
(a) The total error
(b) h = 0.10
h 0.01
(d) / 100
Figure 5: The total MMD error and example outputs of WEvidential for different kernel scales. Panel
(a): the total MMD error for different scale values h = 0.001, 0.01, 0.1, 1.0, 10, 100 both plotted in the common log scale. Panel (b): the output of WEvidential for h 0.1. Panel (c): the output of WEvidential for h = 0.01. Panel (d): the output of WEvidential for h = 100.
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 (i))) in Algorithm | to, respectively,
1. the smoothed Wasserstein gradient in 77:
2. the diagonal approximate Wasserstein Newton direction in
(8)
3. the full approximate Wasserstein Newton direction in I7
The fourth algorithm is LGBoost in Appendix B which is rather 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 implemented by removing h," and replacing g." @ h with g in Algorithm 2 The full-Newton WEvidential is implemented by replacing of h" in Algorithm with of computed by the following Algorithm where (0) denotes the Hessian matrix a function f:
at 6.
We prepared the same synthetic dataset #:) as Appendix D.I 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 [.1] The decision tree regressor with the maximum depth was used for each weak learner, and the learning rate was set to 0. We used increasing number of weak leaners up 100 weak learners, in order to observe the decay of the approximation error and the increase the computational time. The initial constant {0" ) 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 clear.
Ossoows the approximation error and the computational time of each algorithm with respect 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 x d) (N d) block matrices, where 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 the Gaussian noise used in the algorithm. We therefore recommend our default WEvidential for better performance and efficient computation.
20

---

Algorithm 4: Computation of Full Approximate Wasserstein Newton Direction
Input: input 2i, output distribution Jii outputs of boostings (F (Ti)..... Fm (x;) Output: Wasserstein Newton direction (u:.....; uf") evaluated at (F,, (8.))...,. Fr (2:4)
(0) ..... ) (F)) (24), ..... F. (2:))
for n - 1..... do
Oe - 1
N Σド=1 V log p. (4, 0T) VK(O,.
for k 5- 1,..., N do
11 - N -0
-
0.) 03) Vk(0, VK(0,,0))
Kn <- Ia 6(0,,00)
end
end
K
KP N
H
H
K N1
K K
H
HNW
maximum mean discneiseeeppppp discrepancy
0.20
Arst order WEvidentel Default WEvidential
Common lg computationn timee
0.15
LGBoost
-0.5
-1.0
0.10
-1.5
0.00
1.6
weak learner number
common--og weak leamer number
(a) Approximation Error
(b) Computational Time
(c) First-order
(d) Gongonal-Newtord
(e) Full-Newton
(f) LGBoost
Figure 6: The approximation error and computational time of the four different WGBoost algorithms. of each algorithm measured by the MMD averaged the inputs
Panel (a): the approximation with respect to the number of weak learners. Panel (b): the computational time with respect to the number of weak learners in common logarithm scale. Panel (c)-(f): the outputs of the four algorithms each with 100 weak learners used.
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 provided in the source code. Appendices E.
I CPU, The scripts reproduce all the experiments to E.3 describe additional details of the applications in Sections 4.1 4.3 respectively. Appendix E.4 describes a choice of initial constants { === for the WGBoost algorithm used in Section
E.1 Detail of Section 4.1
The normal response distribution A(y m, o) 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 TR. We reparametrised it as one the Fuclidean space by the log transform o' == log(o), which is the standard practice in Bayesian computation (20). The inverse of the log transform
21

---

is the exponential transform exp(o'). The change of variable formula tells the form of the )') up to the normalising constant. Under the prior in Example | the individual-level posterior on (m, o') conditional on each individual response yi is given by
individual-Ievel posterior on (m, o')
p(m, o' 1.) exp (-2 (yi 1-4-ma m) exp(o')²
1 m 2 10a
0.01
exp
exp(o')LU0 exp
(eeplo')
where the Jacobian determinant |dr/ /dol exp(o") is used for the change of variable. Figure shows the output of WEvidential for the old faithful geyser dataset and the conditional density estimated through (6).
3-
3-
50 60
70
90
50 90
70
80
90
Figure 7: Conditional density estimation for the old faithful geyser dataset (grey dots) by WEvidential. Left: distributional estimate (10 particles) of the location parameter for each input. Right: estimated density by the predictive distribution (6) based on the output partieles.
E.2
Detail of Section 4.2
Section 4.2
used the same reparametrisation of the normal response parameter (m, or) as that of The NLL and RMSE of each algorithm on test data {x,, Di) were computed by
Appendix E.2
1 D NLL = ALL=-D2 Σ 1=1
Logp(uu x.)) and RMSE
L (u D.) D
where p(y | m44) and úJ 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 3/} (y/ train ) /uu with the mean train and standard deviation train of the training responses. Hence WEvidential provided the predictive distribution p(y/ xx)) and the point prediction V for the standardised responses y/- With no loss of generality, the NLL and RMSE for the original responses y; be computed follows:
NLL 1
1 D
log x4))
log
train x4)) log Usid
(19)
D
D Σ (o) - 11) i=1
RMSE
(U=-
(yrriin
+ yrain x x))
train Vaid
(20)
D
1=1
where the equality of the NLL follows from the change of variable p(y: xx) = p(u' | xx)) uain and the equality of the RMSE follows from rearranging the terms.
We provide a brief description of each algorithm used in Section 4.2, For all the algorithms, the normal
response distribution p(y m, c)) specified as the response distribution, where the algorithm produces point or distributional estimate of the response parameter (m, or) at each input x.
MCDropout [9] trains a single neural network F while dropping out each network weight with some Bernoulli probability. can be interpreted as a variational approximation of a Bayesian neural network. It generates multiple subnetworks {F" In= by subsampling the network weights by the dropout. The predictive distribution p(y x)) is given by the model averaging (1/N) EN p(y (m,o) = Fz)) for each input
DEnsemble 10 simply trains independent copies [F*) of a neural network F in parallel. is one of the mainstream approaches to uncertainty quantification based on deep learning. The predictive distribution is given by the model averaging in MCDropout.
22

---

CDropout [5 consider a continuous relaxation of the Bernoulli random variable used in MCDropout optimise the Bernoulli probability of the dropout. It generates multiple sub- networks (P*) by subsampling the network weights by the dropout with the optimised probability. The predictive distribution is the same as MCDropout.
NGBoosting 32 is family of gradient booting that use the natural gradient [66] of the response distribution as a target variable each weak learner. contrast to other methods, NGBoost outputs single value F(x) be plugged into the response-distee distriiieedttiitiinnnn parameter.
The predictive distribution p(y x) is given by p(y (m, σ) = F())) for each input æ.
DEvidential 13] extends deep evidential learning [1], originally proposed in classification settings, to regression settings. considers the case where the individual-leve posterior of the response distribution falls into a conjugate parametric form, and predicts the hyperpa- rameter 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 limited yet additional comparison of WEvidential with other Bayesian neural networks (BNNs) and kernel-based model, Gaussian process (GP). We used a subset of the datasets in Section 4 which were used in [77 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). T summarises the result.
Table 4: The NLLS and RMSEs for each dataset, where the best score is underlined. The results other than that of WEvidential were reported in 67
Dataset
Criteria
WEvidential
VI
DVI
SWAG
PCA+VI
DGP1-50
boston
2.47 + 0.16
2.43 0.03
2.41 0.02
2.76 0.13
2.13
2.33
concrete energy
0.11
3.04 0.02
3.06
3.01 0.09
0.10
3.13
NLL
0.53 0.08
2.38 0.02
1.01 0.06
1.70
1.72 1.59
1.32
0.03
naval yacht
4.03
-5.87 +: 0.29
-6.29 + 0.04
+6.11
-6.11
-3.60
0.24
1.68 0.04
0.03
0.42
0.40
keggd
-0.91 0.04
0.04
-1.09
boston
2.78 0.60
0.98
0.95
concrete
0.52
0.43
5.14
energy
RMSE
0.42
1.59
1.59
naval
0.00
0.00
0.00
yacht
0.48 0.18
0.97 + 0.38
0.38
keggd
0.24
0.13 0.03
0.13
E.3
Detail of Section 4.3
We reparametrised the parameter of the categorical response distribution C(y g) used in Section 4.3 similarly to the normal response distribution used Sections [4. a.2 The categorical response distribution has class probability parameter
(q] ... Zz. in the simplex AІ. We applied the log-ratio transform of := [log (2//)))..... ... log((o- 1ogg(/-1//0 //)) ) I that maps from the simplex A to
the Euclidean space IR - 42
The inverse of the log-ratio transform is the logistic transform
exp((1)
exp(16-1)
E A
where
24 = 1+
exp(d)))
2
2
ZR
The logistic normal distribution the original parameter y corresponds to normal distribution on the transformed parameter o' [4). The change of variable formula tells the individual-level posterior on g' up to the normalising constant. Under the prior in Example 2 the individual-lewe posterior
23

---

p(q' y:) conditional on each individual observed response is given by
1
[w=k]
k-1
exp(())
[v.=
1 1171
pld ve) x
X
II j=1
X exp
2k
Zk
2 102
where [w = .] is 1 y the j-th class label and 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 Ax by 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 [4]
considers the case where the individual-level posterior of the categorical response distribution falls into Dirichlet distribution similarly to deep evidential learning [TT]. 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 xx) quantity computable for any proba- bilistic 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 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 was used. T 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.
Table 5: The OOD detection performance of WEvidential, NGBoost, and RForest on the segment dataset, where WEvidential (Entropy) indicates the result of WEvidential based on the entropy.
Dataset
Criteria
WEvidential
WEvidential (Entropy)
NGBoost
RForest
Accuracy
96.57 + 0.6
96.57 + 0.6
94.09 + 0.6
88.18 18 1.6
segment
OOD
99.67 + 0.2
98.96 + 0.4
89.96 # 1.3
66.51 1.1
To investigate on the effective learning rate WEvidential for classification, we performed 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 the test classification accuracy. We trained WEvidential using different learning rates from 0. to 1.0 while fixing the number of weak learners to 4000. 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 0.4 for classification.
test set to measure
that the result for the learning rate
Table 6: The classification accuracy and OOD detection performance of WEvidential on the synthetic dataset for different learning rates, where the number of weak learners is fixed to 4000.
Criteria
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
Accuracy
92.50
92.50
92.50
93.75
91.25
86.25
90.00
75.00
85.00
80.00
OOD
29.85
34.79
35.26
41.77
40.52
42.69
35.21
35.58
42.04
36.05
24

---

Table 7: The classification accuracy and OOD detection performance of WEvidential on the synthetic dataset for different numbers of weak learners, where the learning rate is fixed to 0.4.
Criteria
1000
2000
3000
4000
Accuracy
91.25
92.50
91.25
93.75
OOD
43.96
37.70
39.84
41.77
E,4 Choice of Initial State of WGBoost
In standard gradient boosting, the initial state at step m specified by constant that fits given data. Similarly, we specified the initial state {10 ) of WEvidential by set of constants that fits the output distributions in average. We find such a of constants by performing an approximate Wasserstein gradient flow averaged over all the output distributions. Specifically,
given the term (1) in Algorithm 1 define ((p) (1/D) E (i()) and perform the update
scheme of set of N particles (0 IN
0 m
(#t (())
+ vo
IN m++11
99 mm
with the learning rate in = 0.01 to the maximum step number m = 5000. The initial particle locations for this update scheme were sampled from standard normal distribution. We specified the initial state {19 1==1 by the set of particles {1im IN obtained though this scheme at m = 5000.
25

---

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 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 long 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 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 (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
26

---

Answer: [NA].
Justification: This work primarily focused on the methodological development of the proposed algorithm, empirically demonstrating the performance the various benchmarks. Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide 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
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
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 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 novel architecture, describing the architecture fully might suffice, or if the contribution 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 accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions 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) the contribution is primarily 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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which 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 have some path reproducing or verifying the results.
5. Open access to data and code
27

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
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.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "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 exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
y /nips. /nips.cc//ppp http:ppblliggggieeeeeeeeeea
for more details.
The authors should provide instructions on data access and preparation, including how 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. only 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) recommended, but including URLs to data and code permitted.
6. Experimental Setting/ettails
Question:
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary 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 the core the paper to 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
28

---

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.
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 96% CL, 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 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
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 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 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 failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws 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 no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
29

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
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 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include 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 package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
30

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 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 included in the source code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part 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 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. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According 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 equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending 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.
31


Reviews:
[{"review_id": "AS90DN492A", "reviewer": "Reviewer_jx13", "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. \n\nPaper'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.", "questions": "I would be curious to know about the scalability and the computational complexity of the proposed algorithm.", "limitations": "NA", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n\nS2: 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.\n\nMisc:\nThere is a typo in the line no 284, root is written as room.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your strong endorsement of our work. We also believe that WGBoost is a promising direction that connects Wasserstein dynamics and gradient boosting ensemble for the first time.\n\n&nbsp;\n\n> Authors did not specify details about the hyperparameter selection for Conditional Density Estimation, Classification and OOD Detection tasks.\n\nThank you for raising this point. We will add a clarification about the hyperparameter selection for the experiments other than Section 4.2.\n\n&nbsp;\n\n> Misc: There is a typo in the line no 284, root is written as room.\n\nThank you for pointing this out. We have corrected the typo.\n\n&nbsp;\n\n> I would be curious to know about the scalability and the computational complexity of the proposed algorithm.\n\nThank you for this point. In Appendix E, we have a simulation study to compare the loss decay and computational time of WGBoost under four different choices of the approximate Wasserstein gradient. We will expand the discussion about the computational cost of the algorithm in the main text (e.g. adding some big O cost order).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eoc5egmyqC", "reviewer": "Reviewer_UM7J", "summary": "This 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.", "questions": "Algorithm 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?", "limitations": "n/a", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Originality: \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.", "weaknesses": "Originality: \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,", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback to our work. We respond to each of your comments and concerns below. We hope that these responses will open up room for an upward revision of the score of our work. (In our response below, citation numbers such as [9] and [10] correspond to references in our manuscript.)\n\n&nbsp;\n\n> For a paper on trees, I find the experiments quite small scale and limited ... / ... large scale uncertainty quantification experiments are lacking ...\n\nThank you for the opportunity for us to strengthen our experiment. We will add new experiments on large scale datasets, e.g. HIGGS, KEGG, Buzz in social media, Year Prediction MSD datasets which all have very large size n>500000. Meanwhile, the current set of the experiments would be a common benchmark in UQ literatures [13, 33, 9, 10] and include medium-large datasets of size n≈50000 or n≈10000 (For example, a relevant gradient-boosted tree UQ method, Natural gradient boosting [33] in NeurIPS, used only those datasets in Section 4.2). We believe that adding the large scale experiment to the current experiments would make our benchmark strong for both the tree and UQ contexts.\n\n&nbsp;\n\n> Comparisons to (approximate) Bayesian neural networks are quite weak ...\n\nThank you for raising the point together with the expert literatures. We will include advanced approximate Bayesian methods beyond MC Dropout in our comparison. The previous manuscript followed our relevant UQ literatures [33, 13, 14] who used MC Dropout and Deep Ensemble in comparison. \n\nHere, we have checked experimental results of the provided literature, Subspace Inference for Bayesian Deep Learning. Their Appendix Table 2&3 shows Log-Likelihood (-1 times NLL) and RMSE scores of several advanced Bayesian methods on some datasets we used. Our method produces the best score in 8 rows out of 10 rows in Table 2&3 combined.\n\n&nbsp;\n\n> The natural missing comparison here is to Gaussian processes and other kernel methods, ... Another missing set of experiments is comparison to quantile regression / pinball loss using trees, ...\n\nThank you for your constructive feedback. We will also include kernel-based and quantile-based nonparametric methods in our experiments. With these methods added, our comparison covers a range of methods: common deep learning UQ methods, other approximate Bayesian methods, other gradient boosting UQ method, and nonparametric regression methods. We believe that this gives us a high-quality benchmark for our proposed method.\n\n&nbsp;\n\n> The code doesn’t seem to provide any clarity here / The current algorithm is so unspecific it’s very hard to follow.\n\nThank you for the opportunity for us to clarify. First, we would like to recap our current paper structure. In Section 2, we firstly explained an abstract framework of WGBoost in Algorithm 1. To use WGBoost, we have to specify \n\n- (a) a loss functional over probability densities;\n- (b) how to estimate/approximate the Wasserstein gradient of the chosen loss. \n\nThese (a) and (b) are user-specified components of WGBoost. Then, in Section 3, we presented the specific setting where we use (a) the KL divergence and (b) the kernel-smoothing estimate of the Wasserstein gradient, for our application. So it was our intention that we described the algorithm in a general manner in Algorithm 1, and then specified (a) and (b) later in Section 3. \n\nWe will more strongly clarify that WGBoost in Algorithm 1 takes (a) and (b) as user-specified components. In Appendix, we have Algorithm 3 that shows the explicit algorithmic table for the specific setting in Section 3. We will clarify this as well. Finally, the terminology \"exact or approximate Wassetstein gradient\" would be confusing so we will change it to \"estimated Wasserstein gradient\". \n\n\n&nbsp;\n\n> My understanding of Wasserstein particle flows is that the particles should interact in some manner during the gradient step. / Is the Wasserstein gradient where the particles manage to interact? \n\nThank you for checking this point with us. In principle, interaction is not compulsory part of Wasserstein particle flows. Whether particle interaction happens or not depends on how to estimate the Wasserstein gradient. If it happens, it is beneficial to have some enforced dispersion of particles. \n\nFor example, the JKO-scheme-based estimation of the Wasserstein gradient does not have interaction term. Langevin Monte Carlo---a stochastic formulation of Wasserstein flow--does not have interaction term neither. On the other hand, our case of the kernel smoothing estimate in Section 3 indeed induces particle interaction. So interaction is not compulsory but the WGBoost under the setting of Section 3 for our application indeed has particle interaction thanks to the choice of the estimate.\n\n&nbsp;\n\n> Table 1: where do the error bars in the baseline methods come from? / Was the same preprocessing done / e.g. same train / test split for these?\n\nWe followed a de-facto standard protocol of [53] to prepare 20 different patterns of training/test sets from the datasets in Table 1. The error bar was created by computing the test error 20 times using the 20 different patterns.\n\n&nbsp;\n\n>  it seems like a mean / variance two stage fit model would be the natural “base” method to compare to ...\n\nWGBoost produces a particle-based \"distributional\" estimate of the mean and variance of the Gaussian-noise regression for each input. Indeed a method that produces a point estimate of mean and variance would be a good base model, and NGBoost is such a tree method in our comparison. We will further investigate on the two stage model too.\n\n&nbsp;\n\n> What would entropies of a traditional multiclass classification tree perform like in terms of an out of distribution detector here?\n\nWe didn't include this in our manuscript but observed that entropies of a traditional classifier shows poor OOD performance during the development of our method. We will add a discussion on this point.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6V97MCQPtt", "reviewer": "Reviewer_q9Jy", "summary": "This 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.", "questions": "I wrote questions in the Weakness.", "limitations": "The limitation is unclear. As far as I read no formal description is presented.", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your assessment of our work. We are afraid that there seems a misinterpretation of our method in the reviewer's concerns: (a) the proposed method seems almost identical to the stein variational gradient descent (SVGD) and (b) the paper lacks the comparison with existing work. Please let us recap the high-level setting of WGBoost in the response that follows next. We hope this clarification will assist in the reassessment of our work and open up room for an upward revision of the score. \n\n&nbsp;\n\n> the proposed method seems almost identical to SVGD / the paper lacks the comparison with existing work\n\n**While SVGD is a sampling method, WGBoost is a new gradient boosting method**. Please note that our setting is **not** about sampling from a posterior over model parameters (e.g. BNN). To see this, we compare (i) learning of BNN and (ii) our setting here:\n\n- **BNN Learning**: We have a dataset $( x_i, y_i )\\_{i=0}^{n}$ of input vector $x_i$ and output vector $y_i$, by which we have a posterior $P(w \\mid ( x_i, y_i )\\_{i=0}^{n} )$ over the network parameter $w$. The goal is to approximate $P(w \\mid ( x_i, y_i )\\_{i=0}^{n} )$ well by sampling or VI;\n- **Our Setting**: We have a dataset $( x_i, \\pi_{x_i} )\\_{i=0}^{n}$ where $\\pi_{x_i}$ is a probability distribution given at each input $x_i$. Our goal is to have a learning algorithm that can predict the unseen distribution $\\pi_x$ for any new input $x$.\n\nOur setting is clearly a less common, challenging regression problem because our target output is now a probability distribution $\\pi_{x_i}$ at each input $x_i$ rather than some finite-dimensional vector output $y_i$. WGBoost is designed to solve this problem. Please find *Section B in our global rebuttal* for more detailed recap.\n\nConnection between WGBoost with SVGD is as follows. In Section 2, we described a general framework of WGBoost in Algorithm 1. To use WGBoost, we need to specify two components:\n\n- (a) a loss functional $\\mathcal{F}(\\hat{\\pi}\\_{x_i} | \\pi\\_{x_i})$ between the WGBoost output $\\hat{\\pi}\\_{x_i}$ and the target distribution $\\pi\\_{x_i}$ at each $x_i$;\n- (b) how to estimate the Wasserstein gradient of $\\mathcal{F}$ at each $x_i$.\n\nIn Section 3.3, we described that the kernel-based functional gradient used in SVGD is one way to approximate the Wasserstein gradient of the KL divergence (see [49]). We then described that we use the kernel-based approximation to specify item (b) for our main applications. In Appendix E, we compared WGBoost under four different choices of the estimate of the Wasserstein gradient: functional gradient in SVGD, one in Stein Newton Method with full Hessian, one in Stein Newton Method with diagonal Hessian, and a stochastic alternative in Langevin diffusion. We applied WGBoost for probabilistic classification/regression. Hence our main experiment was the place to compare WGBoost with relevant UQ competitors.\n\n&nbsp;\n\n> ... the quality of approximation by the WGBoost is strongly affected by the choice of the kernel function. ...\n\nThank you for raising this point. The choice of the kernel is indeed important for most kernel-related methods. We observed that the Gaussian kernel with a scale parameter 0.1 works well through a simulation study. We will add a detailed discussion on the kernel choice and the sensitivity with simulation studies.\n\n&nbsp;\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. \n\nThank you for the opportunity for us to clarify. Our proposed method is a classification/regression algorithm. For experiments using UCI data, reporting the total log-likelihood score or RMSE would be fairly a standard practice. For other simpler datasets, we have visualisation of the output of WGBoost: Figure 2 with the toy target with $\\pi_{x_i} = \\mathcal{N}(\\sin(x_i), 1)$, Figure 1 for data in Section 4.1, and Figure 5 in Appendix for second data in Section 4.1. However, we agree that having more visualisation even for UCI datasets in Section 4.2-4.3 and adding a discussion on the approximation property would be helpful. \n\nFor visualisation of UCI datasets, we will pick a few examples of input $x_i$ and target density $\\pi_{x_i}$ from the dataset. We will then show the output of WGBoost at $x_i$ together with $\\pi_{x_i}$. The input $x_i$ is of arbitrary dimension but the target density $\\pi_{x_i}$ at each input $x_i$ is often 1 or 2 dimensional uni-modal density in our application. So the approximation at each $x_i$ should be less challenging. We will demonstrate this point through our additional visualisation.\n\n&nbsp;\n\n> In addition to the above point, there is no discussion about the computational cost of the proposed algorithm. ... / As for the convergence speed, it has been known that the convergence of the SVGD is slow ... I suspect that the proposed method suffers similar problem. However, no discussion is present about the convergence speed or no numerical comparison ...\n\nThank you for the opportunity for us to clarify. As recapped at the top, WGBoost is not a sampling method such as SVGD. WGBoost is an algorithm that depends on user-specification of how to estimate the Wasserstein gradient. In Appendix E, we have a simulation study to compare both the loss-decay speed and computation time of WGBoost under four different estimates of the Wasserstein gradient. We will make more clear this and further expand the discussion on the speed/computational order in the main text.\n\nSection 3.4 clarified that, for our application, we employ the functional gradient used in Stein Newton Method (with diagonal Hessian approximation) in WGBoost at the end rather than one used in SVGD. It is the second-order version of the functional gradient in SVGD and more efficient. Our simulation study shows that this choice has better loss-decay and computation time.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rR50eVY9ah", "reviewer": "Reviewer_DY4u", "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.\n\nSeveral 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.", "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?\n\nAre there other possible functionals which could be interesting besides the KL or other divregences between probability distributions?\n\nCould the method be used with input distributions, e.g. to do regression on probability distributions?\n\nTypos:\n- Line 145: \"gradint\"\n- Line 520: \"Wasserstien\"", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This paper is well written and proposes a new interesting boosting method to minimize loss over probability distributions.\n\n- The paper is well written\n- A new boosting algorithm guided by functionals on probability distributions\n- Application on real datasets outperforming baseline methods", "weaknesses": "The paper is good overall in my opinion, but still has some weaknesses.\n\n- Experiments focus on KL divergence functional.\n- No theoretical analysis", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your strong endorsement of our work. We are delighted that you have found our methodology well-written and intriguing.\n\n&nbsp;\n\n> 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?\n\nThank you for raising this point. So far we have compared them only through the simulation study to pick up the most reasonable algorithm (i.e. Wasserstein Hessian diagonal preconditioner) for the real-world application. We can add the further comparison of them based on the real-world application.\n\n&nbsp;\n\n> Are there other possible functionals which could be interesting besides the KL or other divregences between probability distributions?\n\nThank you for checking this point with us. Yes, we have a large room to investigate what would be other interesting functionals to use with WGBoost. For example, kernel Stein discrepancy would be one of the examples that we can perform the Wasserstein gradient flow. We will expand the discussion about the other potential functional choices. Although our scope in this work is the KL divergence, it is important feature work to develop a further special case of WGBoost with other functionals.\n\n&nbsp;\n\n> Could the method be used with input distributions, e.g. to do regression on probability distributions?\n\nThank you for your interesting question. Our current thought is that, because WGBoost can use any base learner, we can do that if we use such a base learner that can take a distribution as an input value. This sounds another interesting usage of WGBoost.\n\n&nbsp;\n\n> Typos: ...\n\nThank you. We have corrected the pointed typos.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to express our gratitudes to all the reviewers for their efforts in assessing our work. Our full rebuttal to each reviewer has been provided in each individual rebuttal area. This global rebuttal area is used for the following contents:
- (A) Brief Summary of Each Rebuttal
- (B) Concise Overview of Our Work

&nbsp;

## A. Brief Summary of Each Rebuttal

- **Reviewer jx13** & **Reviewer DY4u**: We appreciate your strong endorsement of our work. We believe that our work represents a promising advance that combines the gradient boosting framework with Wasserstein gradient systems for the first time.

- **Reviewer UM7J**: Thank you for your constructive feedback to our work. About our experiment, we will add (i) larger-scale datasets, (ii) other advanced approximate Bayesian methods, and (iii) other kernel-based/quantile-based nonparametric methods. About clarity of our algorithm (Algorithm 1), please find our full rebuttal and Section B below recapping our paper structure. In short,
  - Section 2 explains a general framework of WGBoost in Algorithm 1, where we need to specify *'how to estimate the Wasserstein gradient'* as a user-specified component;
  - Section 3 then explains that we plug the kernel-smoothing estimate as the Wasserstein gradient estimate into Algorithm 1 for our application.

- **Reviewer q9Jy**: Thank you for your assessment of our work. We are afraid that WGBoost seems misinterpreted as a sampling algorithm (such as SVGD) in raised concerns. Briefly, WGBoost is an algorithm to solve a new form of regression problem, such that, each input $x$ is arbitrary vector and each associated output is a probability distribution $\pi_x$ rather than some vector $y$. Please find our full rebuttal and Section B below recapping the high-level problem setting of WGBoost.

&nbsp;

## B. Concise Overview of Our Work

Our paper structure is mainly twofold. We proposed a general framework of WGBoost in Section 2. We then described how WGBoost can be used for probabilistic classification/regression in Section 3.

**B.1. Problem to Solve**

Firstly, the problem that the general framework of WGBoost can solve is a new form of regression problem below:
 
- We have a dataset $( x_i, \pi_{x_i} )\_{i=0}^{n}$ where input  $x_i$ is arbitrary vector and output $\pi_{x_i}$ is a probability distribution given at each $x_i$. Our goal is to have a learning algorithm that predicts the unseen distribution $\pi_x$ for any new input $x$.

This is a challenging problem because the target output $\pi_{x_i}$ at each input $x_i$ is an infinite-dimensional object (probability distribution) rather than some finite-dimensional vector $y_i$. WGBoost constructs a map from any input $x$ to a set of particles $ \hat{\pi}_x$ that approximates $\pi_x$:

- $x \to \text{WGBoost} \to ( \theta^1(x) , \dots, \theta^N(x) ) =: \hat{\pi}_x \approx \pi_x$.

Our novelty is in bringing the concept of Wasserstein dynamics into gradient boosting, which has never been explored.

**B.2. Gradient Boosting with Wasserstein Gradient**

Gradient boosting is a well-used ensemble method that trains multiple weak learners and combines them iteratively. In standard gradient boosting, each weak learner is trained to approximate the gradient $\nabla_{z_i} L(z_i | y_i)$ of a loss function $L(z_i | y_i)$ between the ensemble output $z_i$ and a target-output vector $y_i$ at each input $x_i$. In WGBoost, we rather use a loss functional $\mathcal{F}(\hat{\pi}\_{x_i} | \pi\_{x_i})$ between the WGBoost output and the target distribution at each input $x_i$. Then, each weak learner is trained to approximate the "Wasserstein gradient" of the loss functional at each input $x_i$. The key innovation is that the Wasserstein gradient evaluated here will be a finite-dimensional vector. Hence, each weak learner in WGBoost can be trained like usual regression algorithms, despite the original target $\pi\_{x_i}$ is an infinite-dimensional object. See also illustrative Figure 1-2 in the attached PDF.

In gradient boosting, users have to specify which loss function $L(z_i | y_i)$ to use. In WGBoost, users have to specify

- (a) which loss functional $\mathcal{F}(\hat{\pi}\_{x_i} | \pi\_{x_i})$ to use at each input $x_i$;
- (b) how to estimate/approximate the Wasserstein gradient of the chosen functional $\mathcal{F}$.

For Wasserstein gradient flows in general, it is common that the Wasserstein gradient is not analytically available (except some nice cases) and needs to be approximated (c.f. third paragraph of Section 2.1). Hence WGBoost needs to specification of item (b) too.

**B.3. Application to Classification/Regression**

Section 3 described application of WGBoost to probabilistic classification/regression, and how we specified item (a)-(b) for it. For illustration, consider classification where we have a dataset $( x_i, y_i )\_{i=1}^{N}$ of input vector $x$ and output label $y$. Assume a categorical distribution $\mathcal{C}(y \mid q)$ over the output label $y$ with the class probability vector $q$. With WGBoost, we can obtain a particle-based distributional estimate $\hat{\pi}\_{x}(q)$ of $q$ for any input $x$ (i.e. predictive uncertainty in $q$); see Figure 3 diagram in the attached PDF. To train WGBoost, we prepare the target distribution $\pi_{x_i}$ at each input $x_i$ as follows:

- Prepare a prior distribution $\nu_{x_i}(q)$ over $q$ at each input $x_i$ (c.f. Section 3.2);
- Get a posterior distribution $\pi_{x_i}(q) \propto \mathcal{C}(y_i \mid q) \times \nu_{x_i}(q)$ over $q$ from the likelihood $\mathcal{C}(y_i \mid q)$ and the prior $\nu_{x_i}(q)$ at each data point $(x_i, y_i)$.

Note that this posterior $\pi_{x_i}(q)$ is made at each data point $(x_i, y_i)$ pointwisely. This is different from a posterior of e.g. Bayesian neural network, which is defined over the network parameter $w$ using all data, $P(w \mid (x_i, y_i)\_{i=1}^{N})$. For regression, we can think of $\mathcal{N}(y \mid m, \sigma)$ instead of $\mathcal{C}(y \mid q)$.

Meta-review: Despite a score discrepancy, all four reviewers agreed that the proposed approach to predict probability distributions, based on Wasserstein Gradient Flow and Boosting, is original.

I have read the manuscript myself, and I assert that it is well-written. This "new perspective on ensemble learning algorithms" (as said by Reviewer jx13) will likely inspire future developments. The experiments are modest in scale but sufficient to show that the approach is sound and applicable to several tasks (probabilistic regression, classification, and out-of-distribution detection). Moreover, following the recommendations of reviewer UM7J, the authors have committed to expanding their experiments. 

I therefore enjoin the authors to carefully take into account the reviewers' comments while preparing the camera-ready version. I also suggest briefly discussing in which manner the approach might be reminiscent of the SVGD sampling method, which will help clarify the contributions, as the authors convincingly did in their answers to Reviewer q9Jy.
Decision: Accept (poster)