Stochastic Multiple Target Sampling Gradient Descent
Authors: Hoang Phan
Ngoc N. Tran
Trung Le
Toan Tran
Nhat Ho
Dinh Phung
Venue: NeurIPS 2022 Accept (2022)
Keywords: Multi objective optimization
Multi task learning
Stein Variational Gradient Descent

Abstract: Sampling from an unnormalized target distribution is an essential problem with many applications in probabilistic inference. Stein Variational Gradient Descent (SVGD) has been shown to be a powerful method that iteratively updates a set of particles to approximate the distribution of interest. Furthermore, when analysing its asymptotic properties, SVGD reduces exactly to a single-objective optimization problem and can be viewed as a probabilistic version of this single-objective optimization problem. A natural question then arises: ``Can we derive a probabilistic version of the multi-objective optimization?''. To answer this question, we propose Stochastic Multiple Target Sampling Gradient Descent (MT-SGD), enabling us to sample from multiple unnormalized target distributions. Specifically, our MT-SGD conducts a flow of intermediate distributions gradually orienting to multiple target distributions, which allows the sampled particles to move to the joint high-likelihood region of the target distributions. Interestingly, the asymptotic analysis shows that our approach reduces exactly to the multiple-gradient descent algorithm for multi-objective optimization, as expected. Finally, we conduct comprehensive experiments to demonstrate the merit of our approach to multi-task learning.

