# Normalizing flow neural networks by JKO scheme

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

## Key facts

- Cited by 39 papers in the map (data through 14 Oct 2026).
- Builds on 56 of its 81 references that are in the map.
- First posted on arXiv 29 Dec 2022.

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering](https://mlsanity.com/paper/1606.09375/) (2016)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [An Introduction to Variational Autoencoders](https://mlsanity.com/paper/1906.02691/) (2019)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Error bounds for approximations with deep ReLU networks](https://mlsanity.com/paper/1610.01145/) (2016)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Invertible Residual Networks](https://mlsanity.com/paper/1811.00995/) (2018)
- [StyleFlow: Attribute-conditioned Exploration of StyleGAN-Generated Images using Conditional Continuous Normalizing Flows](https://mlsanity.com/paper/2008.02401/) (2020)
- [Residual Flows for Invertible Generative Modeling](https://mlsanity.com/paper/1906.02735/) (2019)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [A Variational Perspective on Diffusion-Based Generative Models and Score Matching](https://mlsanity.com/paper/2106.02808/) (2021)
- [Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit](https://mlsanity.com/paper/1905.09883/) (2019)
- [Riemannian Continuous Normalizing Flows](https://mlsanity.com/paper/2006.10605/) (2020)

## Built on by (most cited first, 30 of 39)

- [SAFE: Multitask Failure Detection for Vision-Language-Action Models](https://mlsanity.com/paper/2506.09937/) (2025)
- [A Sharp Convergence Theory for The Probability Flow ODEs of Diffusion Models](https://mlsanity.com/paper/2408.02320/) (2024)
- [Can We Detect Failures Without Failure Data? Uncertainty-Aware Runtime Failure Detection for Imitation Learning Policies](https://mlsanity.com/paper/2503.08558/) (2025)
- [Convergence of flow-based generative models via proximal gradient descent in Wasserstein space](https://mlsanity.com/paper/2310.17582/) (2023)
- [Provable Acceleration for Diffusion Models under Minimal Assumptions](https://mlsanity.com/paper/2410.23285/) (2024)
- [Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling](https://mlsanity.com/paper/2504.10612/) (2025)
- [Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport](https://mlsanity.com/paper/2402.05443/) (2024)
- [On the Relation between Rectified Flows and Optimal Transport](https://mlsanity.com/paper/2505.19712/) (2025)
- [Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approach](https://mlsanity.com/paper/2506.03979/) (2025)
- [Gaussian Interpolation Flows](https://mlsanity.com/paper/2311.11475/) (2023)
- [Flow-based Distributionally Robust Optimization](https://mlsanity.com/paper/2310.19253/) (2023)
- [Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions](https://mlsanity.com/paper/2409.20547/) (2024)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Computing high-dimensional optimal transport by flow neural networks](https://mlsanity.com/paper/2305.11857/) (2023)
- [Distributionally Robust Optimization via Iterative Algorithms in Continuous Probability Spaces](https://mlsanity.com/paper/2412.20556/) (2024)
- [Transformer Neural Autoregressive Flows](https://mlsanity.com/paper/2401.01855/) (2024)
- [Importance Corrected Neural JKO Sampling](https://mlsanity.com/paper/2407.20444/) (2024)
- [Score-based Neural Ordinary Differential Equations for Computing Mean Field Control Problems](https://mlsanity.com/paper/2409.16471/) (2024)
- [Failure Identification in Imitation Learning Via Statistical and Semantic Filtering](https://mlsanity.com/paper/2604.13788/) (2026)
- [Posterior sampling via Langevin dynamics based on generative priors](https://mlsanity.com/paper/2410.02078/) (2024)
- [Local Flow Matching Generative Models](https://mlsanity.com/paper/2410.02548/) (2024)
- [Devil is in the Details: Density Guidance for Detail-Aware Generation with Flow Models](https://mlsanity.com/paper/2502.05807/) (2025)
- [Assured autonomy: How operations research powers and orchestrates generative AI systems](https://mlsanity.com/paper/2512.23978/) (2025)
- [Score-based deterministic density sampling](https://mlsanity.com/paper/2504.18130/) (2025)
- [Generative models for decision-making under distributional shift](https://mlsanity.com/paper/2604.04342/) (2026)
- [Critically Damped Third-Order Langevin Dynamics](https://mlsanity.com/paper/2409.07697/) (2024)
- [Flow-based Generative Modeling of Potential Outcomes and Counterfactuals](https://mlsanity.com/paper/2505.16051/) (2025)
- [Critically-Damped Higher-Order Langevin Dynamics for Generative Modeling](https://mlsanity.com/paper/2506.21741/) (2025)
- [High-dimensional Mean-Field Games by Particle-based Flow Matching](https://mlsanity.com/paper/2512.01172/) (2025)
- [Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity](https://mlsanity.com/paper/2501.14993/) (2025)

## Abstract (arXiv)

Normalizing flow is a class of deep generative models for efficient sampling and likelihood estimation, which achieves attractive performance, particularly in high dimensions. The flow is often implemented using a sequence of invertible residual blocks. Existing works adopt special network architectures and regularization of flow trajectories. In this paper, we develop a neural ODE flow network called JKO-iFlow, inspired by the Jordan-Kinderleherer-Otto (JKO) scheme, which unfolds the discrete-time dynamic of the Wasserstein gradient flow. The proposed method stacks residual blocks one after another, allowing efficient block-wise training of the residual blocks, avoiding sampling SDE trajectories and score matching or variational learning, thus reducing the memory load and difficulty in end-to-end training. We also develop adaptive time reparameterization of the flow network with a progressive refinement of the induced trajectory in probability space to improve the model accuracy further. Experiments with synthetic and real data show that the proposed JKO-iFlow network achieves competitive performance compared with existing flow and diffusion models at a significantly reduced computational and memory cost.