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

Authors: takuo  matsubara
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:41:34

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.

Main paper: 
 introduction: 
 Gradient boosting is a celebrated machine learning algorithm that has achieved considerable success with tabular data [1] . Gradient boosting has been extensively used for point forecasts and probabilistic classification, yet a relatively small number of studies have been concerned with the predictive uncertainty of gradient boosting. Predictive uncertainty of machine learning models plays a growing role in today's real-world production systems [2] . It is vital for safety-critical systems, such as medical diagnoses [3] and autonomous driving [4] , to assess the potential risk of their actions that partially or entirely rely on predictions from their models. Gradient boosting has already been applied in a diverse range of real-world applications, including click prediction [5] , ranking systems [6] , scientific discovery [7] , and data competition [8] . There is a pressing need for methodology to harness the power of gradient boosting to predictive uncertainty quantification.In classification and regression, we typically specify a noise distribution p(y | θ) of a response variable y and use a model to return a point estimate θ(x) of the response parameter for each input x. In recent years, the importance of capturing uncertainty in the model output θ(x) has increasingly been emphasised [2] . A variety of approaches have been proposed to obtain a distributional estimate p(θ | x) of the response parameter for each input x [e.g. 9, 10, 11] . For example, Bayesian neural networks (BNNs) quantify uncertainty in network weights and propagate it to the space of network outputs. Marginalising the predictive distribution p(y | θ) over the distributional estimate p(θ | x) has been demonstrated to confer enhanced predictive accuracy and robustness against adversarial attacks [11] . Furthermore, the dispersion of the distributional estimate has been used as a powerful indicator for out-of-distribution (OOD) detection [12] . In this context, a line of research based on the concept of evidential learning has recently gained significant attention [11, 13, 14, 15] . The idea can be broadly interpreted as making use of the 'individual-level' posterior p(θ | y i ) of the response parameter θ conditional on each individual datum y i , which arises from the response-distribution likelihood p(y i | θ) and a user-specified prior p(θ). If each individual-level posterior falls into a closed form characterised by some hyperparameter, neural networks can be trained by using the finite-dimensional hyperparameter as a target value for each input. Outstanding performance and computational efficiency of the existing approaches have been delivered in a wide spectrum of engineering and medical applications [16, 17, 18, 19] . However, the existing approaches are limited to neural networks and to the case where every individual-level posterior is in closed form so that the finite-dimensional hyperparameter can be predicted by proxy. In general, posterior distributions are known only up to their normalising constants and, therefore, require an approximation typically by particles [20] .Without closed-form expression, each individual-level posterior needs to be treated as an infinitedimensional output for each input. This challenge poses the following fundamental question: Consider a supervised learning setting whose outputs are probability distributions. Given a training set of input values and output distributions {x i , µ i } D i=1 , can we build a model that receives an input x and returns a nonparametric prediction of the output distribution? Motivated by this question, we propose a general framework of Wasserstein gradient boosting (WGBoost). WGBoost receives an input and returns a particle approximation of the output distribution. Figure 1 illustrates inputs and outputs of WGBoost. In this paper, we focus on application of WGBoost to evidential learning, where the individual-level posterior p(θ | y i ) of the response parameter θ is used as the output distribution µ i for each input x i in the training set. Figure 2 compares the pipeline of evidential learning based on WGBoost with that of Bayesian learning.Contributions Our contributions are summarised as follows:1. Methodology of WGBoost: Section 2 establishes the general framework of WGBoost. It is a novel family of gradient boosting that returns a set of particles that approximates an output distribution assigned at each input. In contrast to standard gradient boosting that fits a weak learner to the gradient of a loss function, WGBoost fits a weak learner to the estimated Wasserstein gradient of a loss functional over probability distributions. 2. Application to Evidential Learning: Section 3 establishes tree-based evidential learning based on WGBoost, with the loss functional specified by the Kullback-Leibler (KL) divergence. Following modern gradient-boosting libraries [21, 22] that uses second-order gradient boosting (c.f. Section 2.2), we implement a concrete second-order WGBoost algorithm built on an approximate Wasserstein gradient and Hessian of the KL divergence. 3. Experiment on Real-world Data: Section 4 demonstrates the performance of probabilistic regression, and classification with OOD detection, on real-world tabular datasets. To the author's knowledge, WGBoost is the first framework that enables evidential learning for (i) boosted tree models and (ii) cases without closed form of individual-level posteriors. Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes]Justification: The claim made in the abstract and introduction were elaborated in the subsequent section and empirically demonstrated through the experiments.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
wasserstein gradient boosting: 
 This section establishes the general formulation of WGBoost. Section 2.1 recaps the notion of Wasserstein gradient flows, a 'gradient' system of probability distributions that minimises an objective functional in the space of probability distributions. Section 2.2 recaps the notion of gradient boosting, a sequential ensemble method that fits a new weak learner to the 'gradient' of the remaining loss at each iteration. Section 2.3 combines the above two notions to establish WGBoost, a novel family of gradient boosting that enables to solve distribution-valued supervised learning.Notation and Setting Let X and Y denote the space of inputs and responses in classification and regression. Suppose Θ = R d . Let P 2 be the 2-Wasserstein space i.e. a set of all probability distributions on Θ with finite second moment equipped with the Wasserstein metric [23] . We identify a probability distribution in P 2 with its density whenever it exits. Denote by ⊙ and ⊘, respectively, elementwise multiplication and elementwise division of two vectors in R d . Let ∇ be the gradient operator. Let ∇ 2 d be a second-order gradient operator that takes the second derivative at each coordinate i.e.∇ 2 d f (θ) = [∂ 2 f (θ)/∂θ 2 1 , . . . , ∂ 2 f (θ)/∂θ 2 d ] T ∈ R d .In the Euclidean space, a gradient flow of a function f means a curve of points x t that solves a differential equation (d/dt)x t = -∇f (x t ) from some initial value x 0 . That is the continuous-time limit of gradient descent, which minimises the function f as t → ∞. A Wasserstein gradient flow means a curve of probability distributions µ t minimising a given functional F on the 2-Wasserstein space P 2 from some initial distribution µ 0 . The Wasserstein gradient flow µ t is characterised as a solution of a partial differential equation, known as the continuity equation:EQUATIONwhere ∇ W F(µ) : Θ → Θ denotes the Wasserstein gradient of F at µ [24, 25] . Appendix A recaps the derivation of the Wasserstein gradient, presenting the examples for several functionals.One of the elegant properties of the Wasserstein gradient flow is casting the infinite-dimensional optimisation of the functional F as a finite-dimensional particle update [23] . The continuity equation (1) can be reformulated as a dynamical system of a random variable θ t ∼ µ t , such thatEQUATIONin the sense that the law µ t of such a random variable θ t is a weak solution of the continuity equation. Consider the case where the initial measure µ 0 is set to the empirical distribution μ0 of N particles {θ n 0 } N n=1 . Discretising the continuous-time system (2) by the Euler method with a small step size ν > 0 yields an iterative update scheme of N particles {θ n m } N n=1 from step m = 0:EQUATIONwhere μm denotes the empirical distribution of the particles {θ n m } N n=1 at step m. In practice, it is common that the Wasserstein gradient of a chosen functional F is not well-defined for empirical distributions. In such case, the particle update scheme (3) is not directly applicable because it depends on the Wasserstein gradient ∇ W F(μ m ) at the empirical distribution μm . For example, the KL divergence F(µ) = KL(µ | π) with a reference distribution π leads to the Wasserstein gradient [∇ W F(µ)](θ) = -(∇ log π(θ) -∇ log µ(θ)) ill-defined if µ is an empirical distribution. Hence, the particle update scheme (3) is often performed with the estimated or approximated Wasserstein gradient well-defined for empirical distributions [e.g. 26, 27, 28, 29, 30] . A main application of WGBoost in Section 3 uses the 'smoothed' Wasserstein gradient of the KL divergence [26] .Gradient boosting [31] is a sequential ensemble method of M weak learners f 1 , . . . , f M . It iteratively constructs an ensemble F m of m weak learners f 1 , . . . , f m from step m = 0 to M . Given the current ensemble F m at step m, it trains a new weak learner f m+1 to construct the next ensemble byEQUATIONwhere ν is a shrinkage hyperparameter called a learning rate. The initial state of the ensemble F 0 (x) at step m = 0 is typically set to a constant that best fits the data. Any learning algorithm can be used as a weak learner in principle, although tree-based algorithms are most used [32] .The fundamental idea of gradient boosting is to train the new weak learner f m+1 to approximate the negative gradient of the remaining error of the current ensemble F m . Suppose that a loss function L measures the remaining error R i (F m (x i )) := L(F m (x i ), y i ) for each output vector y i ∈ R d . The new weak learner f m+1 is fitted to the set {x i , g i } D i=1 whose target variable g i is each specified byg i := -∇R i (F m (x i )) ∈ R d .The target g i is often called a pseudo residual. For each data input x i , the boosting scheme (4) updates the output of the current ensemble F m (x i ) in the steepest descent direction of the error R i (F m (x i )).Although [31] originally suggested an additional line search to determine a scaling constant of each weak learner, the line search has been reported to have a negligible influence on performance [33] .In modern gradient-boosting libraries, such as XGBoost [21] and LightGBM [22] , the standard practice is to use the diagonal (coordinatewise) Newton direction of the remaining error R i (F m (x i )) in lieu of the negative gradient g i . The new base leaner f m+1 is instead fitted to the set {x i , g i ⊘h i } n i=1 , where the negative gradient g i is divided elementwise by the Hessian diagonal h i given byh i := ∇ 2 d R i (F m (x i )) ∈ R d .The target variable g i ⊘ h i is the diagonal Newton direction that minimises the second-order Taylor approximation of the remaining error R i (F m (x i )) for each coordinate independently. Combining the second-order gradient boosting framework with tree-based weak learners has demonstrated exceptional scalability and performance [34, 35] . Although it is possible to use the 'full' Newton direction as the target variable of each weak learner, the impracticality of the full Newton direction has been pointed out [e.g. 36, 37] . In addition, the coordinatewise computability of the diagonal Newton direction is suitable for popular gradient-boosting tree algorithms [36] .Now we consider the setting of distribution-valued supervised learning, where we are given a training set of input vectors and output distributions {x i , µ i } D i=1 ⊂ X × P 2 . Our goal is to construct a model that receives an input and returns a set of N particles whose empirical distribution approximates the output distribution. We specify a loss functional D(• | •) between two probability distributions-such as the KL divergence-to measure the remaining error F i (•) = D(• | µ i ) for each i-th training output distribution µ i . Our idea is to combine gradient boosting with Wasserstein gradient, where we iteratively construct a set of N boosting ensembles F 1 m , . . . , F N m from step m = 0 to M . Here, the output F n m (x) of each n-th boosting ensemble represents the n-th output particle for an input x. Given the current set of N ensembles F 1 m , . . . , F N m at step m, WGBoost trains a set of N new weak learners f 1 m+1 , . . . , f N m+1 and computes the next set of N ensembles by    F 1 m+1 (x) . . .EQUATION)where ν is a learning rate. Similarly to standard gradient boosting, we specify the initial state of N ensembles F 1 0 , . . . , F N 0 at step m = 0 by a set of constants. Throughout, denote by μm,i the empirical distribution of the N output particles F 1 m (x i ), . . . , F N m (x i ) for each i-th training input x i . As discussed in Section 2.1, the Wasserstein gradient often needs to be estimated for empirical distributions. For better presentation, let G i (µ) denote an estimate of the Wasserstein gradient ∇ W F i (µ) of the i-th remaining error F i (µ), which is well-defined for any distribution µ. If the Wasserstein gradient ∇ W F i (µ) is originally well-defined for any distribution µ, it is a trivial choice of the estimate, i.e., G i (µ) = ∇ W F i (µ). Otherwise, any suitable estimate can be used as G i (µ). The foundamental idea of WGBoost is to train the n-th new learner f n m+1 to approximate the estimated Wasserstein gradient -G i (μ m,i ) evaluated at the n-th boosting output F n m (x i ) for each x i , so that,   f 1 m+1 (x i ) . . . f N m+1 (x i )    ≈    -[G i (μ m,i )] F 1 m (x i ) . . . -[G i (μ m,i )] F N m (x i )    .For each data input x i , the boosting scheme (5) approximates the particle update scheme (3) for the output particles F 1 m (x i ), . . . , F N m (x i ) under the estimated Wasserstein gradient. The output particles are updated in the direction to decrease the remianing errorF i (μ m,i ) = D(μ m,i | µ i ) at each step m.Algorithm 1 summarises the general procedure of WGBoost. See Figure 1 for illustration of WGBoost. In Section 3, we choose the KL divergence as a loss functional D and use a kernel smoothing estimate of the Wasserstein gradient. See Appendix A for the Wasserstein gradient of other divergences. Remark 1 (Stochastic WGBoost). Stochastic gradient boosting [38] uses only a randomly sampled subset of data to fit a new weak learner at each step m to reduce the computational cost. The same subsampling approach can be applied for WGBoost whenever the dataset is large. Remark 2 (Second-Order WGBoost). If any estimate of the Wasserstein 'Hessian' of the remaining error F i is available, the Newton direction of F i may also be computable [e.g. 39, 40] . Implemention of a second-order WGBoost algorithm is immediate by plugging such a Newton direction into G i (µ) in Algorithm 1. Our default WGBoost algorithm for tree-based evidential learning is built on a diagonal approximate Newton direction of the KL divergence, aligning with the standard practice in modern gradient-boosting libraries to use the diagonal Newton direction.Input: training set {x i , µ i } D i=1 of input x i ∈ X and output distribution µ i ∈ P 2 Parameter :loss D, estimate G i (µ) of the Wasserstein gradient ∇ W D(µ | µ i ), particle number N , iteration M , learning rate ν, weak learner f , initial constants (ϑ 1 0 , . . . , ϑ N 0 ) Output: set of N boosting ensembles (F 1 M , . . . , F N M ) at final step M (F 1 0 (•), . . . , F N 0 (•)) ← (ϑ 1 0 , . . . , ϑ N 0 ) ▷ set initial state of N boosting ensembles for m ← 0, . . . , M -1 do for i ← 1, . . . , D do μm,i ← empirical distribution of output values (F 1 m (x i ), . . . , F N m (x i )) for input x i for n ← 1, . . . , N do g n i ← -[G i (μ m,i )] (F n m (x i )) ▷ compute target value of n-th new weak learner end end for n ← 1, . . . , N do f n m+1 ← fit {x i , g n i } D i=1▷ fit n-th new weak learnerF n m+1 (•) ← F n m (•) + νf n m+1 (•)▷ set next state of n-th boosting ensemble end endQuestion: Does the paper discuss the limitations of the work performed by the authors?Justification: A concise discussion on the important remaining questions and the potential limitation were provided in the last section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
