MuSGD: Multi-Target Sampling Gradient Descent

• MuSGD is a sampling algorithm that employs composite Stein directions to simultaneously target multiple unnormalized distributions.
• It integrates RKHS-based Stein gradients with a convex QP to determine optimal weightings, ensuring descent for all KL divergences.
• Empirical results show that MuSGD outperforms traditional methods like MGDA and linear scalarization in accuracy and calibration on datasets such as CelebA and multi-MNIST.

MuSGD (Stochastic Multiple Target Sampling Gradient Descent) is a sampling algorithm designed to simultaneously handle multiple unnormalized target distributions. It extends Stein Variational Gradient Descent (SVGD) into the multi-target/multi-objective domain by iteratively updating a population of particles via composite Stein directions. MuSGD is primarily intended for probabilistic inference and multi-task learning, with theoretical guarantees and empirical benefits over classical approaches such as linear scalarization and multi-gradient descent algorithms (Phan et al., 2022).

1. Mathematical Formulation

Given $K$ unnormalized target densities $\{\pi_k(x)\}_{k=1}^K$, with $x\in\mathbb{R}^d$, the goal is to construct a sequence of intermediate distributions $(q_0 \to q_1 \to \cdots \to q_L)$ that progressively move closer to the joint high-density region of all $\{\pi_k\}$. Each step is achieved by a push-forward update:

$$q_{t+1} = T_t\# q_t, \quad T_t(x) = x + \epsilon_t \phi_t(x),$$

where $\phi_t(x)$ is a transport field constructed to minimize all KL divergences $$[\mathrm{KL}(q \| \pi_1), \dots, \mathrm{KL}(q \| \pi_K)]$$ simultaneously, treating the task as a multi-objective problem over the space of densities. The optimization objective is:

$$\min_{q \in \mathcal{Q}} [\mathrm{KL}(q \| \pi_1), \dots, \mathrm{KL}(q \| \pi_K)]$$

2. Gradient Flow, Stein Directions, and Update Rule

Continuous-Time Gradient Flow

For each target $k$, define the Stein-variational direction in the reproducing kernel Hilbert space (RKHS) $\{\pi_k(x)\}_{k=1}^K$0:

$\{\pi_k(x)\}_{k=1}^K$1

To ensure descent for all KL objectives, MuSGD finds an optimal convex weight vector $\{\pi_k(x)\}_{k=1}^K$2 by solving the quadratic program:

$\{\pi_k(x)\}_{k=1}^K$3

Construct the composite descent direction:

$\{\pi_k(x)\}_{k=1}^K$4

Within the mean-field limit, the particle SDE is governed by:

$\{\pi_k(x)\}_{k=1}^K$5

Discrete-Time MuSGD Update

Particles $\{\pi_k(x)\}_{k=1}^K$6 represent the empirical distribution $\{\pi_k(x)\}_{k=1}^K$7. The discrete update executes:

1. For all $\{\pi_k(x)\}_{k=1}^K$8 and $\{\pi_k(x)\}_{k=1}^K$9, compute

$x\in\mathbb{R}^d$0

1. Build Gram matrix $x\in\mathbb{R}^d$1 using Monte Carlo inner products.
2. Solve $x\in\mathbb{R}^d$2 (small QP).
3. Compute composite direction $x\in\mathbb{R}^d$3 as above.
4. Update each particle:

$x\in\mathbb{R}^d$4

Expanded form:

$x\in\mathbb{R}^d$5

3. Theoretical Properties and Connections

In the limit of infinite RBF kernel bandwidth (or a single particle), the kernel repulsive term is eliminated and $x\in\mathbb{R}^d$6, yielding:

$x\in\mathbb{R}^d$7

If $x\in\mathbb{R}^d$8, this matches the multi-gradient descent (MGDA) direction. The paper proves that as $x\in\mathbb{R}^d$9, kernel bandwidth $(q_0 \to q_1 \to \cdots \to q_L)$0, and step size $(q_0 \to q_1 \to \cdots \to q_L)$1, MuSGD exactly recovers the classical MGDA update (Phan et al., 2022).

Under standard smoothness and Lipschitz continuity of $(q_0 \to q_1 \to \cdots \to q_L)$2 and the kernel, both the continuous-time flow and discrete MuSGD trajectories converge as $(q_0 \to q_1 \to \cdots \to q_L)$3. All KL divergences decrease at each iteration:

$(q_0 \to q_1 \to \cdots \to q_L)$4

4. Algorithmic Workflow and Computational Complexity

The standard MuSGD pseudocode is: $\{\pi_k\}$5 Complexity per iteration:

• Kernel computations and Stein terms: $(q_0 \to q_1 \to \cdots \to q_L)$5
• Building $(q_0 \to q_1 \to \cdots \to q_L)$6: $(q_0 \to q_1 \to \cdots \to q_L)$7
• QP on simplex: $(q_0 \to q_1 \to \cdots \to q_L)$8

Memory footprint includes $(q_0 \to q_1 \to \cdots \to q_L)$9 particles, $\{\pi_k\}$0 Stein fields at $\{\pi_k\}$1 points, and the $\{\pi_k\}$2 matrix $\{\pi_k\}$3.

5. Empirical Performance and Applications

Sampling Accuracy

On synthetic problems (e.g., mixtures of Gaussians in $\{\pi_k\}$4), MuSGD particles concentrate in the true joint high-density region, outperforming methods such as MOO-SVGD, which scatter across separate modes.

Multi-Task Learning

MuSGD has been evaluated on multi-MNIST, multi-FashionMNIST, CelebA (10 attributes), SARCOS regression datasets. Using architectures such as LeNet or ResNet-18, MuSGD alternates particle-based sampling for shared parameters via MT-SGD and task-specific parameters via SVGD. Metrics considered are ensemble accuracy, Brier score, and Expected Calibration Error (ECE).

Extracted empirical results:

• On CelebA, MuSGD attains highest mean accuracy (89.0% vs. 88.2% for MOO-SVGD) and lowest ECE (2.0% vs. 2.5%).
• On SARCOS regression, MuSGD yields the lowest RMSE for all outputs (0.0428 vs. 0.0515 for MOO-SVGD).
• MuSGD consistently outperforms single-task SGD, linear scalarization, MGDA, Pareto MTL, and MOO-SVGD in both accuracy (+1–2%) and calibration (lower ECE), across varied tasks (Phan et al., 2022).

6. Interpretation and Practical Guidelines

MuSGD generalizes SVGD to the multi-objective setting by:

• Combining multiple Stein directions through simplex QP solving.
• Retaining kernel-based repulsion for diversity among particles.
• Guaranteeing global theoretical descent for all KL objectives.
• Asymptotically coinciding with classical multi-gradient descent.
• Demonstrably enhanced joint sampling and downstream multi-task generalization under standard smoothness conditions.

A plausible implication is that MuSGD represents the first kernelized variant of multi-gradient descent, yielding improved empirical sampling efficiency and predictive calibration in multi-objective scenarios. The algorithm is readily adaptable to modern deep learning architectures, provided access to the Stein gradients and kernel functions.