Tangent-Space Preconditioning of Wasserstein Gradient Flows for Particle Sampling

We propose a simple geometric perspective on Wasserstein gradient flows for designing efficient particle-based sampling algorithms in high-dimensional settings. The central idea is to precondition the Wasserstein gradient flow in its tangent space, compensating for slowly relaxing directions that can impede sampling. Remarkably, we show that using the 2-Wasserstein Hessian of the chi-squared flow at the target as a preconditioner recovers the idealized population dynamics underlying Laplacian-adjusted Wasserstein gradient descent (LAWGD), providing a geometric interpretation of LAWGD as Wasserstein preconditioning. This perspective further identifies a scalar-to-vector transport operator as the key object to approximate, leading to two principled finite-sample realizations based on kernels and neural networks. Experiments demonstrate improved sampling efficiency and reveal complementary regimes for the two realizations, with kernels performing well in low dimensions and neural approximations becoming increasingly advantageous as dimension grows.