application to evidential learning: 
 This section provides our default setting to implement a concrete WGBoost algorithm for evidential learning, which enables classification and regression with predictive uncertainty. The individual-level posterior p(θ | y i ) of a response distribution p(y | θ) is used as the output distribution µ i of the training set {x i , µ i } D i=1 . Section 3.1 recaps derivation of the individual-level posterior p(θ | y i ), followed by Section 3.2 discussing the default choice of the prior. We choose the KL divergence as a loss functional of WGBoost. Section 3.3 recaps a widely-used estimate of the Wasserstein gradient of the KL divergence based on kernel smoothing [26] . A further advantage of the kernel smoothing estimate is that the approximate Wasserstein Hessian is available, with which Section 3.4 establishes a second-order WGBoost algorithm similarly to modern gradient-boosting libraries.Suppose that a response distribution p(y | θ) of a response variable y is specified, as is typically done for probabilistic prediction. Suppose also that a prior distribution p i (θ) of the response parameter θ is specified for each individual data input x i . For each individual data pair (x i , y i ), the responsedistribution likelihood p(y i | θ) and the prior p i (θ) determine the individual-level posteriorp(θ | y i ) ∝ p(y i | θ)p i (θ)by Bayes' theorem. This individual-level posterior is set to the output distribution µ i of the training set {x i , µ i } D i=1 of WGBoost. The framework of WGBoost then constructs a model that returns a particle approximation of the output distributionµ i (•) = p(• | y i ) for each data input x i .For a new input x, the constructed WGBoost model provides a set of particles (θ 1 (x), . . . , θ N (x)) as a distributional prediction p(θ | x) of the response parameter θ. We can define a predictive distribution p(y | x) of the response y for the new input x via marginalisation of the output particles:EQUATIONWe can also define a point prediction ŷ for the new input x via the individual-level Bayes action ŷ = argmin y∈Y Θ U (y, θ)p(θ | x)dθ, which minimises the average of some error U : Y × Θ → R.For example, the Bayes action ŷ is simply the mean of the output particles if U (y, θ) = (y -θ) p(y | θ = f (x, w))p(w | {x i , y i } D i=1 )dw.In contrast, WGBoost returns a distributional prediction p(θ | x) of the response parameter θ, circumventing the marginalisation over the model parameter w that can be ultra-high dimensional.The prior distribution p i (θ) of the response parameter θ is specified at each individual data input x i . The approach to eliciting the prior may differ, depending on whether past data are available. If past data are available, they can be utilised to elicit a reasonable prior for unobserved data. If no past data are available, we recommend the use of a noninformative prior that have been developed as a sensible choice of prior in the absence of past data; see [e.g. 41] for the introduction. To avoid numerical errors, if a noninformative prior is improper (i.e. nonintegrable), we recommend the use of a proper probability distribution that approximates the noninformative prior sufficiently well. Example 1 (Normal Location-Scale Response). Consider a scalar-valued response variable y ∈ R for regression. A normal location-scale response distribution N (y | m, σ) has the mean and scale parameters m ∈ R and σ ∈ (0, ∞). A typical noninformative prior of m and σ are given by, respectively, 1 and 1/σ which are improper. At every data point (x i , y i ), we use a normal prior N (m | 0, σ 0 ) over m and an inverse gamma prior IG(σ | α 0 , β 0 ) over σ, with the hyperparameters σ 0 = 10 and α 0 = β 0 = 0.01, which approximate the non-informative priors. Example 2 (Categorical Response). Consider a label response variable y ∈ {1, . . . , k} for k-class classification. A categorical response distribution C(y | q) has the class probability parameter q = (q 1 , . . . , q k ) in the k-dimensional simplex ∆ k . If k = 2, it corresponds to the Bernoulli distribution. A typical noninformative prior of q is given by 1/(q 1 × • • • × q k ) which are improper. At every data point (x i , y i ), we use the logistic normal prior-a multivariate generalisation of the logit normal distribution [42] -over q with the mean 0 and identity covariance matrix scaled by 10. Remark 4 (Reparametrisation and Standardisation). Section 2 supposed Θ = R d for some dimension d without no loss of generality. Any parameter that lies in a subset of the Euclidean space (e.g. σ in Example 1) can be reparametrised as one in the Euclidean space (e.g. log σ). Appendix D details the reparametrisation used for Examples 1 and 2. In addition, if one's dataset has scalar outputs of a low or high order of magnitude, we recommend standardising the outputs.We consider the KL divergence KL(µ | µ i ) as a loss functional of WGBoost. One challenge of the KL divergence is that the resulting Wasserstein gradientG KL i (µ) (θ) := -(∇ log µ i (θ) -∇ log µ(θ))is not well-defined when µ is an empirical distribution. A particularly successful solution-which originates in [43] and has been applied in wide contexts [26, 44, 45] -is to smooth the Wasserstein gradient through a kernel integral operator [46] . By integration-bypart (see [e.g. 43]), the smoothed Wasserstein gradient, denoted G * i (µ), falls into the following form that is well-defined for any distribution µ:Θ [G KL i (µ)](θ * )k(θ, θ * )dµ(θ * )EQUATIONwhere ∇k(θ * , θ) denotes the gradient of k with respect to the first argument θ * . An approximate Wasserstein gradient flow based on the smoothed Wasserstein gradient G * i (µ) is called the Stein variational gradient descent [43] or kernelised Wasserstein gradient flow [47] . In most cases, the kernel k is set to the Gaussian kernel k(θ, θ * ) = exp(-∥θ -θ * ∥ 2 /h) with the scale h > 0. Appendix B discusses a choice of kernel. This work uses the Gaussian kernel with h = 0.1 throughout.Another common approach to approximating the Wasserstein gradient flow of the KL divergence is the Langevin diffusion approach [48] . The discretised algorithm, called the unadjusted Langevin algorithm [49] , is a stochastic particle update scheme that adds a Gaussian noise at every iteration. However, several known challenges, such as asymptotic bias and slow convergence, often necessitate an ad-hoc adjustment of the algorithm [48] . Appendix B discusses a variant of WGBoost built on the Langevin algorithm, although it is not considered the default implementation.We use a diagonal (coordinatewise) approximate Wasserstein Newton direction of the KL divergence, following the standard practice in modern gradient-boosting libraries [21, 22] to use the diagonal Newton direction of a loss. Similarly to smoothed Wasserstein gradient G * i (µ), the approximate Wasserstein Hessian of the KL divergence KL(µ | µ i ) can be obtained through the kernel smoothing. The diagonal of the approximate Wasserstein Hessian, denoted H * i (µ), is defined byEQUATIONThe diagonal approximate Wasserstein Newton direction of the KL divergence is then defined by-[G * i (µ)] (•) ⊘ [H * i (µ)] (•).Appendix C provides the derivation based on [39] who derived the Newton direction of the KL divergence in the context of nonparametric variational inference.The second-order WGBoost algorithm is established by plugging it into G i (µ) in Algorithm 1, that is,EQUATIONAlgorithm 1 under the choice ( 9) is considered our default WGBoost algorithm for evidential learning.We refer this algorithm to as the Wasserstein-boosted evidential learning (WEvidential). The explicit pseudocode is provided in Algorithm 2 for full clarity. Output: set of N boosting ensembles(F 1 M , . . . , F N M ) at final step M (F 1 0 (•), . . . , F N 0 (•)) ← (ϑ 1 0 , . . . , ϑ N 0 ) ▷ set initial state of N boostings for m ← 0, . . . , M -1 do for i ← 1, . . . , D do (θ 1 i , . . . , θ N i ) ← (F 1 m (x i ), . . . , F N m (x i )) ▷ get output particles for i-th data input for n ← 1, . . . , N do g n i ← 1 N N k=1 ∇ log p(θ k i | y i )k(θ k i , θ n i ) + ∇k(θ k i , θ n i ) h n i ← 1 N N k=1 -∇ 2 d log p(θ k i | y i )k(θ k i , θ n i ) 2 + ∇k(θ k i , θ n i ) ⊙ ∇k(θ k i , θ n i ) end end for n ← 1, . . . , N do f n m+1 ← fit {x i , g n i ⊘ h n i } D i=1▷ fit n-th new weak learnerF n m+1 (•) ← F n m (•) + νf n m+1 (•)▷ set next state of n-th boosting end end Remark 5 (Computation). The diagonal Newton direction (9) has the computational complexity O(N ×d) same as that of the smoothed Wasserstein gradient. Hence, there is essentially no reason not to use the diagonal Newton direction (9) instead of the smoothed Wasserstein gradient. Although it is possible to use the full Newton direction with no diagonal approximation, the computation requires the inverse and product of (N × d) × (N × d) matrices that result in the complexity up to O(N 3 × d 3 ). Appendix D presents a simulation study to compare computational time and convergence speed of four WGBoost algorithms built on different estimates of the Wasserstein gradient.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• 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. Justification: This work proposed a generic algorithm that is not tied to specific applications. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiment on real-world tabular data: 
 We empirically demonstrate the performance of the WGBoost algorithm through three experiments using real-world tabular data. The first application illustrates the output of WGBoost through a simple conditional density estimation. The second application benchmarks the probabilistic regression performance. The third application demonstrates the classification and OOD detection performance. The source code is available in https://github.com/takuomatsubara/WGBoost.Common Hyperparameters Throughout, we set the number of output particles N to 10 and set each weak learner f to the decision tree regressor [50] with maximum depth 1 for Section 4.1 and 3 for the rest. We set the learning rate ν to 0.1 for regression and 0.4 for classification, unless otherwise stated. Appendix E contains further details, including a choice of the initial constant {ϑ n 0 } N n=1 .We illustrate the output of WEvidential by estimating a conditional density p(y | x) from onedimensional scalar inputs and outputs {x i , y i } D i=1 . The normal output distribution N (y | m, σ) and the prior p i (m, σ) in Example 1 were used to define the individual-level posterior p(m, σ | y i ), in which case the output of the WGBoost algorithm is a set of 10 particles {(m n (x), σ n (x))} 10 n=1 of the mean and scale parameters for each input x. We chose the number of weak learners M , drawing on an early-stopping approach used in [32] , where we held out 20% of the training set as a validation set and chose the number 1 ≤ M ≤ 4000 achieving the least validation error. Once the number M was chosen, WEvidential was trained again using all the entire training set.The conditional density is estimated using the predictive distribution (6) by WEvidential. We used two real-world datasets, bone mineral density [51] and old faithful geyser [52] . Figure 3 depicts the result for the former dataset, demonstrating that the WGBoost algorithm captures the heterogeneity of the conditional density on each input well. The result for the latter dataset is contained in Appendix E.1. We examine the regression performance of WEvidential using a standard benchmark protocol that originated in [53] and has been used in a number of subsequent works [10, 9, 32] . The benchmark protocol uses real-world tabular datasets from the UCI machine learning repository [54] , each with one-dimensional scalar responses. As in Section 4.1, the normal response distribution N (y | m, σ) and the prior p i (m, σ) in Example 1 were used to define the individual-level posterior p(m, σ | y i ).We followed the data splitting protocol in [53] and randomly held out 10% of each dataset as a test set. The negative log likelihood (NLL) is measured by using the predictive distribution (6) . The root mean squared error (RMSE) is measured by using the point prediction by the mean value of the predictive distribution. We chose the number of weak learners M by the same approach as in Section 4.1. We repeated this procedure 20 times for each dataset, except the protein and year msd datasets for which we repeated five times and once. For the year msd dataset only, we subsampled 10% of data to fit each weak learner and used the learning rate 0.01 due to the large dataset size.Table 1 compares the performance of WEvidential with five other methods: Monte Carlo Dropout (MCDropout) [9] , Deep Ensemble (DEnsemble) [10] , Concrete Dropout (CDropout) [55] , Natural Gradient Boosting (NGBoost) [32] , and Deep Evidential Regression (DEvidential) [13] . Appendix E provides further details on the experiment and a limited yet additional comparison. The WGBoost algorithm achieves the best score or a score sufficiently close to the best score most often.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, DEnsembles, CDropout, NGBoost, and DEvidential were reported in [9] , [10] , [55] , [32] and [13] respectively. We examine the classification and anomaly OOD detection performance of WEvidential on two realworld tabular datasets, segment and sensorless, following the protocol used in [14] . The categorical response distribution C(y | q) and the prior p i (q) in Example 2 were used to define the individuallevel posterior p(q | y i ), in which case the output of the WGBoost algorithm is a set of 10 particles {q n } 10 n=1 of the class probability parameter q in the simplex ∆ k for each input x. We set the number of weak learners M to 4000 without early stopping to reduce the computational cost.The segment and sensorless datasets have 7 and 11 classes in total. For the segment dataset, the data subset that belongs to the last class was kept as the OOD samples. For the sensorless dataset, the data subset that belongs to the last two classes was kept as the OOD samples. For each dataset, 20% of the non-OOD samples is held out as a test set to measure the classification accuracy. There exist several ways of defining a OOD score for each input [56] . For the WGBoost algorithm, the inverse of the maximum norm of the output-particle variance was used as the OOD score. We measured the OOD detection performance by the area under the precision recall curve (PR-AUC), viewing non-OOD test data as the positive class and OOD data as the negative class. We repeated this procedure five times.Table 2 compares the performance of WEvidential with four other methods: MCDropout, DEnsemble, and Distributional Distillation (DDistillation) [57] , and Posterior Network (PNetwork) [14] . Appendix E provides further details on the experiment. Figure 4 exemplifies how the dispersion of the output particles differ between OOD and non-OOD inputs. WEvidential demonstrates a high classification and OOD detection accuracy simultaneously. Although PNetwork has the best OOD detection performance for the sensorless dataset, the performance of the WGBoost algorithm also exceeds 80%, which is distinct from MCDropout, DEnsemble, and DDistillation.  
