Transport along Local Annealing Paths: Wasserstein Geometry and Multimodal Sampling

Han Chen; Sifan Liu

cs.LGNov 3, 2026

Transport along Local Annealing Paths: Wasserstein Geometry and Multimodal Sampling

Authors: Han Chen, Sifan Liu

Organizations: Department of Statistical Science, Duke University

Abstract

Variational approaches to learning transport maps for multimodal sampling can suffer from mode collapse. One way to mitigate this problem is to learn transport along a prescribed path of intermediate distributions defined by geometric annealing. We show that annealing can take a substantial detour in Wasserstein space: when mass moves between modes, the path can be much longer than the distance between its endpoints. We further link path geometry to transport learning, showing that long annealing paths can slow population optimization, while sufficiently regular paths permit rapid convergence. Motivated by these findings, we propose componentwise annealed transport (CAT) to reduce the need for mass transfer between modes. Starting from a mode-informed mixture reference, CAT constructs an exact mixture decomposition of the target and learns separate transport maps along componentwise annealing paths. We estimate the component weights separately and show that these estimates remain consistent even when the learned flows are inexact. We also develop adaptive annealing schedules guided by transport effort. Experiments on multimodal synthetic distributions and Bayesian inference problems demonstrate improvements in transport learning and sampling efficiency.

Figures & tables

\frac{\mathrm{d}X_{\beta,i}^{j}}{\mathrm{d}\beta}=v_{\psi}^{j}(\beta,X_{\beta,i}^{j}),\qquad\frac{\mathrm{d}\ell_{\beta,i}^{j}}{\mathrm{d}\beta}=-\nabla\cdot v_{\psi}^{j}(\beta,X_{\beta,i}^{j}).

Algorithm 2 Sampling and density evaluation

Distribution/Model	$n$	$d$
Skew-normal mixture (Section B.3.1 )	–	5
Skew- $t$ product mixture (Section B.3.2 )	–	10
Seemingly unrelated regression (SUR, Section B.3.3 )	50	5
Input-output hidden Markov model (IOHMM, 25 , Section B.3.4 )	500	29
Latent Dirichlet allocation (LDA, 25 , Section B.3.5 )	4999	502

Table 1: Summary of synthetic targets and models used in the experiments, with the number of observations

n

(when applicable), and the number of dimensions

d

Figure 1: Training dynamics of CAT on IOHMM (top) and LDA (bottom). Columns show the training loss, normalized importance-weight ESS, and ELBO against optimizer steps; ESS and ELBO are computed using calibrated weights. Open circles mark the initial base distribution

b

at iteration 0.

Figure 2: Wall-clock comparison of CAT, GAT, and RKL-CNF on the LDA posterior. Curves show normalized importance-weight ESS, sliced 2-Wasserstein distance (SW2), and ELBO. Open circles give the values for the shared base distribution

b

before training; insets enlarge the early training period.

Figure 3: Learned trajectories and cumulative transport lengths for CAT and GAT at

\gamma=1

(top) and

\gamma=4

(bottom). Left and middle: trajectories projected onto the first coordinate, colored by initial component. Gray dashed curves show oracle trajectories and target marginals; histograms show terminal samples. Annotations give CAT’s calibrated mixture weights and GAT’s switching rates. Right: cumulative lengths computed from the full ten-dimensional velocity fields.

Method	Recovery $\uparrow$	Mode TV $\downarrow$	SW2 $\downarrow$	Time (s)	ESS $/N$ $\uparrow$	$\operatorname{Var}_{q}(\log w)$ $\downarrow$
Skew-normal mixture
CAT	1.00	0.004 $\pm$ 0.004	0.254 $\pm$ 0.121	61 $\pm$ 5	0.994 $\pm$ 0.002	0.01 $\pm$ 0.00
GMM	1.00	0.409 $\pm$ 0.009	2.591 $\pm$ 0.091	11 $\pm$ 1	0.536 $\pm$ 0.007	52.85 $\pm$ 2.31
FAB	1.00	0.408 $\pm$ 0.008	2.753 $\pm$ 0.100	36 $\pm$ 2	0.559 $\pm$ 0.008	18.73 $\pm$ 2.96
TF	1.00	0.407 $\pm$ 0.008	2.861 $\pm$ 0.099	37 $\pm$ 2	0.539 $\pm$ 0.050	1.23 $\pm$ 0.12
APT	0.97	0.139 $\pm$ 0.109	1.332 $\pm$ 0.713	17 $\pm$ 1	–	–

Table 2: Comparison on four targets over 100 runs. Entries report means

\pm

std, except Recovery. Arrows indicate the preferred direction; dashes denote unreported quantities.

