MWGraD: Multi-Objective Wasserstein Descent

• MWGraD is a particle-based optimization method that leverages the Riemannian geometry of Wasserstein space for multi-objective distributional optimization.
• It computes projections of negative Wasserstein gradients onto the convex hull of functional gradients, ensuring Pareto-optimal updates for conflicting objectives.
• Kernel-based approximations and quadratic programming enable efficient particle updates, with accelerated variants like A-MWGraD offering improved convergence rates.

Multiple Wasserstein Gradient Descent (MWGraD) refers to a class of particle-based optimization algorithms for simultaneous minimization of multiple objective functionals over probability measures in Wasserstein space. Unlike classical multi-objective gradient methods in Euclidean space, MWGraD exploits the Riemannian geometry of the Wasserstein-2 space $\mathcal{P}_2(\mathcal{X})$, projecting negative Wasserstein gradients onto the convex hull of multiple functional gradients and providing a framework for multi-objective distributional optimization (MODO). MWGraD unifies concepts from optimal transport, multi-objective optimization, and statistical inference, enabling efficient sampling and inference in settings with potentially conflicting distributional objectives (Nguyen et al., 24 May 2025, Nguyen et al., 27 Jan 2026).

1. Mathematical Foundations and Problem Statement

The goal of multi-objective distributional optimization is to find a probability measure $\rho \in \mathcal{P}_2(\mathcal{X})$ that is Pareto-optimal with respect to $K$ smooth functionals $F_k:\mathcal{P}_2(\mathcal{X})\to \mathbb{R}$. The space $\mathcal{P}_2(\mathcal{X})$ is equipped with the 2-Wasserstein metric

$$\mathcal{W}_2^2(\rho,\rho') = \inf_{\gamma \in \Pi(\rho, \rho')}\int \|x - y\|^2 \gamma(dx, dy),$$

where $\Pi(\rho, \rho')$ denotes couplings of $\rho$ and $\rho'$. Pareto optimality requires that no alternative $\rho'$ strictly improves all $F_k$.

The Wasserstein geometry induces gradient flows for each $F_k$ via the continuity equation: $$\frac{d\rho_t}{dt} = -\operatorname{div}(\rho_t v_{F_k})$$ with $$v_{F_k}(x) = -\nabla_x \left[\delta_\rho F_k\right](x)$$, where $\delta_\rho F_k$ is the first variation of $F_k$ at $\rho$ (Nguyen et al., 24 May 2025, Nguyen et al., 27 Jan 2026).

2. Core MWGraD Algorithm: Continuous and Discrete Time

Continuous-Time Formulation

The continuous-time MWGraD flow seeks a velocity field $v^*(x)$ such that the induced flow decreases all $F_k$ as much as possible. It does so by projecting the origin onto the convex hull of gradients in the cotangent space: $$\dot{\rho}_t + \nabla \cdot (\rho_t \nabla \Phi_t) = 0,\qquad \Phi_t + \operatorname{proj}_{\mathcal{C}(\rho_t),\rho_t}[0]=0,$$ where $$\mathcal{C}(\rho) = \operatorname{conv}\{\delta_\rho F_k[\rho]\}$$, and $\operatorname{proj}_{\mathcal{K}, \rho}[f]$ denotes the metric projection in the cotangent space $L^2(\rho)$ (Nguyen et al., 27 Jan 2026).

Discrete-Time Particle MWGraD

At each iteration $n$, MWGraD approximates gradients using a set of $m$ particles $\{x^{(n)}_i\}_{i=1}^m$ sampled from the current empirical measure. The algorithm performs:

1. Computation of each functional’s Wasserstein gradient $\Delta_k^{(n)}(x_i)$ via kernel methods (SVGD or Blob).
2. Solution of a convex quadratic program for $w^{(n)} \in \Delta^K$:

$$w^{(n)} = \arg\min_{w \in \Delta^K} \frac12 \int \left\| \sum_{k=1}^K w_k \Delta_k^{(n)} \right\|^2 d\rho_n.$$

1. Formation of the combined descent vector $v_n(x) = \sum_{k=1}^K w_k^{(n)} \Delta_k^{(n)}(x)$.
2. Particle update using a forward Euler step

$x^{(n+1)}_i = x^{(n)}_i - \eta v_n(x^{(n)}_i).$

This mirrors the continuous-time projected-gradient flow and can be efficiently implemented when $K$ is moderate (Nguyen et al., 27 Jan 2026, Nguyen et al., 24 May 2025).

3. Geometric and Algorithmic Structure

The theoretical justification of this approach relies on the Riemannian structure of $(\mathcal{P}_2, \mathcal{W}_2)$:

• The tangent space at $\rho$ can be identified as $-\operatorname{div}(\rho u)$ for $u: \mathcal{X} \to \mathbb{R}^d$.
• The inner product is $\int \langle \nabla u_1, \nabla u_2 \rangle d\rho$.
• The multi-objective descent vector is obtained by minimizing the $L^2(\rho)$-norm of convex combinations of the individual gradients.

Algorithmically, MWGraD generalizes the Multiple Gradient Descent Algorithm (MGDA) to probability measures, ensuring descent in all objectives through min-norm convex aggregation (Nguyen et al., 24 May 2025).

4. Convergence Theory and Limitations

Under geodesic convexity in $\mathcal{W}_2$, the multi-objective “merit” function

$$\mathcal{M}(\rho) = \sup_{q \in \mathcal{P}_2} \min_{k \in [K]} \{ F_k(\rho) - F_k(q) \}$$

decays as $O(1/t)$ for the continuous MWGraD flow. Specifically, if all $F_k$ are geodesically convex and sublevel sets are bounded (with Wasserstein diameter $R$), then $\mathcal{M}(\rho_t) \leq R/(2t)$. No inertial effect or accelerated convergence is present in the original MWGraD; rates are limited by the underlying geometry and the method of approximation of log-density and gradients (Nguyen et al., 27 Jan 2026).

A-MWGraD, an accelerated variant inspired by Nesterov’s momentum, introduces auxiliary momentum fields and achieves faster convergence rates—$O(1/t^2)$ for geodesically convex objectives, and exponential decay under strong convexity. The absence of such mechanisms is a primary limitation of the unaccelerated MWGraD, especially in high-precision or high-dimensional regimes (Nguyen et al., 27 Jan 2026).

5. Relation to Other Wasserstein-based Methods

The MWGraD paradigm extends and interacts with several streams in the literature:

Method (Reference)	Scope	Gradient Type
Bures-Wasserstein GD (Chewi et al., 2020)	Gaussian barycenters	Single-objective
Product-space MWGraD (Chen et al., 31 Oct 2025)	Coupled distributional evolution	Two-marginal (opposite-flux)
Classical SVGD	Particle inference	Euclidean functional kernel
MWGraD (Nguyen et al., 24 May 2025, Nguyen et al., 27 Jan 2026)	Multi-objective over $\mathcal{P}_2$	Convex-hull Wasserstein
A-MWGraD (Nguyen et al., 27 Jan 2026)	Accelerated multi-objective	Momentum-augmented Wasserstein

In contrast to SVGD-based multi-objective methods, MWGraD couples updates via the geometry of the Wasserstein space. In product-space Wasserstein gradient flows, as described in (Chen et al., 31 Oct 2025), an “equal and opposite flux” structure emerges when coupling two marginals via relative entropy, relevant in control-theoretic applications.

6. Practical Implementations and Applications

In practical settings, MWGraD employs kernel-based approximations for the Wasserstein gradients, using either SVGD-style or mass-transport (“Blob”) methods. The quadratic optimization over $K$ weights is efficiently tractable for moderate $K$, and the overall computational complexity is dominated by kernel operations and, when used, Sinkhorn or entropic-OT solvers (Nguyen et al., 27 Jan 2026, Nguyen et al., 24 May 2025).

MWGraD has been evaluated on tasks including:

• Mixture-of-Gaussians synthetic sampling, demonstrating efficient particle concentration in joint high-density regions.
• Dissimilarity-based distributional matching (KL, JS), where MWGraD outperforms naive multi-objective extensions of SVGD.
• Multi-task Bayesian learning (e.g., Multi-MNIST), showing superior average test accuracy relative to MOO-SVGD and MT-SGD—suggesting efficiency for shared-parameter Bayesian models (Nguyen et al., 24 May 2025, Nguyen et al., 27 Jan 2026).

Kernel hyperparameters, particle ensemble size, and step sizes are empirically tuned; choices affect both accuracy and computational cost.

7. Open Challenges and Future Directions

The principal limitations of MWGraD stem from:

• The $O(1/t)$ convergence rate in the absence of acceleration.
• The need to efficiently approximate gradients and log-densities, particularly in high dimensions.
• Computational overhead increasing with $K$ for solving the min-norm quadratic program in aggregation.

Proposed advancements include momentum-based acceleration (A-MWGraD), adaptive step sizes, second-order schemes, and distributed particle methods. Potential applications extend to fairness-aware sampling, multi-objective generative modeling, and decentralized control of interacting particle systems.

A plausible implication is that the geometric projection approach of MWGraD and its close variants will remain central in scaling distributional multi-objective optimization and probabilistic inference for large-scale, multi-criteria learning problems (Nguyen et al., 27 Jan 2026, Nguyen et al., 24 May 2025, Chen et al., 31 Oct 2025).