discussion: 
 This work established the general framework of WGBoost and developed the concrete algorithm WEvidential for evidential learning. The established framework of WGBoost offers exciting avenues for future research. Important directions for future study include (i) exploring alternative loss functionals to the KL divergence, (ii) investigating the convergence properties, and (iii) evaluating robustness of obtained predictive uncertainty in comparison to other methods. A particular limitation of WGBoost may arise when data are not tabular, as is the case of standard gradient boosting. These questions require careful examination and are critical for future study.This appendix contains the technical and experiment details referred to in the main text. Appendix A recaps derivation of the Wasserstein gradient. Appendix B discusses the variant of WGBoost built on the unadjusted Langevin algorithm. Appendix C derives the diagonal approximate Wasserstein Newton direction used for WEvidential. Appendix D provides simulation studies on kernel choice of WEvidential and comparison of WGBoost algorithms built on different estimate of the Wasserstein gradient. Appendix E contains the additional details of the experiment presented in the main text.This section recaps derivation of the Wasserstein gradient of a functional F, with examples of common divergences. The Wasserstein gradient depends on a function on Θ called the first variation [24] . The first variation δF(µ)/δµ of the functional F at µ is a function on Θ that satisfieslim ϵ→0 + F(µ + ϵν) -F(µ) ϵ = Θ δF(µ) δµ (θ)ν(θ)dθfor all signed measure ν s.t. µ + ϵν ∈ P 2 for all ϵ sufficiently small. The Wasserstein gradient ∇ W F(µ) of the functional F at µ is derived as the gradient of the first variation (see [e.g. 24]):[∇ W F(µ)](θ) := ∇ δF(µ) δµ (θ).It is common to suppose that the functional F consists of three energies, which are determined by functions U : R → R, V : Θ → R, and W : Θ → R respectively, such thatF(µ) = Θ U (µ(θ))dθ internal energy + Θ V (θ)µ(θ)dθ potential energy + 1 2 Θ×Θ W (θ -θ ′ )µ(θ)dθµ(θ ′ )dθ ′ interaction energy.For a functional F that falls into the above form, the Wasserstein gradient is derived as[∇ W F(µ)] (θ) = ∇U ′ (µ(θ)) + ∇V (θ) + Θ ∇W (θ -θ ′ )µ(θ ′ )dθ ′where U ′ is the derivative of U : R → R [23] . The KL divergence F(µ) = KL(µ | π) of a distribution π falls into the form with U (x) = x log x, V (θ) = -log π(θ), and W (θ) = 0, whereKL(µ | π) = Θ log µ(θ)µ(θ)dθ + Θ -log π(θ)µ(θ)dθ.Table 3 presents examples of Wasserstein gradients of common divergencesF(µ) = D(µ | π).In the context of Bayesian inference, the KL divergence is particularly useful among many divergences. The Wasserstein gradient of the KL divergence requires no normalising constant of a posterior distribution π. This is because the Wasserstien gradient depends only on the log-gradient of the posterior ∇ log π(θ) = ∇π(θ)/π(θ) of the target π, in which case the normalising constant of the target π is cancelled out by fraction. Hence, any posterior known only up to the normalising constant can be used as the target distribution π in the Wasserstein gradient of the KL divergence.Table 3 : Wasserstein gradients of four divergences: the KL divergence [58] , the chi-squared divergence [47] , the alpha divergence [59] , and the maximum mean discrepancy [60] .DivergenceF(µ) = D(µ | π) Wasserstein gradient [∇ W F(µ)] (θ) KL(µ | π) -(∇ log π(θ) -∇ log µ(θ)) Chi 2 (µ | π) 2∇(µ(θ)/π(θ)) Alpha(µ | π) (µ(θ)/π(θ)) α-1 ∇(µ(θ)/π(θ)) MMD(µ | π) Θ ∇k(θ, θ ′ )µ(θ)dθ -Θ ∇k(θ, θ ′ )π(θ)dθIf a chosen functional F on P 2 is the KL divergence F(µ) = KL(µ | π) of a target distribution π, the continuity equation ( 1) admits an equivalent representation as the Fokker-Planck equation [61] :EQUATIONwhere ∆ denotes the Laplacian operator. Recall that the original continuity equation ( 1) can be reformulated as the deterministic differential equation ( 2) of a random variable θ t ∼ µ t . In contrast, the Fokker-Planck equation ( 10) can be reformulated as a stochastic differential equation of a random variable θ t ∼ µ t , known as the overdamped Langevin dynamics [62] :EQUATIONwhere B t denotes a standard Brownian motion. Note that the deterministic system (2) in the case of the KL divergence and the above stochastic system (11) are equivalent at population level, in a sense that the law of the random variable θ t in both the systems solves the two equivalent equations.At the algorithmic level, however, discretisation of each system leads to different particle update schemes. Set the initial distribution µ 0 in (11) to the empirical distribution μ0 of N initial particles {θ n 0 } N n=1 . Discretising the stochastic system (11) by the Euler-Maruyama method with a step size ν > 0 yields a stochastic update scheme of particles {θ n m } N n=1 from step m = 0:   θ 1 m+1 . . . θ N m+1    =    θ 1 m . . . θ N m    + ν    ∇ log π(θ 1 m ) + 2/ν ξ 1 . . . ∇ log π(θ N m ) + 2/ν ξ N    ,where each ξ n denotes a realisation from a standard normal distribution on R d . The above updating scheme of each n-th particle is known as the unadjusted Langevin algorithm [49] . We can define a variant of WGBoost by replacing the term G i (µ) in Algorithm 1 with ∇ log µ i (•) + 2/ν ξ i where µ i is an output distribution at each x i and ξ i is a realisation from a standard normal distribution. The procedure is summarised in Algorithm 3, which we call Langevin gradient boosting (LGBoost).Input: training set {x i , µ i } D i=1 of input x i ∈ X and output distribution µ i ∈ P 2 Parameter :particle number N , iteration M , rate ν, weak learner f , initial constants(ϑ 1 0 , . . . , ϑ N 0 ) Output: set of N boosting ensembles (F 1 M , . . . , F N M ) at final step M (F 1 0 (•), . . . , F N 0 (•)) ← (ϑ 1 0 , . . . , ϑ N 0 ) for m ← 0, . . . , M -1 do for n ← 1, . . . , N do for i ← 1, . . . , D do g n i ← ∇ log µ i (F n m (x i )) + 2/ν ξ n i where ξ n i ∼ N (0, I d ) end f n m+1 ← fit {x i , g n i } D i=1 F n m+1 (•) ← F n m (•) + νf n m+1 (•) end endThis section derives the diagonal approximate Wasserstein Newton direction based on the kernel smoothing. The approximate Wasserstein Newton direction of the KL divergence was derived in [39] under a different terminology-simply, the Newton direction-from a viewpoint of nonparametric variational inference. We place their result in the context of approximate Wasserstein gradient flows. Appendix C.1 shows the derivation of the smoothed Wasserstein gradient and Hessian. Appendix C.2 defines the Newton direction built upon the smoothed Wasserstein gradient and Hessian, following the derivation in [39] . Appendix C.3 derives the diagonal approximation of the Newton direction.Consider the one-dimensional case Θ = R for simplicity. For a map T : R → R and a distribution µ ∈ P 2 , let µ t be the pushforward of µ under the transform θ → θ + tT (θ) defined with a timevariable t ∈ R. This means that µ t is a distribution obtained by change-of-variable applied for µ. The Wasserstein gradient of a functional F(µ) can be associated with the time derivative (d/dt)F(µ t ) [23] . In what follows, we focus on the KL divergence F(µ) = KL(µ | π) as a loss functional. Under a condition T ∈ L 2 (µ), the time derivative at t = 0 satisfies the following equalityEQUATIONwhere G KL (µ) denotes the Wasserstein gradient of F(µ) = KL(µ | π) with the target distribution π made implicit. It gives an interpretation of the Wasserstein gradient as the steepest-descent direction because the decay of the KL divergence at t = 0 is maximised when T = -G KL (µ).The 'smoothed' Wasserstein gradient can be derived by restricting the transform map T to a more regulated Hilbert space than L 2 (µ EQUATIONwhere[G * (µ)](•) := Θ [G KL (µ)](θ)k(θ,•)dµ(θ) corresponds to the smoothed Wasserstein gradient used in the main text. The decay of the KL divergence at t = 0 is maximised by T = -G * (µ).Similarly, the Wasserstein Hessian of the functional F(µ) can be associated with the second time derivative (d 2 /dt 2 )F(µ t ) [23] . As discussed in [e.g. 46], the Wasserstein Hessian of the KL divergence, denoted Hess(µ), is an operator over functions T ∈ L 2 (µ) that satisfiesEQUATIONSee [46] for the explicit form of the Wasserstein Hessian. In the same manner as the smoothed Wasserstein gradient, applying the reproducing property in (14) under the condition T ∈ H and exchanging the integral order, the second time derivative satisfies an alternative equality as follows:EQUATIONwhere[Hess * (µ)](⋆ 1 , ⋆ 2 ) := ⟨k(⋆ 1 , •), Hess(µ)k(⋆ 2 , •)⟩ L 2 (µ)is the smoothed Wasserstein Hessian and the symbols ⋆ 1 and ⋆ 2 denote the variables to which each of the two inner products is taken.In the multidimensional case Θ = R d , the transport map T is a vector-valued function T : R d → R d , where a similar derivation can be repeated by replacing L 2 (µ) and H with the product space of d independent copies of L 2 (µ) and H. It follows from Proposition 1 and Theorem 1 in [39] -which derives the explicit form of ( 13) and ( 15) under their terminology, first and second variations-that the explicit form of the smoothed Wasserstein gradient and Hessian is given by[G * (µ)] (•) = E θ∼µ -∇ log π(θ)k(θ, •) -∇k(θ, •) ∈ R d , [Hess * (µ)] (⋆ 1 , ⋆ 2 ) = E θ∼µ -∇ 2 log π(θ)k(θ, ⋆ 1 )k(θ, ⋆ 2 ) + ∇k(θ, ⋆ 1 ) ⊗ ∇k(θ, ⋆ 2 ) ∈ R d×dwhere ∇ 2 denotes an operator to take the Jacobian of the gradient-i.e., ∇ 2 f (θ) is the Hessian matrix of f at θ-and ⊗ denotes the outer product of two vectors. Note that both the smoothed Wasserstein gradient and Hessian are well-defined for any distribution µ including empirical distributions.In the Euclidean space, the Newton direction of an objective function is a direction s.t. the secondorder Taylor approximation of the function is minimised. Similarly, [39] characterised the Newton direction T * : R d → R d of the KL divergence KL(µ | π) as a solution of the following equation[Hess * (µ)](⋆ 1 , ⋆ 2 ), T * (⋆ 2 ) H + [G * (µ)](⋆ 1 ), V (⋆ 1 ) H = 0 for all V ∈ H.Here Θ = R d and H is the product space of d independent copies of the RKHS of a kernel k. To obtain a closed-form solution, [39] supposed that the Newton direction T * can be expressed in a formT * (•) = n i=1 W n k(•, θ n )dependent on a set of each particle θ n ∈ Θ and associated vector-valued coefficient W n ∈ R d . Once the set of the particles is given, the set of the associated vector-valued coefficients is determined by solving the following simultaneous linear equationEQUATION)These equations ( 16) can be rewritten in a matrix form [64] . Let K := N × d. Define a block matrix H ∈ R K×K and a block vector G ∈ R K by the following partitioningH =    H 11 • • • H 1N . . . . . . . . . H N 1 • • • H N N    and G =    G 1 . . . G N   with each block specified asH ij := [Hess * (µ)](θ i , θ j ) ∈ R d×d and G i := [G * (µ)](θ i ) ∈ R d .Define a block matrix K ∈ R K×K and a block vector W ∈ R K by the following partitioningK :=    K 11 • • • K 1N . . . . . . . . . K N 1 • • • K N N    and W :=    W 1 . . . W N   with each block of K specified as K ij := I d × k(θ i , θ j ) ∈ R d×d , where I d denotes the d × d identity matrix. Notice that W is a block vector that aligns the vector-valued coefficients {W n } N n=1 . Using these notations, the optimal coefficients that solve ( 16) is simply written as W = -H -1 G [64] .Given the optimal coefficients W = -H -1 G, the Newton direction T * (θ n ) evaluated at the given particle θ n for each n = 1, . . . , N can be written in the following block vector form   T * (θ 1 ). . .EQUATIONTo distinguish from the standard Newton direction in the Euclidean space, we call (17) the approximate Wasserstein Newton direction. The approximate Wasserstein Newton direction yields a second-order particle update scheme. Suppose we have particles {θ n m } N n=1 to be updated at each step m. At each step m, define the above matrices H and G with the empirical distribution µ = πm of the particles {θ n m } N n=1 . Replacing the Wasserstein gradient in the particle update scheme (3) by the approximate Wasserstein Newton direction (17) provides the second-order update scheme in [39] .We derive the diagonal approximation of the approximate Wasserstein Newton direction, which we used for our second-order WGBoost algorithm. A few approximations of the approximate Wasserstein Newton direction were discussed in [39] for better performance of their particle algorithm. We derive the diagonal approximation so that no matrix product and inversion will be involved. Specifically, we replace the matrices K and H in ( 17) by the diagonal approximations K and Ĥ, that is,K =    I d • • • 0 . . . . . . . . . 0 • • • I d    and Ĥ =    h 11 • • • 0 . . . . . . . . . 0 • • • h N N    ,where K nn = I d × k(θ n , θ n ) = I d for the Gaussian kernel k used in this work, and the matrix h nn ∈ R d×d denotes the diagonal approximation of the diagonal block H nn of H.Recall that H nn = [Hess * (µ)](θ n , θ n ). Denote by Diag(A) the diagonal of a square matrix A.The diagonal approximation h nn is a diagonal matrix whose diagonal is Diag(H nn ). We plug the diagonal approximations K and Ĥ in (17) . It follows from inverse and multiplication properties of diagonal matrices that the approximate Wasserstein Newton direction turns into a formEQUATION)At an arbitrary particle location θ, denote by [H * (µ)](θ) the diagonal of the smoothed Wasserstein Hessian [Hess * (µ)](θ, θ). It is straightforward to see that the diagonal can be written as[H * (µ)] (•) = E θ∼µ -∇ 2 d log π(θ)k(θ, •) 2 + ∇k(θ, •) ⊙ ∇k(θ, •) .Notice that Diag (H nn ) = [H * (µ)] (θ n ) by definition. It therefore follows from the formula (18) withG n = [G * (µ)](θ n ) and Diag (H nn ) = [H * (µ)] (θ n) that the diagonal approximate Wasserstein Newton direction at an arbitrary particle location θ can be independently computed by-[G * (µ)](θ) ⊘ [H * (µ)] (θ).We used this direction in Section 3. In the main text, the diagonal approximate Wasserstein Newton direction is defined for each loss functionalF i (•) = D(• | µ i ), with π = µ i , using the smoothed Wasserstein gradient G * i (µ)and the diagonal of the smoothed Wasserstein Hessian H * i (µ) for each i-th output distribution µ i .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.We used the kernel smoothing estimate of the Wasserstein gradient of the KL divergence in order to built WEvidential. The smoothed Wasserstein gradient originates in an approximate Wasserstein gradient flow called Stein variational gradient descent (SVGD) [43] . It is fairly common in SVGD in practice to use the Gaussian kernel k(θ, θ * ) = exp(-∥θ -θ * ∥ 2 /h) with the scale h > 0. For WEvidential, the scale h may be viewed as a hyperparameter to choose. We recommend using the value of the scale s.t. 0.01 ≤ h ≤ 1.0 in general. This work uses the value h = 0.1 throughout the experiments in the main text. One may opt for performing more advanced tuning of the scale h.We provide a simulation study to examine sensitivity to the scale value. We prepared a synthetic dataset {x i , µ i } D i=1 whose inputs are 200 gird points on the interval [-3.5, 3.5] and output distributions are normal distributions µ i (θ) = N (θ | sin(x i ), 0.5) conditional on each x i . We used WEvidential with different kernel scales h = 0.001, 0.01, 0.1, 1.0, 10, 100 and fitted it to the synthetic dataset {x i , µ i } D i=1 . The decision tree regressor with the maximum depth 3 was used for each weak learner. The learning rate and the number of weak learners were set to 0.1 and 100. The initial constant {ϑ n } 10 n=1 of WEvidential was set to 10 grid points in the interval [-10, 10]. We computed the output of WEvidential for 500 grid points in the interval [-3.5, 3.5] . We used the maximum mean discrepancy (MMD) [65] to measure the approximation error between the empirical distribution μi of the output of WEvidential and the output distribution µ i :MMD 2 (μ i , µ i ) = E θ∼μi,θ ′ ∼μi [k D (θ, θ ′ )] -2E θ∼μi,θ ′ ∼µi [k D (θ, θ ′ )] + E θ∼µi,θ ′ ∼µi [k D (θ, θ ′ )]where k D is a Gaussian kernel k D (θ, θ ′ ) = exp(-(θ -θ ′ ) 2 /s) with the scale s = 0.025. The total approximation error was measured by the MMD averaged over all the inputs. Figure 5 shows the total approximation error of WEvidential for each scale value h in the common log scale, together with examples of the output of WEvidential for different values of h. It demonstrates that the total error is minimised by h = 0.1 and stays in a relatively small value range for 0.01 ≤ h ≤ 1.0. We compare four WGBoost algorithms built on different estimates of the Wasserstein gradient of the KL divergence. The first three algorithms set the estimate G i (µ) in Algorithm 1 to, respectively, 1. the smoothed Wasserstein gradient in (7); 2. the diagonal approximate Wasserstein Newton direction in (8); 3. the full approximate Wasserstein Newton direction in (17) .The fourth algorithm is LGBoost in Appendix B which is rather a variant of WGBoost. The first and third algorithms are called the first-order WEvidential and the full-Newton WEvidential, respectively. The second algorithm is our default WEvidential presented in Section 3. The first-order WEvidential is implemented by removing h n i and replacing g n i ⊘ h n i with g n i in Algorithm 2. The full-Newton WEvidential is implemented by replacing g n i ⊘ h n i in Algorithm 2 with v n i computed by the following Algorithm 4, where ∇ 2 f (θ) denotes the Hessian matrix of a function f : Θ → R at θ.We prepared the same synthetic dataset {x i , µ i } D i=1 as Appendix D.1 and fitted each algorithm to the dataset. We computed the output of each algorithm for 500 grid points in the interval [-3.5, 3.5] and measured the approximation error by the MMD in the same manner as Appendix D.1. The decision tree regressor with the maximum depth 3 was used for each weak learner, and the learning rate was set to 0.1. We used an increasing number of weak leaners up to 100 weak learners, in order to observe the decay of the approximation error and the increase of the computational time. The initial constant {ϑ n } 10 n=1 for each algorithm was set to 10 grid points in the interval [-10, 10], which sufficiently differs from the output distributions so that the decay of the approximation error is clear.Figure 6 shows the approximation error and the computational time of each algorithm with respect to the increasing number of weak learners, together with the output of each algorithm with 100 weak learners trained. It demonstrates that WEvidential and full-Newton WEvidential reduce the approximation error most efficiently, while full-Newton WEvidential takes the longest computational time among others. As in Algorithm 4, the computation of the full approximate Wasserstein Newton direction requires the inverse and product of the (N × d) × (N × d) block matrices, where N denotes the particle number N and d denotes the particle dimension. The error decay of LGBoost is not only slow but also shows stochasticity due to the Gaussian noise used in the algorithm. We therefore recommend our default WEvidential for better performance and efficient computation. EQUATIONThis section describes additional details of the applications in Section 4. All the experiments were performed with x86-64 CPUs, where some of them were parallelised up to 10 CPUs and the rest uses 1 CPU. The scripts to reproduce all the experiments are provided in the source code. Appendices E.1 to E.3 describe additional details of the applications in Sections 4.1 to 4.3 respectively. Appendix E.4 describes a choice of initial constants {ϑ n 0 } N n=1 for the WGBoost algorithm used in Section 4.The normal response distribution N (y | m, σ) in Example 1 was used in Section 4.1. The normal response distribution has the scale parameter σ that lies only in the positive domain of the Euclidean space R. We reparametrised it as one in the Euclidean space R by the log transform σ ′ := log(σ), which is the standard practice in Bayesian computation [20] . The inverse of the log transform is the exponential transform σ = exp(σ ′ ). The change of variable formula tells the form of the individual-level posterior on (m, σ ′ ) up to the normalising constant. Under the prior in Example 1, the individual-level posterior on (m, σ ′ ) conditional on each individual response y i is given byp(m, σ ′ | y i ) ∝ exp - 1 2 (y i -m) 2 exp(σ ′ ) 2 × exp - 1 2 m 2 10 2 × 1 exp(σ ′ ) 1.01 exp - 0.01 exp(σ ′ ) ,where the Jacobian determinant |dσ/dσ ′ | = exp(σ ′ ) is used for the change of variable. Figure 7 shows the output of WEvidential for the old faithful geyser dataset and the conditional density estimated through (6) . where p(y i | x i ) and ŷi denote the predictive distribution and the point prediction provided by the algorithm. For WEvidential, we standardised the responses in the training data, in which case standardisation of test data were performed as y ′ i = (y i -y train mean )/y train std with the mean y train mean and standard deviation y train std of the training responses. Hence WEvidential provided the predictive distribution p(y ′ i | x i ) and the point prediction ŷ′ i for the standardised responses y ′ i . With no loss of generality, the NLL and RMSE for the original responses y i can be computed as follows:EQUATIONEQUATION)where the equality of the NLL follows from the change of variable p(y i | x i ) = p(y ′ i | x i )/y train std and the equality of the RMSE follows from rearranging the terms.We provide a brief description of each algorithm used in Section 4.2. For all the algorithms, the normal response distribution p(y | m, σ) was specified as the response distribution, where the algorithm produces a point or distributional estimate of the response parameter (m, σ) at each input x.• MCDropout [9] trains a single neural network F while dropping out each network weight with some Bernoulli probability. It can be interpreted as a variational approximation of a Bayesian neural network. It generates multiple subnetworks {F n } N n=1 by subsampling the network weights by the dropout. The predictive distribution p(y | x) is given by the model averaging (1/N ) N i=1 p(y | (m, σ) = F n (x)) for each input x.• DEnsemble [10] simply trains independent copies {F n } N n=1 of a neural network F in parallel. It is one of the mainstream approaches to uncertainty quantification based on deep learning. The predictive distribution is given by the model averaging as in MCDropout.• CDropout [55] consider a continuous relaxation of the Bernoulli random variable used in MCDropout to optimise the Bernoulli probability of the dropout. It generates multiple subnetworks {F n } N n=1 by subsampling the network weights by the dropout with the optimised probability. The predictive distribution is the same as MCDropout.• NGBoosting [32] is a family of gradient booting that use the natural gradient [66] of the response distribution as a target variable of each weak learner. In contrast to other methods, NGBoost outputs a single value F (x) to be plugged into the response-distribution parameter. The predictive distribution p(y | x) is given by p(y | (m, σ) = F (x)) for each input x.• DEvidential [13] extends deep evidential learning [11] , originally proposed in classification settings, to regression settings. It considers the case where the individual-level posterior of the response distribution falls into a conjugate parametric form, and predicts the hyperparameter of the individual-level posterior by a neural network. The predictive distribution is also given in a conjugate closed-form.The algorithms used in Section 4.2 are computationally efficient uncertainty quantification methods that are commonly-used in practice. In the following, we provide a limited yet additional comparison of WEvidential with other Bayesian neural networks (BNNs) and a kernel-based model, Gaussian process (GP). We used a subset of the datasets in Section 4, which were used in [67] who employed BNNs with one hidden layer of 50 units. We additionally used a large-scale dataset, keggd. For the keggd dataset only, we reported the NLL and RMSE of the normalised outputs without the adjustment (19) and (20) , in line with [67] . WEvidential was compared with BNNs learned by variational inference with the reparametrisation trick (VI), deterministic variational inference (DVI), stochastic weight averaging Gaussian (SWAG), principal component analysis subspace inference (PCA+VI), and two-layer deep Gaussian process with 50 induced points trained via expectation propagation (DGP1-50). Table 1 summarises the result. We reparametrised the parameter of the categorical response distribution C(y | q) used in Section 4.3, similarly to the normal response distribution used in Sections 4.1 and 4.2. The categorical response distribution has a class probability parameter q = [q 1 , . . . , q k ] in the simplex ∆ k . We applied the log-ratio transform q ′ := [log(q 1 /q k ), . . . , log(q k-1 /q k )] ∈ R k-1 that maps from the simplex ∆ k to the Euclidean space R k-1 [42] . The inverse of the log-ratio transform is the logistic transformq = exp(q ′ 1 ) z k , . . . , exp(q ′ k-1 ) z k , 1 z k ∈ ∆ k where z k = 1 + k-1 j=1 exp(q ′ j ).The logistic normal distribution on the original parameter q corresponds to a normal distribution on the transformed parameter q ′ [42] . The change of variable formula tells the individual-level posterior on q ′ up to the normalising constant. Under the prior in Example 2, the individual-level posterior p(q ′ | y i ) conditional on each individual observed response y i is given byp(q ′ | y i ) ∝ 1 z k [yi=k] × k-1 j=1 exp(q ′ j ) z k [yi=j] × exp - 1 2 ∥q ′ ∥ 2 10 2 ,where [y i = j] is 1 if y i is the j-th class label and 0 otherwise.We provide a brief description of each algorithm used in Section 4.3. MCDropout and DEnsemble are described in Appendix E.2.• DDistillation [57] predicts the parameter of a Dirichlet distribution over the simplex ∆ k by a neural network using the output of DEnsemble. The output of multiple networks in DEnsemble is distilled into the Dirichlet distribution controlled by one single network.• PNetwork [14] considers the case where the individual-level posterior of the categorical response distribution falls into a Dirichlet distribution similarly to deep evidential learning [11] . It predicts the hyperparameter of the individual-level posterior given in the form of the Dirichlet distribution.For the OOD detection, the OOD score we used was the inverse of the maximum norm of the variance of the WEvidential output. There are other quantities that are possible to use as the OOD score. For example, the entropy of the predictive distribution p(y | x) is a quantity computable for any probabilistic classification method. We compared the OOD detection performance of WEvidential with that of NGBoost and Random Forest (RForest) based on the entropy of their predictive distributions. For reference, we also computed the OOD detection performance of WEvidential based on the entropy of the predictive distribution. For NGBoost, the decision tree regressor was used for each weak learner and 4000 weak learners were trained with the learning rate 0.01. For RForest, the maximum depth 3 was used. Table 5 shows the result on the segment dataset, demonstrating that the performance of WEvidential is higher, to a large margin, than that of NGBoost and RForest based on the entropy. To investigate on the effective learning rate of WEvidential for classification, we performed a simple simulation study of WEvidential using the synthetic madelon dataset. We created 600 data points of three-dimensional inputs and three-class labels. The data subset that belongs to the last class was kept as the OOD samples. We randomly held out 20% of non-OOD samples as a test set to measure the test classification accuracy. We trained WEvidential using different learning rates from 0.1 to 1.0 while fixing the number of weak learners to 4000. Table 6 shows that the result for the learning rate 0.4 demonstrates the best classification accuracy and the third best OOD detection performance. We also trained WEvidential using different numbers of weak learners from 1000 to 4000 while fixing the learning rate to 0.4. Table 7 shows that the result for the number of weak learners 4000 demonstrates the best classification accuracy and the second best OOD detection performance. While the most effective learning rate differs depending on each dataset, we drew on the insight obtained from the simulation study and set the learning rate to 0.4 for classification. In standard gradient boosting, the initial state at step m = 0 is specified by a constant that most fits given data. Similarly, we specified the initial state {ϑ n 0 } N n=1 of WEvidential by a set of constants that most fits the output distributions in average. We find such a set of constants by performing an approximate Wasserstein gradient flow averaged over all the output distributions. Specifically, given the term G i (µ) in Algorithm 1, we define Ḡ(µ) := (1/D) D i=1 G i (µ) and perform the update scheme of a set of N particles { θn m } N n=1 :    θ1 m+1 . . . θN m+1    =    θ1 m . . . θN m    + ν 0    -[ Ḡ(π m )]( θ1 m ) . . . -[ Ḡ(π m )]( θN m )   with the learning rate ν 0 = 0.01 up to the maximum step number m = 5000. The initial particle locations for this update scheme were sampled from a standard normal distribution. We specified the initial state {ϑ n 0 } N n=1 by the set of particles { θn m } N n=1 obtained though this scheme at m = 5000.NeurIPS Paper Checklist 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This work proposed a generic algorithm. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: The credit and license were appropriately included in the text and source code. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix:  


