Wasserstein Gradient Flows and Forward-Only Diffusion Are Not Enough for Multimodal Sampling
Abstract
There has been a plethora of algorithms proposed that leverage Wasserstein gradient flows or forward-only diffusion processes to perform sampling tasks. These approaches are often characterized by theoretical guarantees of exponentially fast convergence to the target distribution. In this work, we study the mixing behavior of this family of sampling methods. By invoking the Jordan--Kinderlehrer--Otto scheme and Otto calculus, we first establish that Wasserstein gradient flow and forward diffusion-based samplers share the same density evolution. Consequently, their convergence can be characterized by studying the convergence behavior of the associated diffusion process, which has been extensively analyzed in nonequilibrium statistical physics. We perform two independent and well-established analyses for this purpose, namely spectral analysis and mean first-passage time analysis. We show that, although the sampling distributions converge to the target distribution exponentially fast, the presence of multimodality can lead to exponentially long mixing times associated with small spectral gaps. Further analysis elucidates that even when combined with annealing, this family of samplers still requires impractically long time to sample multimodal distributions. We discuss the origin of this behavior as a consequence of the local (i.e., gradient-driven) dynamics underlying this class of samplers, and highlight the need for non-local mechanisms to facilitate the transport of probability mass across modes, enabling efficient exploration in multimodal settings.