Research map / Optimization / Nonconvex and smooth optimization
Mirror and Preconditioned Gradient Descent in Wasserstein Space
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Lifts mirror descent and preconditioned gradient descent to the Wasserstein-2 space via L2(mu) Bregman divergences, proving convergence under relative smoothness/convexity along specific curves, with experiments on ill-conditioned Gaussian/interaction-energy tasks and single-cell perturbation prediction.
Problem
Optimizing functionals over the Wasserstein space with explicit schemes that adapt to the geometry of the objective, since standard Wasserstein gradient descent is slow on ill-conditioned problems and existing discrete-time analyses are limited to specific objectives such as KL.
Approach
Generalizes mirror descent and preconditioned gradient descent to P2(R^d) by replacing the squared L2(mu) cost with Bregman divergences of pushforward-compatible potentials (or translation-invariant costs), and proves convergence via relative smoothness and convexity assessed along acceleration-free curves S + t(T-S).
Key results
- On 20 random ill-conditioned Gaussian targets in R^10, negative-entropy mirror descent converges almost as fast as ideally preconditioned forward-backward and much faster than forward-backward.
- For single-cell treatment prediction with Sinkhorn divergence, preconditioned GD converges in about 10 times fewer iterations than vanilla GD over 43 datasets.
- Preconditioned GD attains better optimal values of the objective than vanilla GD on average across the 43 treatment datasets.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use preconditioned gradient descent instead of vanilla gradient descent | Single-cell response prediction · 43 4i and scRNA-seq datasets | convergence speed | Helped |
| Use preconditioned gradient descent instead of vanilla gradient descent | Single-cell response prediction with Sinkhorn divergence · Single-cell datasets | iterations to convergence: 10 times less iterations | Helped |
| Use preconditioned gradient descent instead of vanilla gradient descent | Single-cell response prediction · 43 4i and scRNA-seq datasets | attained objective | Helped |
| Use negative entropy as the Bregman potential (NEM) instead of the Forward-Backward scheme | KL minimization for ill-conditioned Gaussian targets · 20 Gaussian targets in R^10 | convergence speed | Helped |
| Use negative entropy as the Bregman potential (NEM) instead of the ideally preconditioned Forward-Backward scheme | KL minimization for ill-conditioned Gaussian targets · 20 Gaussian targets in R^10 | convergence speed | No clear effect |
| Use interaction Bregman potentials that account for the ill-conditioned covariance matrix Σ | Mirror descent on interaction energies · Ill-conditioned interaction-energy task | convergence | Helped |
Surprising findings
- Negative-entropy mirror descent, which does not use an ideal preconditioner, is nearly as fast as the ideally preconditioned scheme on Gaussian targets.
- The Gaussian-restricted NEM scheme may fall outside the theoretical framework yet still shows benefits.
Limitations
- Relative smoothness and convexity assumptions may be hard to verify for general functionals such as Sliced-Wasserstein or Sinkhorn divergence.
- Preconditioned GD analysis is restricted to translation-invariant costs h, and general phi is left for future work.
- Implicit MD updates need Newton root-finding, which scales poorly with dimension and sample size.
Experimental setup
Datasets: 4i cell line data (34 treatments) scRNA sequencing cell line data (9 treatments)
Benchmarks: Gaussian target convergence (KL over time) single-cell treatment response prediction (4i, scRNAseq)
Metrics: KL divergence interaction energy value sliced-Wasserstein distance Sinkhorn divergence energy distance number of iterations to tolerance
Baselines: Wasserstein gradient descent forward-backward scheme SVGD
Builds on: mirror descent preconditioned gradient descent Wasserstein gradient descent JKO scheme Mirror Langevin
Closest papers by meaning
- On the Convergence of Projected Bures-Wasserstein Gradient Descent under Euclidean Strong Convexity (ICML 2024 · similarity 0.87)
- Mirror Descent Beyond Euclidean Stability: An Exponential Separation in Initialization Sensitivity (NeurIPS 2026 · similarity 0.87)
- A Unified Approach for Computing Wasserstein Barycenters of Discrete and Continuous Measures (NeurIPS 2026 · similarity 0.87)
- Spectral Preconditioning for Gradient Methods on Graded Non-convex Functions (ICML 2024 · similarity 0.87)
- Towards Understanding Gradient Dynamics of the Sliced-Wasserstein Distance via Critical Point Analysis (ICML 2025 · similarity 0.86)
- Hessian-guided Perturbed Wasserstein Gradient Flows for Escaping Saddle Points (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.