Decision: 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.

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

Paper's major contributions include, the introduction of the Wasserstein gradient flow framework into gradient boosting; the development of an approximate algorithm for posterior regression using the KL divergence; and the demonstration of WGBoost's performance on regression, classification, and out-of-distribution (OOD) detection tasks. Authors also propose a second-order WGBoost algorithm built on the approximate Wasserstein gradient and Hessian of the KL divergence.
strength_weakness: strength_weakness: Authors did not specify details about the hyperparameter selection for Conditional Density Estimation, Classification and OOD Detection tasks.

Misc:
There is a typo in the line no 284, root is written as room.
questions: questions: I would be curious to know about the scalability and the computational complexity of the proposed algorithm.
limitations: limitations: NA
paper_summary: paper_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.
strength_weakness: strength_weakness: Originality: 

-	Nothing really noted. Being a straightforward application of some machinery is totally fine.

Quality:

-	For 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. 

-	Comparisons 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

-	The natural missing comparison here is to Gaussian processes and other kernel methods, which are naturally probabilistic and similarly nonparametric (like trees).

-	Another 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.

Clarity:

-	My 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.

         o	The code doesn’t seem to provide any clarity here.

-	Overall, 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. 

      o	For 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.



-	The writing is extremely passive and non-specific. Suggestions below.

       o	L150: “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.
       o	L113-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 …”

