---
title: 'MuSGD: Multi-Target Sampling Gradient Descent'
url: https://www.emergentmind.com/topics/musgd-optimizer
type: topic
---

# MuSGD: Multi-Target Sampling Gradient Descent

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 [2206.01934].

## 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) $\mathcal{H}_k^d$:
$$
\psi_k(x) = \mathbb{E}_{y \sim q}\left[ k(y, x) \nabla_y \log\pi_k(y) + \nabla_y k(y, x) \right]
$$
To ensure descent for all KL objectives, MuSGD finds an optimal convex weight vector $w^* = (w_1, \dots, w_K) \in \Delta^{K-1}$ by solving the quadratic program:
$$
w^* = \arg\min_{w \geq 0,\, \sum w_i = 1} w^T U w, \quad U_{ij} = \langle \psi_i, \psi_j \rangle_{\mathcal{H}_k^d}
$$
Construct the composite descent direction:
$$
\phi^*(x) = \sum_{k=1}^K w_k^* \psi_k(x)
$$
Within the mean-field limit, the particle SDE is governed by:
$$
\frac{d}{dt} X_i(t) = \phi^*(X_i(t)) = \sum_{k=1}^K w_k^*\,\mathbb{E}_{y \sim q}\bigl[ k(y, X_i) \nabla\log\pi_k(y) + \nabla_y k(y, X_i) \bigr]
$$

### Discrete-Time MuSGD Update

Particles $\{x_i\}_{i=1}^M$ represent the empirical distribution $q_t$. The discrete update executes:
1. For all $k$ and $i$, compute
   $$
   \psi_k(x_i) \approx \frac{1}{M} \sum_{j=1}^M \bigl[ k(x_j, x_i) \nabla\log\pi_k(x_j) + \nabla_{x_j} k(x_j, x_i) \bigr ]
   $$
2. Build Gram matrix $U$ using Monte Carlo inner products.
3. Solve $w^{(t)} = \arg\min_{w \in \Delta^{K-1}} w^T U w$ (small QP).
4. Compute composite direction $\phi^*(x_i)$ as above.
5. Update each particle:
   $$
   x_i^{(t+1)} = x_i^{(t)} + \eta_t \phi^*(x_i)
   $$
Expanded form:
$$
x_i^{(t+1)} = x_i^{(t)} + \eta_t \sum_{k=1}^K w_k^{(t)} \frac{1}{M} \sum_{j=1}^M \left[ k(x_j, x_i) \nabla\log\pi_k(x_j) + \nabla_{x_j} k(x_j, x_i) \right]
$$

## 3. Theoretical Properties and Connections

In the limit of infinite RBF kernel bandwidth (or a single particle), the kernel repulsive term is eliminated and $k(x,y) \rightarrow 1$, yielding:
$$
\psi_k(x) \to \nabla\log\pi_k(x), \quad \phi^*(x) \to \sum_{k=1}^K w_k^* \nabla\log\pi_k(x)
$$
If $\pi_k \propto e^{-\ell_k}$, this matches the multi-gradient descent (MGDA) direction. The paper proves that as $M \to \infty$, kernel bandwidth $\sigma \to \infty$, and step size $\eta \to 0$, MuSGD exactly recovers the classical MGDA update [2206.01934].

Under standard smoothness and Lipschitz continuity of $\nabla\log\pi_k$ and the kernel, both the continuous-time flow and discrete MuSGD trajectories converge as $\eta \to 0$. All KL divergences decrease at each iteration:
$$
D_{KL}(q^{[T]} \| \pi_k) = D_{KL}(q \| \pi_k) - \epsilon \|\phi^*\|_{\mathcal{H}_k^d}^2 + O(\epsilon^2), \quad \forall\, k
$$

## 4. Algorithmic Workflow and Computational Complexity

The standard MuSGD pseudocode is:
```
Input:   Unnormalized densities {π_k}_{k=1}^K, kernel k,
         M particles {x_i}₁ᴹ, stepsizes {η_t}, #iterations T
Output:  Particles approximating the joint high–density region

for t = 0 ... T-1 do
    For each k = 1...K, compute ψ_k at particles:
        ∀i, ψ_k(x_i) ← (1/M) ∑_{j=1}^M [k(x_j,x_i)∇log π_k(x_j) + ∇_{x_j} k(x_j,x_i)]
    Build U∈ℝ^{K×K}, U_{ij}← ⟨ψ_i,ψ_j⟩ via the MC formula
    Solve w^{(t)}=argmin_{w∈Δ^{K-1}} wᵀ U w (QP on simplex)
    φ^*(x_i) ← ∑_{k=1}^K w_k^{(t)} ψ_k(x_i)
    x_i ← x_i + η_t φ^*(x_i), for i=1…M
end for
return {x_i}
```
Complexity per iteration:
- Kernel computations and Stein terms: $O(K M^2 d)$
- Building $U$: $O(K^2 M^2 d)$
- QP on simplex: $O(K^3)$

Memory footprint includes $M \times d$ particles, $K$ Stein fields at $M$ points, and the $K \times K$ matrix $U$.

## 5. Empirical Performance and Applications

### Sampling Accuracy

On synthetic problems (e.g., mixtures of Gaussians in $\mathbb{R}^2$), 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 [2206.01934].

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

Source: https://www.emergentmind.com/topics/musgd-optimizer