cs.LGNov 6, 2026

Fixed-point acceleration on Neural Wasserstein manifolds

Authors: Vitalii Aksenov, Sebastian Gutierrez Hernandez

Organizations: Weierstrass Institute for Applied Analysis and Stochastics (WIAS), Address: Anton-Wilhelm-Amo-Str. 39, 10117 Berlin, Germany. · Georgia Institute of Technology, Address: 686 Cherry St NW, GA 30332, Atlanta, United States

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

We propose a fixed-point acceleration method for optimization over probability measures, applicable to a variety of machine learning problems. We represent probability measures as pushforwards of a reference distribution through a neural network map and equip the parameter space with the induced pullback Wasserstein metric. This geometric structure allows natural gradient descent to be interpreted as a fixed-point iteration on a Riemannian manifold. We accelerate this iteration using Riemannian Anderson mixing method, which exploits historical residuals to capture approximate second-order information without forming Hessians. In the present parametric setting, the manifold is represented by Euclidean coordinates endowed with a point-dependent metric, simplifying retraction and vector transport. We evaluate the method on multimodal sampling, PDE-based inverse problems, aggregation dynamics, Bayesian logistic regression, and generative modeling. Across the tested tasks, Anderson mixing reduces the number of natural-gradient iterations and achieves competitive performance relative to SGD and Adam in terms of both loss-gradient evaluations and wall-clock time. These results suggest that fixed-point acceleration is an efficient tool for improving the computational efficiency of training on parametric Wasserstein methods.

Figures & tables

 Xk=[θi−θi−1], Rk=[ri−ri−1], i=k,…,k−mk+1.X_{k}=\left[\theta_{i}-\theta_{i-1}\right],\ R_{k}=\left[r_{i}-r_{i-1}\right],\ i=k,\ldots,k-m_{k}+1. 
Algorithm 1 Riemannian Anderson Mixing on Θ\Theta
Figure 1: Empirical assessment of Assumption 3.1 . Estimated minimal eigenvalue and condition number stay bounded away from 00 and ∞\infty in the unit ball around the fixed-point. See Appendix C.3 for detailed description.
Figure 2: Multimodal sampling. (a,c): samples from best trained models and contour lines of target potentials. (b,d): energy decay vs. iteration. RAM improves the convergence of NGD.
Figure 4
Figure 4: Bayesian logistic regression convergence across eight benchmark datasets. Each panel compares NGD with the top five RAM configurations from the hyperparameter sweep for 3 seeds. RAM reaches the low-energy regime in fewer iterations than NGD.
Figure 5: Convergence diagnostics for the aggregation-energy problem.
Figure 6: Performance in generative problems.
ProblemMMDk2⁡\operatorname{MMD_{k}^{2}}NGDRAMβ\beta
   β=1\beta=1best 
Checkerboard0.01850040005003
Two spirals0.0112000380022003
8 Gaussians0.001250017506003
Table 2: Number of iterations to convergence in the generative problems
Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

 α(k)=arg⁡min⁡v∈Rmk+1{∥Rkv∥2s.t. ∑i=0mkvi=1\alpha^{(k)}=\arg\min_{v\in\mathbb{R}^{m_{k}+1}}\left\{\begin{aligned} &\|R_{k}v\|_{2}\\ \text{s.t. }&\sum_{i=0}^{m_{k}}v_{i}=1\end{aligned}\right. (A.1.2)
Appendix
Algorithm 2 Anderson mixing in Rd\mathbb{R}^{d}
 Γk=arg⁡min⁡Γ∈Rmk∥rk−RkΓ∥xk2\Gamma_{k}=\arg\min_{\Gamma\in\mathbb{R}^{m_{k}}}\|r_{k}-R_{k}\Gamma\|^{2}_{x_{k}} (A.3.2)
Appendix
Algorithm 3 Riemannian Anderson Mixing Method
Problemnhiddenn_{\text{hidden}}DhiddenD_{\text{hidden}}Activationdt\mathrm{d}{t}hhΛ\LambdaNCGN_{\text{CG}}
SamplingStyblinsky–Tang22128128tanh⁡\operatorname{tanh}0.10.10.0010.00110−110^{-1}5050
Double banana221616tanh⁡\operatorname{tanh}0.10.10.0010.00110−110^{-1}5050
Aggregation 223232tanh⁡\operatorname{tanh}0.250.255×10−35\times 10^{-3}5×10−45\times 10^{-4}20
BLR 22128128tanh⁡\operatorname{tanh}0.250.251×10−31\times 10^{-3}5×10−25\times 10^{-2}20
GenerativeCheckerboard33128128swish⁡\operatorname{swish}0.10.10.0010.00110−310^{-3}5050
2 spirals33128128swish⁡\operatorname{swish}0.10.10.0010.00110−310^{-3}5050
Appendix
Table 3: Network architectures and optimizer parameters used in the experiments.
Figure 7: RAM and NGD compared with SGD and Adam on DB.
Figure 8: (a-c): Scatter plots for Adam after 10410^{4} iterations with varied learning rate hh . Mode collapse and numerical artifacts are clearly present in the samples. (d) Solution with RAM exhibits correct structure and relative weight of the modes.
DatasetNumber of featuresDataset size
banana25300
diabetes8768
german201000
image182086
ringnorm207400
splice602991
Appendix
Table 4: Dataset statistics for the Bayesian logistic regression benchmark.
 Accuracy ↑\uparrowLog-likelihood ↑\uparrow
DatasetNGDRAMVWGFJKO-ICNNSVGDNGDRAMVWGFJKO-ICNNSVGD
banana0.5650.5400.5500.5500.540-0.688-0.692-0.690-0.690-0.690
diabetes0.7530.7530.7860.7750.780-0.558-0.557-0.450-0.450-0.460
german0.7420.7500.6700.6700.650-0.552-0.561-0.590-0.600-0.600
image0.8210.8200.8660.8200.815-0.435-0.435-0.394-0.430-0.440
ringnorm0.7580.7640.7300.7400.740-0.520-0.509-0.500-0.500-0.500
Appendix
Table 5: Bayesian logistic-regression test performance. We report test accuracy and average test log-likelihood. RAM uses the representative configuration (m,β,λ)=(12,2.0,0.01)(m,\beta,\lambda)=(12,2.0,0.01) . VWGF and JKO-ICNN results are from [ 28 ] ; SVGD results are from [ 57 ] .
Figure 9: Bayesian logistic regression convergence as a function of MVP. RAM reaches a low-energy regime with fewer MVPs.
More in this topic

Inference Acceleration

Inference acceleration reduces AI model latency, memory use, and serving cost while preserving output quality.

See all 20 papers

Explore similar work

CardsList

  1. Random Coordinate Descent on the Wasserstein Space of Probability Measures

    Apr 2, 2026Yewei Xu, Qin LiNonconvex OptimizationRiemannian Optimization

    Department of Mathematics, University of Wisconsin-Madison, Madison, WI 53706-1388, USA

  2. Inversion-Free Natural Gradient Descent on Riemannian Manifolds

    Apr 3, 2026Dario Draca, Takuo Matsubara, Minh-Ngoc TranNatural Gradient MethodsRiemannian Optimization

    School of Mathematics and Statistics, University of Sydney, Australia · University of Sydney Business School, Australia

  3. Neural solver for Wasserstein Geodesics and optimal transport dynamics

    Feb 25, 2026Hailiang Liu, Yan-Han ChenOptimal Transport

    Department of Mathematics, Iowa State University, Ames, IA 50011, USA. · Department of Statistics, Iowa State University, Ames, IA 50011, USA.