Significance:

-	Part 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.

      o	Bayesian neural nets: https://proceedings.mlr.press/v115/izmailov20a.html, https://proceedings.mlr.press/v130/immer21a/immer21a.pdf, 

      o	Gaussian processes: https://arxiv.org/abs/1809.11165,
questions: questions: Algorithm 1: 

-	Is the Wasserstein gradient where the particles manage to interact? Otherwise, I see nowhere else the various particles / learners would interact.

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?

-	In 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)

-	the 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.

Section 3.3: 
-	so 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”)

Section 3.1:

L190: “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.

Table 1: what are the error bars?

Table 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: limitations: n/a
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 
- 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. 
- 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).

- 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.

- 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.
- 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.
- 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.

[1] Stein Points 

[2] Measuring Sample Quality with Kernels

[3] Kernel Stein Discrepancy Descent

[4] On the geometry of Stein variational gradient descent

[5] FUNCTION SPACE PARTICLE OPTIMIZATION FOR BAYESIAN NEURAL NETWORKS

[6] Feature Space Particle Inference for Neural Network Ensembles

[7] Repulsive Deep Ensembles are Bayesian
questions: questions: I wrote questions in the Weakness.
limitations: limitations: The limitation is unclear. As far as I read no formal description is presented.
paper_summary: paper_summary: This work proposes to perform boosting with base learner which are fitted to the Wasserstein gradient of a loss function on the space of probability distributions, which can be useful in particular to capture uncertainty of the models.

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

