Technical Summary: Accelerated Multiple Wasserstein Gradient Flows for Multi-objective Distributional Optimization

1. Problem Statement

The paper addresses Multi-Objective Distributional Optimization (MODO), a problem defined over the space of probability distributions $\mathcal{P}_2(\mathcal{X})$ equipped with the Wasserstein-2 metric. The goal is to find a distribution $\rho$ that simultaneously minimizes a vector of $K$ objective functionals $\mathbf{F}(\rho) = (F_1(\rho), \dots, F_K(\rho))$.

This problem is a natural extension of classical multi-objective optimization (MOO) from Euclidean spaces to probability spaces. A primary motivating application is multi-target sampling, where the objective is to generate particles that approximate multiple unnormalized target distributions simultaneously. In this context, each objective $F_k$ typically represents the Kullback–Leibler (KL) divergence between the current distribution and a specific target distribution.

The challenge lies in the fact that different distributions may perform better on different objectives, necessitating the search for Pareto optimal or weakly Pareto optimal distributions rather than a single global minimum.

2. Methodology

2.1 Background: MWGraD

The work builds upon the Multiple Wasserstein Gradient Descent (MWGraD) algorithm (Nguyen et al., 2025a). MWGraD constructs a flow of distributions that decreases all objectives by:

1. Estimating the Wasserstein gradient $\nabla \delta F_k$ for each objective.
2. Aggregating these gradients into a single descent direction via a convex combination.
3. Updating particles along this direction.

In the continuous-time limit, MWGraD corresponds to a gradient flow where the velocity field is determined by projecting the zero vector onto the convex hull of the objective gradients.

2.2 Proposed Method: A-MWGraD

The authors propose A-MWGraD, an accelerated variant of MWGraD inspired by Nesterov's Accelerated Gradient (NAG) and interpreted through damped Hamiltonian dynamics.

Continuous-Time Formulation:
The accelerated flow is defined by the following system of equations for the distribution $\rho_t$ and a momentum potential $\Phi_t$:
$$\begin{aligned} \dot{\rho}_t + \nabla \cdot (\rho_t \nabla \Phi_t) &= 0 \\ \dot{\Phi}_t + \alpha_t \Phi_t + \frac{1}{2}\|\nabla \Phi_t\|_2^2 + \text{proj}_{\mathcal{C}(\rho_t), \rho_t}[0] &= 0 \end{aligned}$$
where $\mathcal{C}(\rho_t)$ is the convex hull of the first variations of the objectives, and $\alpha_t$ is a time-dependent damping coefficient.

Discrete-Time Implementation:
To make the method practical, the authors derive a particle-based discretization. Let $x_i^{(n)}$ be the position of the $i$-th particle and $v_i^{(n)}$ its velocity at iteration $n$. The update rules are:
$$\begin{aligned} x_i^{(n+1)} &= x_i^{(n)} + \sqrt{\eta} v_i^{(n)} \\ v_i^{(n+1)} &= \alpha_n v_i^{(n)} - \sqrt{\eta} \sum_{k=1}^K w_{n,k} \bar{\Delta}_k^{(n)}(x_i^{(n)}) \end{aligned}$$
where:

• $\eta$ is the step size.
• $\bar{\Delta}_k^{(n)}$ is an approximation of the Wasserstein gradient of the $k$-th objective, computed using Stein Variational Gradient Descent (SVGD) or Blob methods to handle empirical measures.
• $w_n$ is a weight vector minimizing the squared norm of the combined descent direction.
• $\alpha_n$ is the discrete damping parameter, chosen as $\frac{n-1}{n+2}$ for geodesically convex objectives or $\frac{1-\sqrt{\beta}\eta}{1+\sqrt{\beta}\eta}$ for $\beta$-strongly geodesically convex objectives.

3. Key Contributions

1. Continuous-Time Analysis of MWGraD: The authors formalize MWGraD as a continuous-time flow and establish its convergence to weakly Pareto optimal points. They define a merit function $M(\rho)$ and prove that MWGraD achieves a convergence rate of $O(1/t)$ for geodesically convex objectives.
2. Acceleration in Probability Spaces: The paper introduces A-MWGraD and proves it achieves faster convergence rates:
   • $O(1/t^2)$ for geodesically convex objectives.
   • $O(e^{-\sqrt{\beta}t})$ for $\beta$-strongly geodesically convex objectives.
   • This represents the first acceleration result for multi-objective optimization specifically within probability spaces.
3. Explicit Characterization in Gaussian Families: The authors provide an explicit reduction of the infinite-dimensional flows to finite-dimensional Ordinary Differential Equations (ODEs) governing the evolution of covariance matrices when the distributions are restricted to the Gaussian family and objectives are KL divergences.
4. Practical Discretization: A sampling-efficient particle implementation is developed, incorporating kernel-based approximations (SVGD and Blob) for gradient estimation and a strategy for computing optimal aggregation weights.

4. Experimental Results

The authors validate their theoretical findings through numerical experiments on synthetic and real-world datasets:

• Toy Examples (Multi-Target Sampling): Using mixtures of Gaussians, the study compares A-MWGraD against MWGraD and MOO-SVGD.
  • Convergence Speed: A-MWGraD variants consistently reach low values of the gradient norm (GradNorm) significantly faster than MWGraD. For instance, in one setting, A-MWGraD-Blob covered the high-density region in 50 iterations, whereas MWGraD-Blob required ~500.
  • Particle Distribution: While MOO-SVGD scattered particles across all modes, A-MWGraD and MWGraD successfully concentrated particles in the shared high-density region, with A-MWGraD doing so more rapidly.
• Bayesian Multi-Task Learning: Experiments were conducted on Multi-MNIST, Multi-Fashion, and Multi-Fashion-MNIST datasets.
  • Performance: A-MWGraD variants (both SVGD and Blob) achieved higher test accuracies faster than their unaccelerated counterparts.
  • Final Accuracy: In several tasks, A-MWGraD achieved the best ensemble accuracy (e.g., 96.4% on Multi-Fashion+MNIST Task 1), outperforming baselines like MOO-SVGD and MWGraD.

5. Significance and Claims

The paper claims that A-MWGraD successfully extends the principles of Nesterov acceleration to the domain of multi-objective optimization over probability distributions. Its significance lies in:

• Theoretical Improvement: Providing the first proof of accelerated convergence rates ($O(1/t^2)$ and exponential rates) for multi-objective distributional optimization, improving upon the $O(1/t)$ rate of standard MWGraD.
• Practical Efficiency: Demonstrating that acceleration leads to tangible improvements in sampling efficiency and convergence speed for multi-target sampling tasks, which are critical in Bayesian inference and generative modeling.

Limitations and Future Work:
The authors modestly acknowledge that:

• The discrete-time convergence rates for A-MWGraD remain unestablished (only continuous-time rates are proven).
• The theoretical analysis assumes exact computation of Wasserstein gradients, whereas the practical implementation relies on kernel approximations.
• Future work will address these gaps and explore broader applications beyond multi-task learning.