Figure 4: First-coordinate marginals for the skew-normal mixture target. Each panel compares a method’s sample histogram (blue) with the exact target marginal density (red). The two modes reveal differences in relative mass and in the shape of the skewed mode.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Target	Refreshes $\times$ Adam steps	Trajectories	Minibatch	$H$
Skew-normal mixture	$10\times 100$	1024	256	64
SUR	$20\times 100$	768	512	64
IOHMM	$20\times 250$	2560	512	32
LDA	$20\times 250$	6144	512	64

Appendix

Table 3: CAT training settings. Trajectory counts and minibatch sizes are per mode; minibatches contain collocation points. The width

H

is shared by CAT and RKL-CNF.

$\lambda$	Skew-normal mixture	SUR	IOHMM	LDA
$0$	$0.995\pm 0.004$	$0.763\pm 0.206$	$0.948\pm 0.094$	$0.958\pm 0.038$
$10^{-4}$	$0.995\pm 0.004$	$0.797\pm 0.217$	$0.948\pm 0.095$	$0.970\pm 0.012$
$3\times 10^{-4}$	$0.995\pm 0.004$	$0.786\pm 0.208$	$0.948\pm 0.094$	$0.971\pm 0.014$
$\boldsymbol{10^{-3}}$	$\boldsymbol{0.995\pm 0.004}$	$\boldsymbol{0.787\pm 0.214}$	$\boldsymbol{0.948\pm 0.094}$	$\boldsymbol{0.969\pm 0.012}$
$3\times 10^{-3}$	$0.995\pm 0.004$	$0.788\pm 0.214$	$0.948\pm 0.089$	$0.967\pm 0.013$
$10^{-2}$	$0.995\pm 0.003$	$0.786\pm 0.218$	$0.947\pm 0.081$	$0.950\pm 0.046$

Appendix

Table 4: Sensitivity of CAT to the regularization coefficient

\lambda

. Entries report mean

\pm

std of relative ESS over 5 independent runs, separate from those in Section 6 . The default value is shown in bold.

Multimodal Distribution Sampling

See all 8 papers

Wasserstein Gradient Flows and Forward-Only Diffusion Are Not Enough for Multimodal Sampling
Oct 1, 2026Daniel McBride, Pratik Khandagale, Cristina Garcia-Cardona +1Dynamical Systems Wasserstein Gradient Flows
Los Alamos National Laboratory, Los Alamos, NM 87545, USA
Sampling from multimodal distributions with warm starts: Non-asymptotic bounds for the Reweighted Annealed Leap-Point Sampler
Dec 19, 2025Holden Lee, Matheau Santana-GijzenMultimodal Distribution Sampling Markov Chain Monte Carlo
Fixed-point acceleration on Neural Wasserstein manifolds
Nov 6, 2026Vitalii Aksenov, Sebastian Gutierrez HernandezInference Acceleration Riemannian Optimization
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

Oct 1, 2026J/K move · Enter open · S save

Wasserstein Gradient Flows and Forward-Only Diffusion Are Not Enough for Multimodal Sampling

Daniel McBride, Pratik Khandagale, Cristina Garcia-Cardona, Yen Ting Lin

Los Alamos National Laboratory, Los Alamos, NM 87545, USA

Dynamical Systems Wasserstein Gradient Flows Diffusion Sampling Multimodal Distribution Sampling

Open paper AskPDF arXiv

There has been a proliferation of sampling algorithms based on Wasserstein gradient flows (WGF) and forward-only diffusion processes (FODP), often accompanied by theoretical guarantees of exponentially fast convergence to the target distribution. These guarantees are frequently interpreted as evidence that such methods can efficiently sample complex multimodal distributions, often supported by empirical results. In this work, we argue that this interpretation is fundamentally misleading. By invoking the Jordan-Kinderlehrer-Otto (JKO) scheme and Otto calculus, we establish that the canonical WGF sampling dynamics and overdamped forward diffusion share the same density evolution and therefore inherit the same metastability and slow-mixing phenomena long understood in nonequilibrium statistical physics. We analyze this family of samplers using two complementary tools -- spectral analysis and mean first-passage time (MFPT) analysis -- and show that well-separated multimodality can induce exponentially long mixing times associated with small spectral gaps and rare inter-mode transitions. For the commonly adopted log-linear annealing schedule studied here, we find that introducing intermediate distributions does not remove the exponential scaling of the total transport time. The limitation is structural rather than implementation-specific: purely local, gradient-driven transport mechanisms can require exponentially long times to transport probability mass across well-separated modes. We argue that this represents a fundamental limitation of WGF- and FODP-based sampling in their standard forms, and motivates future development of fundamentally nonlocal mechanisms for efficient multimodal sampling.

Abstract

Multimodal Distribution Sampling

Explore similar work

Wasserstein Gradient Flows and Forward-Only Diffusion Are Not Enough for Multimodal Sampling

Sampling from multimodal distributions with warm starts: Non-asymptotic bounds for the Reweighted Annealed Leap-Point Sampler

Fixed-point acceleration on Neural Wasserstein manifolds