- Experiments focus on KL divergence functional.
- No theoretical analysis
questions: questions: Did you compare between the different algorithms (i.e. with and without Wasserstein Hessian preconditioner, and Kernel vs Langevin approximation) on the benchmark of Section 4.2?

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

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

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

Comments: title: Final Author Comment
Response to Reviewer
Response to Reviewer
Thank you for the reply
Response to Reviewer Comment
Response to Reviewer
comment: We would like to express our gratitudes for all the discussions. While Reviewer jx13 & DY4u agree on the novelty of our work, the remaining concern of Reviewer UM7J & q9Jy is their interpretation that WGBoost is essentially same as SVGD. 

In our view WGBoost is clearly different from SVGD as attempted to clarify in the rebuttal, but some of our explanation made from the viewpoint of regression so far may sound differently to some reviewer e.g. expert in sampling contexts. In this comment, we wanted to summarise our last explanation about the difference of WGBoost and SVGD from the viewpoint of sampling. 


## Difference of WGBoost and SVGD from Viewpoint of Sampling

While SVGD samples from one distribution, WGBoost samples from a **conditional** distribution as follows:

- **SVGD** gets samples $(\theta_1, \dots, \theta_N)$ from a distribution $P(\theta)$, where let's say $P(\theta)$ is an arbitrary distribution over an arbitrary variable $\theta$. 
- **WGBoost** gets samples $(\theta_1(x), \dots, \theta_N(x))$ from a conditional distribution $P(\theta \mid x)$ for any input $x$. WGBoost learns the conditional distribution in the setting where we can know the form of $P(\theta \mid x_i)$ only at given data inputs $( x_i )\_{i=1}^{n}$. In the algorithm, WGBoost **predicts** the Wasserstein gradient used to sample $(\theta_1(x), \dots, \theta_N(x))$ for each input $x$ by tree-model. 

SVGD do not have input variable $x$ so do not have the prediction stage of the Wasserstein gradient unlike WGBoost. In this sense, WGBoost is a conditional version $P(\theta \mid x)$ of a Wasserstein gradient flow of $P(\theta)$, where the Wasserstein gradient of the particles $(\theta_1(x), \dots, \theta_N(x))$ will be predicted for each input $x$ because the form of $P(\theta \mid x)$ can be known only at $x=x_i$.

We hope the contrast from SVGD is clear in this explanation. As the author-reviewer discussion period comes to end, we will leave the final decision to the AC and all the reviewers. Thank you once again for all your effort.
Thank you for addressing my comments. I do not have further questions and I will keep my score at 8.
Thank you for your response and the opportunity to clarify.

> the algorithms and their mathematical foundations might be essentially similar (to SVGD)

Please let us make our last attempt to clarify the difference. Here we will compare them from the viewpoint of sampling. While SVGD samples from one distribution, WGBoost samples from a **conditional** distribution as follows:

- **SVGD** gets samples $(\theta_1, \dots, \theta_N)$ from a distribution $P(\theta)$, where let's say $P(\theta)$ is an arbitrary distribution over an arbitrary variable $\theta$. 
- **WGBoost** gets samples $(\theta_1(x), \dots, \theta_N(x))$ from a conditional distribution $P(\theta \mid x)$ for any input $x$. WGBoost learns the conditional distribution in the setting where we can know the form of $P(\theta \mid x_i)$ only at given data inputs $( x_i )\_{i=1}^{n}$. In the algorithm, WGBoost **predicts** the Wasserstein gradient used to sample $(\theta_1(x), \dots, \theta_N(x))$ for each input $x$ by tree-model. 

SVGD do not have input variable $x$ so do not have the prediction stage of the Wasserstein gradient unlike WGBoost. In this sense, WGBoost is a conditional version $P(\theta \mid x)$ of a Wasserstein gradient flow of $P(\theta)$, where the Wasserstein gradient of the particles $(\theta_1(x), \dots, \theta_N(x))$ will be predicted for each input $x$ because the form of $P(\theta \mid x)$ can be known only at $x=x_i$ in our setting.

I hope this way of explanation concisely shows the difference. We are keen to hear if the reviewer's score remains unchanged.
Thank you for your response and the additional opportunity to discuss.

> However, it's not clear to my why a well specfied probabilistic model can't just be used for OOD, like the early deep net based approaches were - by entropy of the multi-class predictive distribution (which NGBoost can surely gather).

Thank you for the suggestion. As in the previous comment, we observed that entropy of multi-class classifier didn't perform well and decided to focus only on the uncertainty done in the current paper for the page limit. We will include entropy of multi-class classifier in the comparison. In our understanding, our experiment has already shown better performance of WGBoost on small to medium-large datasets, so adding large-scale dataset and other comparison would complement our experiment at sufficient level.


> underlying theory is quite similar (with SVGD)

Thank you for the discussion about this point. Infinite-dimensional SVGD has been studied and the algorithm/theory is different [1]. Please let us attempt to clarify the difference again. What WGBoost learn is a conditional distribution $P(y | x) = \int P(y | q) P(q | x) dq$, where let's say $P(y | q)$ is a categorical distribution and $P(q | x)$ is a conditional distribution of the class probability vector $q$ given $x$. In this view, WGBoost is a sampler from a conditional distribution with tree-model incorporated in the procedure as follows: 

- **SVGD**: Let's say $P(\theta)$ is an arbitrary distribution over an arbitrary variable $\theta$. SVGD gets samples $(\theta_1, \dots, \theta_N)$ from a distribution $P(\theta)$ (regardless of $\theta$ is finite-dimensional or replaced with infinite-dimensional variable $f$); 
- **WGBoost** gets samples $(\theta_1(x), \dots, \theta_N(x))$ from a conditional distribution $P(\theta \mid x)$ for any input $x$. WGBoost learns the conditional distribution in the setting where we can know the form of $P(\theta \mid x_i)$ only at given data inputs $( x_i )\_{i=1}^{n}$. In the algorithm, WGBoost **predicts** the Wasserstein gradient used to sample $(\theta_1(x), \dots, \theta_N(x))$ for each input $x$ by tree-model. 

