Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow
This work addresses a computational bottleneck in computing Wasserstein gradient flows, which is important for researchers in machine learning and applied mathematics, but it is incremental as it builds on the existing JKO framework with a novel learning approach.
The authors tackled the high computational cost of solving JKO subproblems for Wasserstein gradient flows by proposing a self-supervised method to learn a JKO solution operator, which maps input densities directly to minimizers and iteratively generates gradient-flow evolution efficiently, achieving accuracy, stability, and robustness in numerical experiments.
The Jordan-Kinderlehrer-Otto (JKO) scheme provides a stable variational framework for computing Wasserstein gradient flows, but its practical use is often limited by the high computational cost of repeatedly solving the JKO subproblems. We propose a self-supervised approach for learning a JKO solution operator without requiring numerical solutions of any JKO trajectories. The learned operator maps an input density directly to the minimizer of the corresponding JKO subproblem, and can be iteratively applied to efficiently generate the gradient-flow evolution. A key challenge is that only a number of initial densities are typically available for training. To address this, we introduce a Learn-to-Evolve algorithm that jointly learns the JKO operator and its induced trajectories by alternating between trajectory generation and operator updates. As training progresses, the generated data increasingly approximates true JKO trajectories. Meanwhile, this Learn-to-Evolve strategy serves as a natural form of data augmentation, significantly enhancing the generalization ability of the learned operator. Numerical experiments demonstrate the accuracy, stability, and robustness of the proposed method across various choices of energies and initial conditions.