# Learning diffusion at lightspeed

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

## Key facts

- Cited by 22 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 54 references that are in the map.
- First posted on arXiv 18 Jun 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Manifold Interpolating Optimal-Transport Flows for Trajectory Inference](https://mlsanity.com/paper/2206.14928/) (2022)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Deep Momentum Multi-Marginal Schr"odinger Bridge](https://mlsanity.com/paper/2303.01751/) (2023)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [Challenging Common Assumptions in Convex Reinforcement Learning](https://mlsanity.com/paper/2202.01511/) (2022)
- [Riemannian Metric Learning via Optimal Transport](https://mlsanity.com/paper/2205.09244/) (2022)
- [First-order Conditions for Optimization in the Wasserstein Space](https://mlsanity.com/paper/2209.12197/) (2022)
- [Trust Region Policy Optimization with Optimal Transport Discrepancies: Duality and Algorithm for Continuous Actions](https://mlsanity.com/paper/2210.11137/) (2022)
- [Improved Complexity Analysis of the Sinkhorn and Greenkhorn Algorithms for Optimal Transport](https://mlsanity.com/paper/2305.14939/) (2023)
- [Dynamic Programming in Probability Spaces via Optimal Transport](https://mlsanity.com/paper/2302.13550/) (2023)
- [Variational Analysis in the Wasserstein Space](https://mlsanity.com/paper/2406.10676/) (2024)
- [Non-Parametric Learning of Stochastic Differential Equations with Non-asymptotic Fast Rates of Convergence](https://mlsanity.com/paper/2305.15557/) (2023)

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

- [Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling](https://mlsanity.com/paper/2504.10612/) (2025)
- [Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling](https://mlsanity.com/paper/2505.13413/) (2025)
- [Pinet: Optimizing hard-constrained neural networks with orthogonal projection layers](https://mlsanity.com/paper/2508.10480/) (2025)
- [Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots](https://mlsanity.com/paper/2410.22729/) (2024)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [DICE: Discrete inverse continuity equation for learning population dynamics](https://mlsanity.com/paper/2507.05107/) (2025)
- [Generative emulation of chaotic dynamics with coherent prior](https://mlsanity.com/paper/2504.14264/) (2025)
- [A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots](https://mlsanity.com/paper/2605.08550/) (2026)
- [Gradient-flow SDEs have unique transient population dynamics](https://mlsanity.com/paper/2505.21770/) (2025)
- [Flows and Diffusions on the Neural Manifold](https://mlsanity.com/paper/2507.10623/) (2025)
- [Learning Explicit Single-Cell Dynamics Using ODE Representations](https://mlsanity.com/paper/2510.02903/) (2025)
- [Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation](https://mlsanity.com/paper/2603.23398/) (2026)
- [Leveraging Gauge Freedom for Learning Non-Gradient Population Dynamics of Stochastic Systems](https://mlsanity.com/paper/2605.25107/) (2026)
- [Two-Parameter Flows for Learning Population Dynamics of Physical Systems](https://mlsanity.com/paper/2605.26285/) (2026)
- [First-Order Trajectory Matching: Fast Ensemble Predictions of Chaotic, Turbulent, Stochastic Systems](https://mlsanity.com/paper/2606.11138/) (2026)
- [Wasserstein Residuals: Learning Gradient Flows from Population Dynamics](https://mlsanity.com/paper/2607.04738/) (2026)
- [Using Linearized Optimal Transport to Predict the Evolution of Stochastic Particle Systems](https://mlsanity.com/paper/2408.01857/) (2024)
- [Multimarginal flow matching with optimal transport potentials](https://mlsanity.com/paper/2606.05327/) (2026)
- [Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization](https://mlsanity.com/paper/2606.27767/) (2026)
- [Beyond Gradient Flow: Identifiability and Recovery from Distribution Snapshots](https://mlsanity.com/paper/2609.35060/) (2026)
- [Discrete Wasserstein Flows for One-Step Generative Modeling](https://mlsanity.com/paper/2610.01355/) (2026)
- [Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals](https://mlsanity.com/paper/2610.03679/) (2026)

## Abstract (arXiv)

Diffusion regulates numerous natural processes and the dynamics of many successful generative models. Existing models to learn the diffusion terms from observational data rely on complex bilevel optimization problems and model only the drift of the system. We propose a new simple model, JKOnet*, which bypasses the complexity of existing architectures while presenting significantly enhanced representational capabilities: JKOnet* recovers the potential, interaction, and internal energy components of the underlying diffusion process. JKOnet* minimizes a simple quadratic loss and outperforms other baselines in terms of sample efficiency, computational complexity, and accuracy. Additionally, JKOnet* provides a closed-form optimal solution for linearly parametrized functionals, and, when applied to predict the evolution of cellular processes from real-world data, it achieves state-of-the-art accuracy at a fraction of the computational cost of all existing methods. Our methodology is based on the interpretation of diffusion processes as energy-minimizing trajectories in the probability space via the so-called JKO scheme, which we study via its first-order optimality conditions.