SVGD do not have input variable $x$ so do not have the prediction of the Wasserstein gradient unlike WGBoost. In this sense, WGBoost is a conditional version $P(\theta \mid x)$ of a Wasserstein gradient flow of $P(\theta)$, where the Wasserstein gradient of the particles $(\theta_1(x), \dots, \theta_N(x))$ will be predicted for each input $x$ (because the form of $P(\theta \mid x)$ can be known only at $x=x_i$).

We hope this way of explanation clarifies the difference better? We are keen to hear if the reviewers opinion still remains unchanged.

[1] Stein variational gradient descent on infinite-dimensional space and applications to statistical inverse problems
Thank you for your response to my questions. 

However, I still have concerns. Specifically, the fundamental and mathematical differences between SVGD and your proposed methods remain unclear. Although you mention that the purpose of the methods differ—sampling versus boosting—the core principle in both cases involves minimizing the objective functional, such as KL, L2, within an RKHS framework [1, 2]. This suggests that, despite the different objectives, the algorithms and their mathematical foundations might be essentially similar. Therefore, I would like to keep the score as is.

[1] Stein Variational Gradient Descent as Gradient Flow
[2] On the geometry of Stein variational gradient descent
Thanks for your further responses .

> WGBoost is the first method that enables the tree version of DEL approaches.

Indeed, this is a selling point (and a drawback as the inference should be more expensive..) of your approach as compared to other tree based methods from my understanding. However, it's not clear to my why a well specfied probabilistic model can't just be used for OOD, like the early deep net based approaches were - by entropy of the multi-class predictive distribution (which NGBoost can surely gather).

> Difference between SVGD and WGBoost

My (somewhat non-rigorous) understanding is that with WGBoost you're essentially non-parametrically estimating the integral $p(\hat y | y, \mathcal{M}) = \int p(\hat y | f(x)) p(f(x) | x) df$ where $f$ is the (non-parametric) functional form induced by the model class. SVGD ends up being quite similar, but (in its [original definition](https://arxiv.org/pdf/1608.04471)) operates on a _parametric_ form of model class, estimating $p(\hat y | y, \mathcal{M}) = \int p(\hat y | f_\theta(x) ) p(\theta | x) d\theta \approx p(\hat y | f_\theta(x) ) q(\theta) d\theta$, where $q(\theta) $ is the approximation distribution learnt by SVGD. However, as SVGD is a gradient flow [paper](https://arxiv.org/abs/1704.07520), it is also possible to express as a nonparametric algorithm in my understanding if we express it in terms of function space rather than parameter space.

I guess practically they end up being different algorithms, but the underlying theory is quite similar, if not exactly the same. And there's several different function space algorithms that use stein flows as the other reviewer points out - training either neural networks or GPs with gradient flows like this is fairly well studied.
We appreciate the reviewer for providing further comments and opportunity for us to discuss about them. All your comments will be used and helpful to improve our manuscript.

(PS: Please refresh this webpage in case the latex commands were not rendered as math equations properly.)

&nbsp;

## Difference from NGBoost and Quantile Regressor

To illustrate, we would like to consider a classification problem, where we have input $x$ and output label $y$. We assume a categorical distribution $C(y \mid q)$ over the label $y$ given the class probability vector $q$. What NGBoost (and two-stage fit model) will provide is a **point estimate of the class probability** $q$ at each input $x$, that is

- $x \to \text{NGBoost} \to q(x) \to C(y \mid q(x))$

On the other hand, what WGBoost will provide is a particle-based **distributional estimate of the class probability** $q$:

- $x \to \text{WGBoost} \to ( q_1(x), \dots, q_N(x) ) \to (1 / N) \sum_{i=1}^{N} C(y \mid q_i(x))$

What we show in the paper is that (i) taking the average over the distributional estimate leads to better performance (in Section 4.2 with better performance than NGBoost for regression) and (ii) uncertainty in the distributional estimate can be used for the OOD detection (in Section 4.3). These advantages are also different from quantile regressors since quantile regressors provide quantile of output variables $y$ only i.e. no distribution on $y$ nor distribution of $q$. (Such a distributional estimate of the class probability $q$ has been explored in Deep Evidential Learning (DEL) approaches. In this context, WGBoost is the first method that enables the tree version of DEL approaches.)

&nbsp;

## Difference from SVGD

Firstly, we would like to clarify difference in what WGBoost and SVGD can do:

- **SVGD**: Given one target density $\pi$, we can obtain samples $(\theta_1, \dots, \theta_N)$ form $\pi$.
- **WGBoost**: Given a set of several inputs and output-densities $( x_i, \pi\_{x_i} )\_{i=1}^{n}$, we can obtain **a map from input $x$ to particles $(\theta_1(x), \dots, \theta_N(x) )$** that approximates the given $\pi\_{x_i}$ for in-sample input $x_i$ and also predicts an unseen $\pi\_{x}$ for new out-of-sample input $x$.

Next, we would like to clarify how WGBoost and SVGD is applied for clarification/regression. Let us consider the classification again, where we have input/label dataset $(x_i, y_i)\_{i=1}^{n}$:

- **SVGD**: We have a model $f(x; w)$ with the parameter $w$ that outputs the class probability $q$. We use a **posterior $Pos(w \mid ( x_i, y_i )\_{i=1}^{n} )$ of the model $f(x; w)$ using all data points**. SVGD produces samples $(w_1, \dots w_N)$ from the posterior $Pos(w \mid ( x_i, y_i )\_{i=0}^{n} )$.
- **WGBoost**: WGBoost is a learning model to estimate $q$ distributionally, so there is no other model $f(x; w)$ like the SVGD case. We use a **posterior $\pi\_{x_i}$ of the categorical distribution $C(y \mid q)$ for each single data point** $(x_i, y_i)$ like done in some DEL approaches:

  - $\pi\_{x_i}(q) \propto C(y_i \mid q) \times \nu\_{x_i}(q)$ for each $(x_i, y_i)$, where $\nu\_{x_i}$ is a prior over $q$ given at each $(x_i, y_i)$.

  WGBoost produces a map from input $x$ to particles s.t. $x \to \text{WGBoost} \to ( q_1(x), \dots, q_N(x) ) \approx \pi_x$. 

&nbsp;

Finally, please let us elaborate the intuitive algorithmic idea of WGBoost. The challenge is how to produce particles for unseen input $x$:

1. We have the given set $( x_i, \pi\_{x_i} )\_{i=1}^{n}$. First let's sample from every given $\pi\_{x_i}$ using any user-choice of Wasserstein gradient flow (WGF) such as SVGD;
2. The WGF for each $\pi\_{x_i}$ uses the Wasserstein gradient $g\_{x_i}$, so we get a set of the computed Wasserstein gradients $(x_i, g\_{x_i})\_{i=1}^{n}$. Since $g\_{x_i}$ is finite-dimensional, we can train a ML model that predicts the gradient $g_x$ for new $x$ by fitting it to the set $(x_i, g\_{x_i})\_{i=1}^{n}$.
3. The trained ML model gives us a prediction of the Wasserstein gradient $g_x$ for unseen input $x$. So let's perform the WGF with the predicted gradient and have predictive particles even for unseen input $x$.

This procedure uses a user-specified WGF (e.g. SVGD) for every $\pi\_{x_i}$ and train a ML predictive model with the computed Wasserstein gradient at $x_i$. So SVGD is simply used as an intermediate component of WGBoost. Rigorously formalising this idea requires academic work, for which we showed that it can be formalised as an extension of gradient boosting.

&nbsp;

## Comment

Would this clarify difference from NGBoost and SVGD? We have added these better clarifications in comparison with SVGD to the main text. To date, adapting Wasserstein gradient flows to other gradient-related ML methods is still underdeveloped. We truly believe this work can bring a new inspiration to the ML community, bridging Wasserstein gradient flows and gradient boosting for the first time. We are keen to hear any opinion from the reviewer.
I thank the authors for their clarifications. However, without further experimental comparisons, my opinions on the paper are unchanged.

In terms of experimental benefits, I don't think there's a compelling application as described in the paper to use Wasserstein boosting for tree methods as compared to something like NGBoost (or other probabilistic baselines like 2 stage fits, quantile regression etc.). Wasserstein boosting provides non parametric uncertainty estimates, which NG Boost and quantile regression on trees also provide. I understand that Wasserstein boosting could be used on top of a quantile regressor or the two stage fitter, but these experiments just aren't done presently so there's, to my understanding, not as much of a compelling story here.

re " , Subspace Inference for Bayesian Deep Learning. " I believe the pre-processing may be slightly different (you can see by comparing their "SGD" baseline to your deep ensemble baseline), but this is good to understand that the Wasserstein boosting algorithm performs quite sensibly.

I also think that, even after the rebuttal, the comparison with SVGD is a bit confusing to me (as mentioned by reviwer q9Jy), and would like to see the authors further clarify here. The algorithm seems to be different primarily in its applications to boosting, but in some sense this is quite similar to the application in parametric models as well.
We appreciate your response and again all your efforts in reviewing our manuscript.
Thank you for your answer and addressing my comments. I do not have further questions and I will keep my score at 7.