# Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow

arXiv: https://arxiv.org/abs/2601.05583 · HTML: https://mlsanity.com/paper/2601.05583/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 40 references that are in the map.
- First posted on arXiv 9 Jan 2026.

## Builds on (most cited first, 22 of 22)

- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Fourier Neural Operator with Learned Deformations for PDEs on General Geometries](https://mlsanity.com/paper/2207.05209/) (2022)
- [Performative Prediction](https://mlsanity.com/paper/2002.06673/) (2020)
- [Regularisation of Neural Networks by Enforcing Lipschitz Continuity](https://mlsanity.com/paper/1804.04368/) (2018)
- [Efficient and Accurate Estimation of Lipschitz Constants for Deep Neural Networks](https://mlsanity.com/paper/1906.04893/) (2019)
- [PRNet: Self-Supervised Learning for Partial-to-Partial Registration](https://mlsanity.com/paper/1910.12240/) (2019)
- [In-Context Operator Learning with Data Prompts for Differential Equation Problems](https://mlsanity.com/paper/2304.07993/) (2023)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [Performative Prediction: Past and Future](https://mlsanity.com/paper/2310.16608/) (2023)
- [Normalizing flow neural networks by JKO scheme](https://mlsanity.com/paper/2212.14424/) (2022)
- [Policy Optimization as Wasserstein Gradient Flows](https://mlsanity.com/paper/1808.03030/) (2018)
- [Performative Prediction with Neural Networks](https://mlsanity.com/paper/2304.06879/) (2023)
- [Convergence of flow-based generative models via proximal gradient descent in Wasserstein space](https://mlsanity.com/paper/2310.17582/) (2023)
- [Some Fundamental Aspects about Lipschitz Continuity of Neural Networks](https://mlsanity.com/paper/2302.10886/) (2023)
- [Taming Hyperparameter Tuning in Continuous Normalizing Flows Using the JKO Scheme](https://mlsanity.com/paper/2211.16757/) (2022)
- [Understanding In-Context Learning on Structured Manifolds: Bridging Attention to Kernel Methods](https://mlsanity.com/paper/2506.10959/) (2025)
- [A kernel formula for regularized Wasserstein proximal operators](https://mlsanity.com/paper/2301.10301/) (2023)
- [On Wasserstein Reinforcement Learning and the Fokker-Planck equation](https://mlsanity.com/paper/1712.07185/) (2017)
- [Importance Corrected Neural JKO Sampling](https://mlsanity.com/paper/2407.20444/) (2024)
- [Learning Proximal Operators to Discover Multiple Optima](https://mlsanity.com/paper/2201.11945/) (2022)
- [A Lagrangian approach to totally dissipative evolutions in Wasserstein spaces](https://mlsanity.com/paper/2305.05211/) (2023)

## Built on by (most cited first, 2 of 2)

- [Deep operator learning for efficient sampling from invariant measures of stochastic differential equations](https://mlsanity.com/paper/2609.11376/) (2026)
- [Transformer-based In-Context Score Operator Learning for Diffusion Models](https://mlsanity.com/paper/2610.13690/) (2026)

## Abstract (arXiv)

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.