Wasserstein flows converge sharply for special interaction energies
Sharp Convergence of Wasserstein Gradient Flows for Spectrally Nonnegative Interaction Energies
Machine Learning
Summary
This paper studies how a certain kind of mathematical process, called a Wasserstein gradient flow, evolves over time when influenced by special interaction energies on curved spaces. The authors prove that for these special energies, the flow's energy decreases quickly and predictably, eventually settling down. They show this even when standard theory doesn’t apply and provide precise rates at which convergence occurs. They also build examples demonstrating their rates are almost the best possible. Their results cover kernels from areas like spheres, transformer neural networks, and fractional Laplacians.
What this means in practice
- •For machine learning practitioners: Improve the understanding of convergence behaviors in transformer models' kernel interactions used in attention mechanisms.
- •For data scientists working with spherical data: Apply refined convergence guarantees when analyzing distributions evolving on spheres with zonal kernels in spatial statistics.
Authors
Zhengjiang Lin, Philippe Rigollet
Abstract
We study the long-time behavior of Wasserstein gradient flows for interaction energies \[ \mathsf E[μ] = \frac12\iint_{M\times M}K(x,y)\,\mathrm dμ(x)\,\mathrm dμ(y) \] on a closed manifold $M$. For kernels diagonal in a Laplace eigenbasis with nonnegative spectral coefficients, we prove a differential inequality relating the relative entropy to the energy gap. Consequently, for any nonnegative initial density $u_0\in L^p(M)$, $p>1$, the energy gap is integrable in time and satisfies \[ \mathsf E[μ_t]-\mathsf E_{\min}=o(t^{-1}). \] If all spectral coefficients are positive, the flow converges weakly to the constant measure. These interaction energies need not be geodesically convex in Wasserstein space, and the associated flows contain no diffusion; their global convergence therefore does not follow from standard Wasserstein gradient flow theory. The kernels covered by our results include zonal kernels on spheres, kernels arising in transformer models, regularized Riesz kernels, and inverse fractional Laplacian kernels. We also investigate the sharpness of the $o(t^{-1})$ rate. For any smooth kernel in this class with infinitely many positive spectral coefficients and any $δ>0$, we construct a solution of the linearized flow whose energy is comparable to $t^{-1-δ}$ along a sequence of times tending to infinity. Moreover, for any $δ>0$, by choosing a suitable inverse fractional Laplacian kernel on the flat torus, we construct an exact solution of the nonlinear Wasserstein gradient flow whose energy is comparable to $t^{-1-δ}$. The nonlinear construction is based on uniform-in-time estimates for the evolution of the dyadic Fourier coefficient blocks and a blockwise energy-persistence argument. These estimates also yield a uniform-in-time quantitative comparison between the nonlinear Wasserstein gradient flow and its linearization.