Stochastic Multiple Target Sampling Gradient Descent
Hoang Phan'
Ngoc N. Tran'
Trung Le
Toan Tran
Nhat Ho
Dinh Phung
VinAl Research, Vietnam
2 Monash University, Australia
3 University of Texas, Austin
Abstract
Sampling from an unnormalized target distribution is an essential problem with many applications in probabílistic inference. Stein Variational Gradient Descent (SVGD) has been shown to be powerful method that iteratively updates a set of particles to approximate the distribution of interest. Furthermore, when analysing its asymptotic properties, SVGD reduces exactly to single-objective optimization problem and can be viewed probabilistic version of this single-objective opti- mization problem. A natural question then arises: "Can we derive probabilistic version of the multi-objective optimization?"" ?"" To answer this question, we propose Stochastic Multiple Target Sampling Gradient Descent (MT-SGD), enabling us sample from multiple unnormalized target distributions. Specifically, our MT- SGD conducts a flow of intermediate distributions gradually orienting to multiple target distributions, which allows the sampled particles move to the joint high- likelihood region of the target distributions. Interestingly, the asymptotic analysis shows that our approach reduces exactly to the multiple-gradien descent algorithm for multi-objective optimization, as expected. Finally, we conduct comprehensive experiments to demonstrate the merit of our approach multi-task learning.
Introduction
Sampling from an unnormalized target distribution that we know the density function up to a scaling
factor is a pivotal problem with many applications in probabilistic inference [2] I8 25] For this purpose, Markov chain Monte Carlo (MCMC) has been widely used to draw approximate posterior samples, but unfortunately, is often time-consuming and has difficulty accessing the convergence [14). Targeting an efficient acceleration of MCMC, some stochastic variational particle-basse basec approaches have been proposed, notably Stochastic Langevin Gradient Descent (26] and Stein Variational Gradient Descent (SVGD) [4]. Outstanding among them is SVGD, with a solid theoretical guarantee of the convergence of the set of particles to the target distribution by maintaining flow of distributions. More specifically, SVGD starts from an arbitrary and -o--amll initial distribution and learns the subsequent distribution in the flow by push-forw forhddorwr the current one using function T (x) = x + cΦ (z),, where x E Rd c > 0 is the learning rate, and á E Hd with H4 to be the Reproducing Kernel Hilbert Space corresponding
kernel k. It is well-known that for the case of using Gaussian RBF kernel, by letting the kernel width approach +oo, the update formula of SVGD at each step asymptotically reduces to the typical gradient descent (GD) [14, showing the connection between probabilistic framework like SVGD and a single-objective optimization algorithm. In other words, SVGD can be viewed probabilistic version of the GD for single-objective optimization.
On the other side, multi-objective optimization (MOO) 7 aims to optimize a set of objective functions and manifests itself in many real-world applications problems, such as in multi-task learning (MTL) [ 24]. natural language processing M. and reinforcement learning 9 22 211.
Leveraging the above insights, it is natural to ask: "Can we derive a probabilistic version of multi-
36th Conference on Neural Information Processing Systems (NeurIPS 2022).

---

objective optimization?". By answering this question, we enable the application of the Bayesian inference framework to the tasks inherently fulfilled by the MOO framework.
Contribution. In this paper, we provide an affirmative answer to that question. In particular, we go beyond the SVGD to propose Stochastic Multiple Target Sampling Gradient Descent (MT-SGD), enabling us to sample from multiple target distributions. By considering the push-forward map (x) = + eω (x)) with Φ E H.g we find closed-form for the optimal push-forward map T pushing the current distribution on the flow simultancously closer to all target distributions. Similar to SVGD, in the case of using Gaussian RBF kernel, when the kernel width approaches toc, MT-SGD reduces exactly to the multiple-gradient descent algorithm (MGDA) [7] for multi-objective optimization (MOO). Our MT-SGD, therefore, can be considered as probabilistic version of the GD multi-objective optimization 7 as expected.
Additionally, in practice, we consider flow of empirical distributions, in which, each distribution is presented as a set of particles. Our observations indicate that MT-SGD globally drives the particles to close to all target distributions. leading them to diversify on the joint high- high-likelihood region for all distributions. It worth noting that, different from other multi-partich leading the particles to diversify on a Pareto front, our MT-SGD orients the particle to diversify on the so-called Pareto common (i.e., the joint high-likelihood region for all distributions) (cf. Section 2.4 for more discussions). We argue and empirically demonstrate that this characteristic is essential for the Bayesian setting, whose main goal estimate the ensemble accuracy and the uncertainty calibration of a model. In summary, we make the following contributions in this work:
approaches [13 15 [7
Propose a principled framework that incorporates the power of Stein Variational Gradient Descent into multi-objective optimization. Concretely, our method is motivated by the theoretical analysis of SVGD, and we further derive the formulation that extends the original work and allows to sample from multiple unnormalize distributions.
Demonstrate our algorithm is readily applicable in the context of multi-task learning. The benefits of MT-SGD are twofold: i) the trained network is optimal, which could not be improved in any task without diminishing another, and ii) there is no need for predefined preference vectors in previous works [7] MT-SGD implicitly learns diverse models universally optimizing for all tasks.
Conduct comprehensive experiments to verify the behaviors of MT-SGD and demonstrate the superiority of MT-SGD to the baselines in Bayesian setting, with higher ensemble performances and significantly lower calibration errors.
Related works. The work of 7 proposed a multi multi-grrddent descent algorithm for multi-objective optimization (MOO) which opens the door for the applications of MOO in machine learning and deep learning. Inspired by 77 MOO has been applied in multi-task learning (MTL) [I7 24), few-shot learning [8) and knowledge distillation [. Specifically, in an earlier attempt at solving MTL, 24 viewed multi-task learning as multi-objective optimization problem, where a task network consists of shared feature extractor and a task-specifie predictor. In another study, [17] developed gradient-based multi-objective MTL algorithm to find a set of solutions that satisfies the user preferences. Also follows the idea of learning neural networks conditioned on pre-defined preference vectors, [3] proposed Pareto MTL, aiming to find set of well-distributed Pareto solutions, which can represent different tradc-off's among different tasks. Recently, the work of [15) leveraged MOO with SVGD [1 and Langevin dynamics [26] diversify the solutions of MOO. In another line of work, 28 proposed a bi-level MOO that can be applied to few-shot learning. Furthermore, somewhat different result was proposed, 8 applied MOO to enable the knowledge distillation from multiple teachers and find a better optimization direction in training the student network.
The paper is organized follows. In Section 2 first present our theoretical contribution by reviewing the formalism and providing the point of view adopted to generalize SVGD in the context of MOO. Then, Section 3 introduces algorithm to showcase the application of our proposed method in the multi-task learning scenario. We report the results of extensive experimental studies performed on various datasets that demonstrate the behaviors and efficiency of MT SGD in Section 4 Finally, we conclude the paper in Section 5 The complete proofs and experiment setups are deferred to the supplementary material.
Outline.

---

Multi-Target Sampling Gradient Descent
We first briefly introduce the formulation of the multi-target sampling in Section 2.1 Second, Section 2.2 presents our theoretical development and shows how our proposed method is applicable to this problem. Finally, we detail how train the proposed method in Section 2.3 and highlight key differences between our method and related work in Section 2.4
2.1 Problem Setting
Given set of target distributions PI:K (0) == find the optimal distribution q" E that minimizes a vector-valued objective function whose k-th component is D (glp))))):
(0)} with parameter 0 E Rd we aim
min [DKE (allpi) ...., D (qllpk)] gEQ
where DR. represents Kullback-Leibler divergence and Q is family of distributions.
When there exists objective function that conflicts with each other, there will be trade-off between these two objectives. Therefore, no "optimal solution" exists in such cases. Alternatively, we are often interested in seeking a of solutions such that each does have any better solution (i.e. achieves lower loss values in all objectives)
13 15 (24). The optimization problem (OP) in T thus be viewed multi-objective OP 7 the probability distribution space. Let denote H&
by the Reproducing Kernel Hilbert Space (RKHS) associated with a positive semi-definite (p.s.d.) kernel and Hgg by the d-dimensional vector function:
f = [/,... Sa), (Ji E Hc.)
Inspired by [14], we construct flow of distributions go: .1. q departed from simple distribution qo. that gradually move closer to all the target distributions. In particular, at each step, assume that q the current obtained distribution, and the goal is learn transformation T id + ed so that the feed-forward distribution q(7) T#q moves closer to Pı:K simultancously. Here use id to denote the identity operator, e > is a step size, and á E Hi is a velocity field. Particularly, the problem of finding the optimal transformation T for the current step is formulated
min DKL
DEL (19]
(2)
2.2
Our Theoretical Development
It is worth noting that the transformation T defined above is injective when is sufficiently small 14 We consider each DrL. (114 ; 1,... K as function w.r.t. by applying the first-order
Taylor expansion 0, we have:
DKL (q I1P) DEL (alle)) V.DKL (p)) ==
c+0(2).
where lim (2) / const.
Given that the velocity field E Hd.. similar 14. the gradient V.DKL (1 Ie=o can be calculated as provided in
V.D.T. (1 |--
(ò, che)
c=
where v/s (.) Ro [k: (0.-) Vo log (0) Vok (0.:) and (:, 'uu the dot product in the RKHS.
This means that, for each target distribution pa, the steepest descent direction oं = th, in which KL divergence of interest DKL (pp:) gets decreased roughly by -e || l| toward the target distribution P;. However, this only guarantees a divergence reduction for single target distribution p: itself. Our next aim is hence to find a common direction o" to reduce the KL divergences w.r.t. all target distributions, which is reflected in the following lemma, showing how to combine the individual steepest descent direction o = th, to yield the optimal direction o as summarized in Figure

---

Pi
PR
1!1
ii
6 H
1,
K
K
1 of
w" argmin w
Uw
Figure 1: How to find the optimal descent direction Ø².
Lemma 1. Let w be the optimal solution of the optimization problem w = argmin T Uw and wEak
6" Ek wf where Ax = =t E IR4 || = 1 and U E RKXK with Uúj = (0) 07) Hg. then we have
(6 ,00,) 1100 1104 = L....K.. K.
Lemma [ provides a common descent direction 6² so that all KL divergences w.r.t. the target distributions are consistently reduced by e 11 roughly and Theorem 2 confirms this argument.
Theorem 2. If there does not exist w E Ax such that Σ wi = 0, given sufficiently small step size e, all KL divergences w.r.. the target distributions are strictly decreased by at least A |||0"|| 12 > 0 where A is positive constant.
The next arising question is how to evaluate the matrix U E RRKK with U.j = (0) (01,07) 35) for He
solving the quadratic problem: min w Uw. To this end, using some well-known equalities in ] we arrive at the following formula:
Un (01,0)) 05) Eg, oy k (0, 0') (V)),,((), V (0')
Ōk(0, 0')
+ V + (0') Ok(0,0))
(02k(0,0))
V log
00 + tr
(3)
00'
0000'
where tr(.) denotes the trace (square) matrix.
2.3
Algorithm for MT-SGD
For the implementation of MT-SGD, we consider as a discrete distribution over a set of M, ,
(M E N*) particles 01, .2,..., O0 q. The formulation evaluate Uti in Equation. 3 becomes:
k (0a, 0b) (V log V (logp;(0.)) + V (og); ak(O,, 0z)
Un
M
00,
Ok(O., 0b)
(O 0m)
V (0.),
(4)
a0,
30,00g
The optimal solution of then can be computed as:
M
- = M
M
(0,, Ve, logpi (0,) + (0,,
(5)
1=1
'll proof's and derivations be found in the supplementary material.

---

The key steps of our MT-SGD are summarized in Algorithm where the set of particles 01:M is updated gradually
approach the multiple distributions P1:K- Furthermore, the update formula consists of two terms: (i) the first term (i.e., relevant to k (0). -) Ve, log Ps (f,)) helps to push the particles to the joint high-like gighliikelhood region for all distributions and (ii) the second term (i.e., relevant to Ve, & (0)s-)) which is repulsive term to push away the particles when they reach out each other.
Finally, we note that our proposed MT-SGD can be applied in the context where we know the target distributions up to scaling factor (e.g., in the posterior inference).
Algorithm Pseudocode for MT-SGD.
Input: Multiple unnormalized target densities PI:K. Output: The optimal particles 0,, ....... Oas-
1: Initialize set of particles 01, 02.0............... OM 40
2: for to I. do
3: Form the matrix U E RR with the element U2j computed as in Equation. 4
4: Solve the QP min w"w to find the optimal weights w" E Ax..
5: Compute the optimal direction o3 (+) EE wi (-), where of is defined in Equation. 5
6: Update 0g = 0; Exp (0s), 1,..,,
7: end for
8: return 01,02. 02..................... 01, ass-
Analysis for the case of RBF kernel. We now consider a radial basis-function (RBF) kernel of bandwidth a: (0, 8') exp
||0 0'1 / (2o²) and examine some asymptotic behaviors.
General case: The elements of the matrix U become
110 - 1²
}
Uij = E0,00 B0,0r~g m
exp
V Voog p: (f), V Vlog (B'))
202
1 (Vlog,,())
d |0 8/(3)
+ σ²
V (logp;(')), 0¹
o² a
Single particle distribution g 50: The elements of the matrix U become
Ui = < log (6), V 1o
and our formulation reduces exactly to MOO in 7
When a - oct The elements the matrix U become
Ui = Eg,0r wa V (Vlogp(()) Vlog,,())
2.4 Comparison to MOO-SVGD and Other Works
The most closely related work to ours is MOO-SVGD [5]. In a nutshell, ours is principally different from that work and we show a fundamental difference between our MT-SGD and MOO-SVGD in Figure 2 Our MT-SGD navigates the particles from one distribution to another distribution consecutively with a theoretical guarantee of globally getting closely multiple target distributions. By contrast, while MOO-SVGD also uses the MOO [7] to update the particles. their employed principle control
repulsive term encourages the particle diversity without any enoorttial-guuarantee the repulsive term, hence it can force the particles to scatter on the multiple distributions. In fact, they aim profile the whole Pareto front, which is preferred when users want obtain collection of diverse Pareto optimal solutions with different trade-off's among all tasks.
Furthermore, it expects that MT-SGD globally moves the set of particles to the joint high- likelihood region for all target distributions. Therefore, we do not claim our MT-SGD as a method to diversify the solution can generate diverse particles on the so-called Pareto common (i.c., the joint high-likelihood region for all target distributions). We argue and empirically demonstrate that by finding and diversifying
Pareto front for user preferences, as in [15] 17]. Alternativel, our MT-SGD
5

---

MOO-SVGD
MT-SGD
Step
Step
Figure 2: Our MT-SGD moves the particles from distribution another distribution to globally get closer to two target distributions (i.e., the blue and green ones). Differently, MOO-SVGD uses MOO M71 to move the particles individually and independently. The diversity enforced by the repulsive forces among particles. There no principle to control these repulsive forces, hence they push the particles scattering on two distributions.
the particles on Pareto common for the multiple posterior inferences, our MT-SGD can outperform the baselines Bayesian-inference metries such the ensemble accuracy and the calibration error.
Morcover, MOO-SVGD is not computationally efficient when the number of particles is high because it requires solving independent quadratic programming problem for each particle (ef. Section 4.1.] and Figure for the experiment on a synthetic dataset). Instead, our work solves a single quadratic programming problem for all particles, then update them accordingly. We verify this computational improvement and present the training time of all baselines in the supplement material.
3 Application to Multi-Task Learning
K and training set {(z,, Vi.,....kk.).. M..) where r, is a data example and yii..... V&K are the labels for the tasks. The model for each task 6i = [a, 33] .j = 1,... K consists of the shared part a and non non-shared part 3i targeting the task The posterior (0) DD) for each task reads
For multi-task learning, we assume to have K tasks (Ta)
p (08
D
xp 0*) p (00) x I p (yis Iz. 04)
N x exp (Viy: x4; 00) = exp Σ
N
(Vis. Tz 83)
where e is a loss function and the predictive likelihood p(ut examined. Note that the prior (0') here is retained from previous studies 13 15). which is uniform and innimmformttve prior and can be treated as a constant term in our formulation.
23, 05) x exp (Bij. 2i 8*) is
For our approach, we maintain a set of models O. = [0%] K
with m = 1,... M, where 0f
j=1
[amm ]. At each iteration, given the non-shared parts [[0]] j=1 with Bi [[ =[
..., K M sample
the shared parts from the multiple distributions (a _ D) .j
Omm ~ p (a (3,D) x ( D) x p (0' ID) x exp
Σ (vi my;
(6)
We now apply our proposed MT-SGD to sample the shared parts [amlm= from the multiple distribu- tions defined in (6) as
Qm am + wj0z (am),

---

where 0j (am)
M Di= k1 (at, am) Va, logp (as D + Va, (ar, am) and w²
1
[ug]k- are the weights received from solving the quadratic programming problem. Here we note that Va, logp (ce IB. can be estimated via the batch gradient of the loss using Equation (6)
Given the updated shared parts [0km)==1+ for each task j. we update the corresponding non-shared parts
[8% 1 m=1 by sampling
N P p (30 a, 1) & p8 (D) x (0' (D) x exp E((
(8)
We now apply SVGD 14 to sample the non-shared parts [3m]== for each task j from the distribution defined in 8
8 3 + 5 (B%m)
where 4i5 (8%)
3े Ear Ka (9l, 8%,) V log p (3d a. aa, V 0t k2 (36, with which
the term V #g log (B: Cle IB) can be estimated via the batch loss gradient using Equation 8
Algorithm Pseudocode for multi-task learning MT-SGD.
Input: A training set D = { Wa1. Vik )
Output: The models On = [ و-= with m = 1,.., M, where 0f, aams 3
1: Initialize a set of particles 01:M qu
2: for epoch 1 #epoch do
3: for iter = to #iter do
4: Update the shared parts [am]m=1 using Equation.
5: for j = to K do
Update the non-shared part [8) m=1 using Equation. 7: end for
6:
8: end for
9: end for
10:
return 01:M-
Algorithm summarizes the key steps of multi-task MT-SGD, Basically, we alternatively update the shared parts given the non-shared ones and vice versa.
Experiments
In this section, we verify our MT-SGD by evaluating its performance on both synthetic and real- world datasets. For our experiments, detailed training and configuration are given in the supplementary material. Our codes are available /ttttps://githu.. com http:/ttwb..oom//iiitttlonnnggg/ /ttmm//ieetooonn151///1----
the RBF kernel (e, 8') exp { ||0 8'112/ (2o*)} } The
4.1 Experiments on Toy Datasets
4.1.1 Sampling from Multiple Distributions
We first qualitatively analyze the behavior of the proposed method on sampling from three target
distributions. Each target distribution is mixture of two Gaussians as (8) = Ta (0
Pa. E,1
N2N (0 P:2, Σ.²) (i 1,2,3) where the mixing proportions Ta 0.7, Vi, n2 0.3, Vi.
the means An = [4, -4 1112 = [0,0.5]" Pm1 = [-4,4] 4 J222 = [0.5, 0 and Pml =
[-3, -3 P32 = [o, 0j* and the common covariance matrix Σ;
0.5
0
, = 1,2,3
0 0.5
and = 1, 2. It can be seen from Figure g that there is a common high-densio region spreading

---

around the origin. The fifty particles are drawn randomly in the space, and the initialization is retained across experiments for a fair comparison.
Figure 3: Sampling from three mixtures of two Gaussian distributions with a joint high-likelihool region. We rum MOO-SVGD (top) and MT-SGD (bottom) to update the initialized particles (left-most figures) until convergence using Adam optimizer [[]. While MOO-SVGD transports the initialized particles scattering on the distributions, MT-SGD perfectly drives them to diversify in the region of interest.
Figure 3 shows the updated particles by MOO-SVGD and MT-SGD at selected iterations, we observe that the particles from MOO-SVGD spread out and tend to characterize all the modes, some of them even scattered along trajectories due to the conflict in optimizing multiple objectives. By contrast, our method is able to find and cover the common high density region among target distributions with well-distribtted particles, which illustrates the basic principles of MT-SGD, Additionally, the 1, .000th step, the training time for ours is 0.23 min, whereas that for MOO-SVGD is 1.63 min. The reason is that MOO-SVGD requires solving an independent quadratic programming problem for each particle at each step.
4.1.2 Multi-objective Optimization
The previous experiment illustrates that MT-SGD can used to sample from multiple target distributions, next test method on the other low-dimensional multi-objectives OP from [29] In particular, we use the two objectives ZDT3, whose Pareto front consists of non-contiguous convex parts, show method simultancously minimizes both objective functions. Graphically, the simulation results from Figure 4 show the difference in the convergence behaviors between MOO-SVGD and MT-SGD: the solution set achieved by MOO-SVGD covers the entire Pareto front, while ours distributes and diversifies on the three middle curves (mostly concentrated the middle curve) which are the Pareto common having low values for two objective functions in ZDT3.
35
5.0
-.2
Figure 4: Solutions obtained by MOO-SVGD (mid) and MT-SGD (right) on ZDT3 problem after 10,000 steps, with blue points representing particles and red curves indicating the Pareto front. As expected, from initialized particles (left), MOO-SVGD's solution set widely distributes on the whole Pareto front while the one of MT-SGD concentrates around middle curves (mostly the middle one).

---

4.2 Experiments on Real Datasets
4.2.1 Experiments on Multi-Pashhon-MMMliiMMNNS"" Datasets
We apply the proposed MT-SGD method on multi-task learning, following Algorithm 2 Our method is validated on different benchmark datasets: (i) Sulli-Fasshoon-MNISTTT [23]. (ii) Mulii-MNSST, and (iii) Multi-Fashion Each of them consists of 120,000 training and 20,000 testing images generated from MNIST 12] and FashionMNIST 277 by overlaying an image on top of another: one in the top-left corner and one in the bottom-right corner. Lenet [2] (22,350 params) is employed as the backbone architecture and trained for 100 epochs with SGD in this experimental setup.
96.2
9299
94.3
90
94.4
0291
BAB
3350
85.5
84.6
0391
02
MODA
Fawels
W00-9VG3
MT-4RD
MOSA
Fawels NTT
NOD-SVI0
VI--50
$4.4
93.2
93.0
92.1
3:2-
90.2
33-
3.3
//ear sca otfociion
MODA
Farcio MTL
Moo SV9D
Unear Staloriann
MIDA
Fonds NTL
Noo SYGo
360
15-
84.1
81.9
83.7
833.3
82.8
3.25
4.2
(left), Multi-MNIST (mid), and Multi-Fashion (right). We report the ensemble accuracy (higher is better) and the Brier (lower is better) over independent runs, as well as the standard deviation (the error bars and shaded areas the figures).
Figure 5: Results on
Baselines: In multi-tase task experiments, the introduced MT-SGD is compared with state-of of the-hht-t baselines including MGDA (24), Pareto MTL [13], MOO-SVGD (15). We note that to reproduce results for these baselines, we either use the author's official implementation released on GitHub or ask the authors for their codes. For MOO-SVGD and Pareto MTL, the reported result is from the ensemble prediction of five particle models. Additionally, for linear scalarization and MGDA, we train five particle models independently with different initializations and then ensemble these models.
Evaluation metries: We compare MT-SGD against baselines regarding both average accuracy and predictive uncertainty. Besides the commonly used accuracy metric, we measure the quality and diversity of the particle models by relying on other popular Bayesian metrics: Brier score 3 20 and expected calibration error (ECE) 19
observe that MT-SGD consistently improves model performance across all tasks in both accuracy and Brier score by large margins, compared to existing techniques in the literature. The network trained using linear scalarization, expected, produces inferior ensemble results while utilizing MOO techniques helps yield better performances. Overall, our proposed method surpasses the second-best bascline by at least 1% accuracy in any experiment. Furthermore, Table provides
From Figure

---

comparison between these methods in terms of expected calibration error, in which MT-SGD also consistently provides the lowest expected calibration error, illustrating our method's ability obtain well-calibrated models (the accuracy is closely approximated by the produced confidence score). It is also worth noting that while Pareto MTL has higher accuracy, MOO-SVGD produces slightly better calibration estimation.
Table 1: Expected calibration error (%) (num bin 10) on Multi-MNIST, Multi-Fashion and Multi-Pashion+MNISTT datasets over 3 runs. We use the bold font to highlight the best results
Dataset
Task
Linear scalarization
MGDA
Pareto MTI.
MOO-SVGI
MT-SGD
Top left
21.33 0.83
19.91 0.26
9.44 + 0.65
9.47 = 0.89
4.65 + 0.11
Multi-Fashion-MNNSS
Bottom right
17.76
16.29 1.35
0.46
0.49
3.17 0.20
Top left
17.37 0.62
15.29 0.49
5.45 0.85
5.37 0.51
3.28 + 0.20
Multi-MNIST
Bottom right
18.09 1.11
16.87 = 0.67
7.34 +: 1.08
6.74 = 0.50
4,00 = 0.19
Top left
15.86 1.20
14.48 0,95
8.55 + 0.69
5.48 = 0.53
3.80 + 0.38
Multi-Fashion
Bottom right
15.98 = 1.32
14.70 = 1.63
9.01 +: 1.77
6.11 = 0.54
4.47 = 0.21
4.2.2 Experiment CelebA Dataset
In this experiment, we verify the significance of MT-SGD larger neural network: Resnet18 [10), which consists of 1 11.4M parameters. We take the first 10 binary classification tasks and randomly select subset of 40k images from the CelebA dataset [16]. Note that in this experiment, we consider Single task, which 10 models are trained separately and serves as a strong baseline.
Table 2: Results on CelebA dataset, regarding accuracy and expected calibration error. For the full names of the tasks, please refer to our supplementary material. While MGDA trains single model only to adapt on all tasks, reported performance of MOO-SVOL and MT-SGD is the ensemble results from five particle models.
Method
5S
AE
Att
BUE
Bald
Bangs
BL
BN
BlaH
BloH
Average
Single task
91.8
84.6
80.3
81.9
98.8
94.8
85.8
81.3
89.6
94.2
88.3
MGDA
91.8
84.0
79.0
81.3
98.6
94.6
83.6
81.6
89.8
93.8
87.8 8
Acc (%)
MOO-SYGD
92.3
84.2
78.9
81.2
98.9
94.5
86.4
80.0
90.8
94.8
88.2
MT-SGD
92.6
84.8
80.3
99.1
95.2
86.3
82.6
91.1
95.0
Single task
3.3
2.4
4.4
3.9
0.7
1.6
5.7
6.5
3.1
1.1
3.3
MGDA
1.4
1.1
3.5
7.3
0.3
1.8
6.9
5.4
2.1
1.2
3.1
ECE (%)
MOO-SYGD
2.8
1.9
3.1
5.6
0.3
0.5
4.7
3.3
1.3
1.3
2.5
MT-SGD
1.2
1.4
1.7
2.3
0.6
1.7
6.8
1.2
0.9
2.0
The performance comparison of the mentioned models CelebA experiment is shown in Table 2 clearly seen from the upper part of the table, MT-SGD performs best in all tasks, except in BL, where MOO-SVGD is slightly better (86.4% vs 86. 86.3%). Moreover, our method matches or beats Single task the second-be -bett baseline in all tasks. Regarding the well-calibrated uncertainty estimates, ensemble learning methods exhibit better results. In particular, MT-SGD and (VO-SVVD provide the best calibration performances, which are 2% and 2.5%, respectively, which emphasizes the importance of efficient ensemble learning for enhanced calibration.
As
5 Conclusion
In this paper, we propose Stochastic Multiple Target Sampling Gradient Descent (MT-SGD), allowing us to sample the particles from the joint libeliikd of multiple target distributions. Our MT-SGD is theoretically guaranteed to simultancously reduce the divergences to the target distributions. Inter- estingly. the asymptotic analysis of our MT-SGD reduces exactly to the multi-objective optimization. We conduct comprehensive experiments to demonstrate that by driving the particles to the Pareto common (the joint high-likelihood of multiple target distributions), our MT-SGD can outperform the baselines on the ensemble accuracy and the well-known Bayesian metrics such as the expected calibration error and the Brier score.
10

---

References
[1] Mark Anderson and Carlos GómezzRodriguez. A modest Pareto optimisation analysis of dependency parsers in 2021. In Proceedings of the 17th International Conference on Parsing Technologies and the IWPT 2 Shared Task on Parsing into Enhanced Universal Dependencies (IWPT 2021), pages 119-130, Online, August 2021. Association for Computational Linguistics.
[2] Christopher M. Bishop. Pattern Recognition and Machine Learning (Information Science and Statistics), Springer-Verlag. Berlin, Heidelberg, 2006.
[3] Glenn W Brier et al. Verification forecasts expressed in terms of probability. Monthly weather review, 78(1):1-3, 1950.
[4] Zhengyu Chen, Jixie Ge, Heshen Zhan, Siteng Huang, and Donglin Wang. Pareto sufe-supe training for few-shot learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 33663-13672, 2021.
[5] Sumanth Chennupati, Mohammad Mahdi Kamani, Zhongwei Cheng, and Lin Chen. Adaptive distillation: Aggregating knowledge from multiple paths for efficient distillation. arXiv preprint arXiv:2110.09674, 2021.
[6] Philip Dawid. The well-calibrated bayesian. Journal of the American Statistical Association, 77(379):005-610, 1982.
[7] Jean-Antoine Désidéri. Multiple-gradien descent algorithm (mgda) for multiobjective opti- mization. Comptes Rendus Mathematique, 350(5-6):313-313 2012.
[8] Shangchen Du, Shan You, Xiaojie Li, Jianlong Wu, Fei Wang, Chen Qian, and Changshui Zhang. Agree to disagree: Adaptive ensemble knowledge distillation in gradient space. advances in neural information processing systems, 14:12245-1222 2020.
[9] Shaona Ghosh, Chris Lovell, and Steve R Gunn. Towards pareto descent directions in sampling experts for multiple tasks in on-line learning paradigm. In 2013 AAAI Spring Symposium Series, 2013.
[10] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[11] Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. Proceedings of the International Conference on Learning Representations, 2014.
[12] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278 2324, 1998.
[13] Xi Lin, Hui-Ling Zhen, Zhenhua Li, Qing-Fu Zhang, and Sam Kwong. Pareto multi-task learning. Advances in neural information processing systems, 32, 2019.
[14] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.
[15] Xingchao Liu, Xin Tong, and Qiang Liu. Profiling pareto front with multi-objective stein variational gradient descent. Advances in Neural Information Processing Systems, 34, 2021.
[16] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision, pages 3738, 2015.
[17] Debabrata Mahapatra and Vaibhav Rajan. Multi-task learning with user preferences: Gradient descent with controlled ascent in pareto optimization. In International Conference on Machine Learning, pages 6597-6607. PMLR, 2020.
[18] Kevin P. Murphy. Probabilistic Machine Learning: An introduction. MIT Press, 2022.
11

---

[19] Mahdi Pakdaman Nacini, Gregory Cooper, and Milos Hauskrecht, Obtaining well calibrated probabilities using bayesian binning. Twenty- Ninnt AAAI Conference on Artificial Intelligence. 2015.
[20] Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's 's uncertainty? evaluating predictive uncertainty under dataset shift. Advances in neural information processing systems. 32, 2019.
[21] Simone Parisi, Matteo Pirotta, Nicola Smacchia, Luca Bascetta, and Marcello Restelli. Policy gradient approaches for multi-objective sequential decision making. In 2014 International Joint Conference on Neural Networks (LJCNN), pages 2323-2333 IEEE, 2014.
[22] Matteo Pirotta and Marcello Restelli. Inverse reinforcement learning through policy gradient minimization. In Thirtieth AAAI Conference on Artificial Intelligence, 2016.
[23] Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. Dynamic routing between capsules. In
I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[24] Ozan Sener and Vladlen Koltun. Multi-task learning as multi-objective optimization. Advances in neural information processing systems, 31, 2018.
[25] Martin J. Wainwright and Michael I. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 1(1-2):1-305, 2008.
[26] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference machine learning (ICML-II), pages 681-688. Citeseer, 2011.
[27] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 0777477777, 2017.
(28) Feiyang Ye, Baijiong Lin, Zhixiong Yue, Pengxin Guo. Qiao Xiao, and Yu Zhang. Multi- objective meta learning. Advances in Neural Information Processing Systems, 34, 2021.
[29] Eckart Zitzler, Kalyanmoy Deb, and Lothar Thiele. Comparison of multiobjective evolutionary algorithms: Empirical results. Evolutionary computation, 8(2):173-11)) 195, 2000.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] In the supplementary material,
discuss the computational overhead of our proposed method, compared with related work.
(c) Did you discuss any potential negative societal impacts of your work? [No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] We include assumptions all theoretical statements.
(b) Did you include complete proofs of all theoretical results? [Yes] The main results and lemmas appear in the main paper and their complete proof is put in the supplementary material.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] They will be attached to the supplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] The training details are provided with the code, but the important details should be in the main paper.
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes] We run each experiment over a number of random seeds and report the average result along with the standard deviation.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] The detailed training environment will be attached in the supplementary material.
4. If you using existing assets (e.g., code, data, models) curating/releasing assets....
(a) If your work uses existing assets, did you cite the creators? [Yes] We cited the authors of the codes that are included in our implementations.
(b) Did you mention the license the assets? [No] All the assets we used are published for research purposes.
(c) Did you include any assets either in the supplemental material
URL? [Yes]
The tools used in our implementation are presented our supplementary material.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [No] The data we used open the public.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [No] The data we used are open to the public.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
13


Reviews:
[{"review_id": "yn0fQ1UXHHh", "reviewer": "Reviewer_yFUe", "summary": "This paper proposed a variant of SVGD for multi-task learning. In particular, instead of driving particles towards the high-density region of a single distribution like SVGD, the proposed MT-SGD drives the particles towards the joint region of all densities. The main novelty of this paper is to propose a weighted combination of the driving directions from each SVGD such that the combined direction is guaranteed to decrease all KL divergences between the current distribution and each target distribution. Additionally, like the connection of SVGD to SGD, the author showed the connection of MT-SGD to MOO. Empirically, the proposed MT-SGD achieves better results compared to baselines consistently. ", "questions": "**Clarity**:\n1. From my understanding, the main novelty is the objective in lemma 1. Is there an intuition on why $w^TUw$ is a reasonable objective?. It seems that this objective is very similar to the one used in MOO (Eq.3). But instead replacing the derivative of $L$ with the functional derivative of $D_{KL}$. Is this true? In that case, personally, I think it would be better if MT-SGD is introduced from the MOO point of view. In this way, it is also easier for the reader to understand intuition. Also, can MT-SGD be derived based on the Pareto stationary KKT condition?\n\n2. In Eq.1, what does the comma mean? I understand this is a direct adaptation of MOO, but can you explain this minimization, since the typical minimization is targeted at a single value, but here, we have a list of $D_{KL}$. \n\n3. In line 49, do you mean empirical distributions?\n\n4. In line 107, is the optimal transformation $T$ for the current step?\n\n5. In line 112, should it be $O(\\epsilon^2)$?\n\n6. Just curious, can the objective $w^TUw$ in lemma 1 be translated to an optimization problem of the weighted combination of $D_{KL}$? \n\n7. In line 175, I suggest using a different subscript for $w$, since $t$ is already reserved for the number of SVGD steps in algorithm 1. \n8. For the captions of Figure 4, it should be red curve. \n\n**SVGD-related**:\n1. One of my concerns about using SVGD is its mode collapse problem, which has been extensively reported [1]. This means without proper tuning of the repulsive force, the particles tend to collapse to a single point in high dimensions. Have you noticed this behaviour? If not, why?\n2. In line 143, the author mentioned when $\\sigma \\rightarrow \\infty$, $U$ reduces to the inner product of the score function. But it seems that the update of MT-SGD, in this case, does not recover the MOO? Since the kernel matrix in SVGD is now a matrix with value 1, which takes the gradient of other particles into consideration.\n\n\n[1]: Gong, Wenbo, Yingzhen Li, and José Miguel Hernández-Lobato. \"Sliced kernelized Stein discrepancy.\" arXiv preprint arXiv:2006.16531 (2020).\n", "limitations": "This work does not have any potential negative societal impact. In terms of the potential limitation, the author argues that the proposed MT-SGD tries to find the Pareto common, whereas the other baselines try to find the Pareto front. I am curious under what scenario Pareto front is preferred compared to Pareto common? ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**:\nThis paper adopts the framework of MOO and extends to SVGD. The basic idea is very similar to MOO, the novelty comes from the incorporation of SVGD. The objective for searching $w^*$ can also be viewed as an adaptation of the MOO objective. In summary, the proposed approach is novel to some extent, but as mentioned by the author, this is not a completely novel framework. In terms of significance, the proposed framework demonstrates clear advantages compared to baselines and seems to be a good contribution to Bayesian multi-tasking learning. \n\n**Weakness**:\nI have checked the proofs in the appendix, which seem to be correct. Thus, my main criticism is its presentation. I will elaborate more in the question section. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ZpUQy03IVbA", "reviewer": "Reviewer_6tPS", "summary": "In this work, the authors introduce Stochastic Multiple Target Sampling Gradient Descent (MT-SGD), an extension of Stein variational gradient descent (SVGD) that tries to approximately sample from multiple distributions *simultaneously*. The authors then demonstrate that MT-SGD is well suited for Bayesian multi-task learning.", "questions": "1. In section 2.2.1 in the appendix, there is an interesting phenomenon in Figure 2 where increasing the number of particles, $M$, seems to **deteriorate** the performance substantially.  A prior, I don't see why this should be the case. I would appreciate it if the authors could clear this up.\n2. In  section on Multi-task learning, the authors should state that they are using a non-informative prior for the parameters of the model. ", "limitations": "While on the checklist the authors stated that they talk about the limitations in the appendix, I don't see it anywhere. To me, the two biggest limitations are the scaling with the number of tasks, $K$, and the number of particles, $M$ (which is inherent to SVGD). I think the authors should mention it in the work, though it doesn't detract from the novelty of the work.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\nMT-SGD is a natural extension of SVGD that is theoretically motivated. Moreover, the authors were also able to demonstrate for RBF kernels, as the scale goes to infinity, MT-SGD reduces to multiple objectivation optimization which is amazing! Moreover, the computational complexity of the algorithm seems to be on the order of SVGD, with the main difference being solving the quadratic programming problem for finding the optimal combination of the gradients. While I'm not familiar with the SVGD and multi-task learning literature, this approach seems both highly novel and theoretically grounded. The experiments section is also compelling as well.\n\n# Weaknesses\n\nFirstly, I'm a little confused about the Multi-task learning formulation. Given the data, one would want to compute the following posterior distribution\n$$ p(\\alpha, \\beta^1, \\ldots, \\beta^K \\mid \\mathbb{D})  \\propto p(\\alpha) \\prod_{j=1}^K p(\\beta^j) \\prod_{i=1}^N p(y_{ij} \\mid x_i, \\alpha, \\beta^j) $$\nConditioned on the non-shared parts, $\\beta^1, \\ldots, \\beta^K$, the posterior for the shared parts is\n$$ p(\\alpha \\mid \\mathbb{D}, \\beta^1, \\ldots, \\beta^J) \\propto p(\\alpha) \\prod_{j=1}^K \\prod_{i=1}^N p(y_{ij} \\mid x_i, \\alpha, \\beta^j)  $$\nwhere we can see that this conditional posterior already attempts to find an $\\alpha$ that is good for all tasks simulatensouly. Moreover, just regular SVGD can be used for this. I don't see the advantage, nor the intuition, of formulating the problem as $K$ individual posteriors. I think it would strengthen the approach if they authors compared SVGD on the posterior I have outlined.\n\nNext, while the experiments section is great, it is **very** concerning that there are no error bars for all the experiments in section 4.2, especially given how close the accuracies are for each of these tasks. Thus, I think it is imperative the authors run the experiments over a number of realizations to make me more confident in their results.\n\nPlease let me know if I have misunderstood anything! I will also be happy to raise my score if the authors get back to me about my concerns.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ViKY-WVuwGb", "reviewer": "Reviewer_TqCR", "summary": "This paper proposes MT-SGD for multi-objective optimization problems, which leverages the idea of SVGD to generate diverse particles w.r.t. the joint high-likelihood region for all distributions. Experiments on both synthetic dataset and real dataset (Multi-Fashion+Multi-MNIST and CelebA) demonstrate the effectiveness of MT-SGD compared to MGDA, Pareto MTL and MOO-SVGD.", "questions": "None", "limitations": "None", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n+ Writing is good and the paper is easy to follow\n+ Finding diverse particles on the Pareto front is a well-motivated topic.\n+ The proposed method is novel and inspiring\n\nWeakness:\n- Lack of convergence analysis of MT-SGD. Theorem 1 & 2 proves that MT-SGD reduces every KL divergence simultaneously, but it is unclear how the variational distribution $q$ or the particle set $\\{\\theta_i\\}_{i=1}^M$ will be when it converged.\n- The equation between line 166 and line 167 seems weird. It should be $p(\\theta|D) \\propto p(D|\\theta) p(\\theta)$. Besides, how the prior $p(\\theta)$ is chosen in experiments?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept