Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Accelerated Multiple Wasserstein Gradient Flows for Multi-objective Distributional Optimization

Dai Hai Nguyen, Duc Dung Nguyen, Atsuyoshi Nakamura, Hiroshi Mamitsuka

TL;DR

The paper addresses multi-objective optimization over probability distributions in Wasserstein space and introduces an accelerated variant of MWGraD, called A-MWGraD, which leverages damped Hamiltonian dynamics to achieve faster continuous-time convergence. It provides rigorous results showing $M(\rho_t) = O(1/t^2)$ for geodesically convex objectives and $M(\rho_t) = O(e^{-\sqrt{\beta}t})$ for $\beta$-strongly geodesically convex objectives, improving upon the $O(1/t)$ rate of MWGraD; a kernel-based discretization enables practical particle-based implementations. The authors develop both SVGD- and Blob-based discretizations, derive particle dynamics, and demonstrate empirically that A-MWGraD accelerates convergence and enhances sampling efficiency on multi-target sampling tasks and Bayesian multitask learning. The work advances scalable, multi-objective distributional optimization with theoretical guarantees and practical gains, while outlining avenues for future discrete-time analysis and broader applications.

Abstract

We study multi-objective optimization over probability distributions in Wasserstein space. Recently, Nguyen et al. (2025) introduced Multiple Wasserstein Gradient Descent (MWGraD) algorithm, which exploits the geometric structure of Wasserstein space to jointly optimize multiple objectives. Building on this approach, we propose an accelerated variant, A-MWGraD, inspired by Nesterov's acceleration. We analyze the continuous-time dynamics and establish convergence to weakly Pareto optimal points in probability space. Our theoretical results show that A-MWGraD achieves a convergence rate of O(1/t^2) for geodesically convex objectives and O(e^{-\sqrtβt}) for $β$-strongly geodesically convex objectives, improving upon the O(1/t) rate of MWGraD in the geodesically convex setting. We further introduce a practical kernel-based discretization for A-MWGraD and demonstrate through numerical experiments that it consistently outperforms MWGraD in convergence speed and sampling efficiency on multi-target sampling tasks.

Accelerated Multiple Wasserstein Gradient Flows for Multi-objective Distributional Optimization

TL;DR

The paper addresses multi-objective optimization over probability distributions in Wasserstein space and introduces an accelerated variant of MWGraD, called A-MWGraD, which leverages damped Hamiltonian dynamics to achieve faster continuous-time convergence. It provides rigorous results showing for geodesically convex objectives and for -strongly geodesically convex objectives, improving upon the rate of MWGraD; a kernel-based discretization enables practical particle-based implementations. The authors develop both SVGD- and Blob-based discretizations, derive particle dynamics, and demonstrate empirically that A-MWGraD accelerates convergence and enhances sampling efficiency on multi-target sampling tasks and Bayesian multitask learning. The work advances scalable, multi-objective distributional optimization with theoretical guarantees and practical gains, while outlining avenues for future discrete-time analysis and broader applications.

Abstract

We study multi-objective optimization over probability distributions in Wasserstein space. Recently, Nguyen et al. (2025) introduced Multiple Wasserstein Gradient Descent (MWGraD) algorithm, which exploits the geometric structure of Wasserstein space to jointly optimize multiple objectives. Building on this approach, we propose an accelerated variant, A-MWGraD, inspired by Nesterov's acceleration. We analyze the continuous-time dynamics and establish convergence to weakly Pareto optimal points in probability space. Our theoretical results show that A-MWGraD achieves a convergence rate of O(1/t^2) for geodesically convex objectives and O(e^{-\sqrtβt}) for -strongly geodesically convex objectives, improving upon the O(1/t) rate of MWGraD in the geodesically convex setting. We further introduce a practical kernel-based discretization for A-MWGraD and demonstrate through numerical experiments that it consistently outperforms MWGraD in convergence speed and sampling efficiency on multi-target sampling tasks.
Paper Structure (19 sections, 9 theorems, 108 equations, 3 figures, 1 table, 2 algorithms)

This paper contains 19 sections, 9 theorems, 108 equations, 3 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Gradient Flows on Euclidean space $\mathcal{X}$ and Their Acceleration
  4. Gradient Flows on Probability Spaces $\mathcal{P}_{2}(\mathcal{X})$ and Their Acceleration
  5. Multiple Wasserstein Gradient Descent Algorithm for MODO
  6. Multiple Wasserstein Gradient Flow and Its Acceleration
  7. Convergence Analysis of MWGraD flow
  8. Accelerated MWGraD Flow and Its Convergence Analysis
  9. Particle Implementation of A-MWGraD Flow
  10. Numerical Experiments
  11. Experiments on toy examples
  12. Experiments on Bayesian multi-task learning
  13. Conclusions
  14. Detailed Proof of Lemma \ref{['lemma:overwholespace']}
  15. Detailed Proof of Theorem \ref{['thm:MWGraD_convergencerate']}
  16. ...and 4 more sections

Key Result

Lemma 3.2

Let $\rho_0\in\mathcal{P}_2(\mathcal{X})$ and $\rho\in \Omega_{\textbf{F}}(\textbf{F}(\rho_0))$. Define Then

Figures (3)

  • Figure 1: Convergence comparison: (a) SVGD-based variants, (b) Blob-based variants. The plot show mean and standard deviation of GradNorm over 1000 iterations with three different step sizes of $\eta=0.001, 0.005, 0.01$.
  • Figure 2: Evaluation of the average test accuracies over 40000 training iterations on datasets: (a) Multi-MNIST, (b) Multi-Fashion, and (c) Multi-Fashion-MNIST.
  • Figure 3: Sampling from multiple target distributions, where each target is a mixture of two Gaussians. These targets have a joint high-density region around the origin. Initially, 50 particles are sampled from the standard distribution, and then updated by MWGraD variants, A-MWGraD variants, and MOO-SVGD. While MOO-SVGD tends to scatter particles across all the modes (see the last row), MWGraD and A-MWGraD variants tend to move particles towards the joint high-density region. Furthermore, A-MWGraD variants converge faster than MWGraD variants.

Theorems & Definitions (21)

  • Definition 2.1: Metric Tensor
  • Definition 2.2: Gradient Flow in Probability Spaces
  • Definition 2.3
  • Definition 2.4
  • Lemma 3.2
  • Theorem 3.4
  • Theorem 3.5
  • proof
  • Lemma 2.1
  • Lemma 2.2
  • ...and 11 more