Transport along Local Annealing Paths: Wasserstein Geometry and Multimodal Sampling
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.
Related
- Wasserstein Gradient Flows and Forward-Only Diffusion Are Not Enough for Multimodal Sampling
- Flow Annealing Posterior Sampling for Function-Space Regression and Inverse Problems
- Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme
- Wasserstein Bounds for Euler and Randomized Midpoint Langevin Sampling with Noisy Gradients
- WASABI: Wasserstein Sampling via Bridge Interpolants
BibTeX
@misc{chen2026transport,
title = {{Transport along Local Annealing Paths: Wasserstein Geometry and Multimodal Sampling}},
author = {Han Chen and Sifan Liu},
howpublished = {arXiv},
year = {2026},
month = nov,
doi = {10.48550/arxiv.2611.05605},
eprint = {2611.05605},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2611.05605},
}
RIS
TY - UNPB
TI - Transport along Local Annealing Paths: Wasserstein Geometry and Multimodal Sampling
AU - Han Chen
AU - Sifan Liu
PY - 2026
DA - 2026/11/03
PB - arXiv
DO - 10.48550/arxiv.2611.05605
UR - https://arxiv.org/abs/2611.05605
AB - 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.
ER -