mlsanityData to 13 Oct 2026

Normalizing flow neural networks by JKO scheme

arXiv 2212.14424 ↗First posted 29 Dec 2022stat.ML cs.LGcited by 39

Not on the map. Why?

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.

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

Builds on

56 of 81 references.

References · 56

  1. 2020Denoising Diffusion Probabilistic Modelscited by 20,952
  2. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,223
  3. 2021High-Resolution Image Synthesis with Latent Diffusion Modelscited by 18,777
  4. 2014Very Deep Convolutional Networks for Large-Scale Image Recognitioncited by 15,321
  5. 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 9,028
  6. 2022Flow Matching for Generative Modelingcited by 5,542
  7. 2016Image-to-Image Translation with Conditional Adversarial Networkscited by 4,516
  8. 2019Generative Modeling by Estimating Gradients of the Data Distributioncited by 4,077
  9. 2018Neural Ordinary Differential Equationscited by 3,750
  10. 2020Taming Transformers for High-Resolution Image Synthesiscited by 3,237
  11. 2016Improved Techniques for Training GANscited by 2,946
  12. 2019Fast Graph Representation Learning with PyTorch Geometriccited by 2,422
38 more
  1. 2017Improved Training of Wasserstein GANscited by 2,226
  2. 2016Convolutional Neural Networks on Graphs with Fast Localized Spectral Filteringcited by 2,086
  3. 2016Density estimation using Real NVPcited by 2,033
  4. 2018Glow: Generative Flow with Invertible 1x1 Convolutionscited by 1,542
  5. 2014NICE: Non-linear Independent Components Estimationcited by 1,054
  6. 2019An Introduction to Variational Autoencoderscited by 998
  7. 2022Building Normalizing Flows with Stochastic Interpolantscited by 979
  8. 2017Masked Autoregressive Flow for Density Estimationcited by 680
  9. 2016Error bounds for approximations with deep ReLU networkscited by 570
  10. 2016Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithmcited by 502
  11. 2016Input Convex Neural Networkscited by 408
  12. 2018Invertible Residual Networkscited by 364
  13. 2020StyleFlow: Attribute-conditioned Exploration of StyleGAN-Generated Images using Conditional Continuous Normalizing Flowscited by 348
  14. 2019Residual Flows for Invertible Generative Modelingcited by 257
  15. 2020How to train your neural ODE: the world of Jacobian and kinetic regularizationcited by 228
  16. 2021A Variational Perspective on Diffusion-Based Generative Models and Score Matchingcited by 182
  17. 2019Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limitcited by 173
  18. 2020Riemannian Continuous Normalizing Flowscited by 145
  19. 2020OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transportcited by 143
  20. 2020Generative Modeling with Denoising Auto-Encoders and Langevin Samplingcited by 119
  21. 2019Graph Normalizing Flowscited by 103
  22. 2016Generative Models and Model Criticism via Optimized Maximum Mean Discrepancycited by 99
  23. 2019Theoretical guarantees for sampling and inference in generative models with latent diffusionscited by 97
  24. 2019Wasserstein-2 Generative Networkscited by 96
  25. 2019A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problemscited by 95
  26. 2021Proximal Optimal Transport Modeling of Population Dynamicscited by 93
  27. 2020Interacting particle solutions of Fokker-Planck equations through gradient-log-density estimationcited by 87
  28. 2021Generative Modeling with Optimal Transport Mapscited by 81
  29. 2018Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusionscited by 77
  30. 2021Diffusion Normalizing Flowcited by 77
  31. 2021Large-Scale Wasserstein Gradient Flowscited by 76
  32. 2021MMD Aggregated Two-Sample Testcited by 74
  33. 2021Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networkscited by 53
  34. 2021Variational Wasserstein gradient flowcited by 53
  35. 2022Probability flow solution of the Fokker-Planck equationcited by 46
  36. 2018Monge-Amp`ere Flow for Generative Modelingcited by 33
  37. 2017On the ability of neural nets to express distributionscited by 26
  38. 2022Self-Consistency of the Fokker-Planck Equationcited by 18

50 most cited of 56 All 56 (JSON)

Built on by

Cited by year (2026 to 13 Oct)

Citing papers per year: 2023: 4; 2024: 11; 2025: 18; 2026: 6.

39 papers.

Sort:

By area:

  1. 2025SAFE: Multitask Failure Detection for Vision-Language-Action Modelscited by 77
  2. 2024A Sharp Convergence Theory for The Probability Flow ODEs of Diffusion Modelscited by 64
  3. 2025Can We Detect Failures Without Failure Data? Uncertainty-Aware Runtime Failure Detection for Imitation Learning Policiescited by 51
  4. 2023Convergence of flow-based generative models via proximal gradient descent in Wasserstein spacecited by 34
  5. 2024Provable Acceleration for Diffusion Models under Minimal Assumptionscited by 27
  6. 2025Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modelingcited by 25
  7. 2024Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transportcited by 22
  8. 2025Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approachcited by 18
  9. 2023Gaussian Interpolation Flowscited by 17
  10. 2025On the Relation between Rectified Flows and Optimal Transportcited by 16
  11. 2023Flow-based Distributionally Robust Optimizationcited by 13
  12. 2024Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributionscited by 11
27 more
  1. 2025Learning of Population Dynamics: Inverse Optimization Meets JKO Schemecited by 9
  2. 2023Computing high-dimensional optimal transport by flow neural networkscited by 8
  3. 2024Distributionally Robust Optimization via Iterative Algorithms in Continuous Probability Spacescited by 8
  4. 2024Transformer Neural Autoregressive Flowscited by 7
  5. 2024Importance Corrected Neural JKO Samplingcited by 7
  6. 2024Score-based Neural Ordinary Differential Equations for Computing Mean Field Control Problemscited by 7
  7. 2026Failure Identification in Imitation Learning Via Statistical and Semantic Filteringcited by 6
  8. 2024Posterior sampling via Langevin dynamics based on generative priorscited by 5
  9. 2024Local Flow Matching Generative Modelscited by 5
  10. 2025Devil is in the Details: Density Guidance for Detail-Aware Generation with Flow Modelscited by 5
  11. 2025Score-based deterministic density samplingcited by 4
  12. 2025Assured autonomy: How operations research powers and orchestrates generative AI systemscited by 4
  13. 2026Generative models for decision-making under distributional shiftcited by 4
  14. 2024Critically Damped Third-Order Langevin Dynamicscited by 3
  15. 2025Critically-Damped Higher-Order Langevin Dynamics for Generative Modelingcited by 3
  16. 2025High-dimensional Mean-Field Games by Particle-based Flow Matchingcited by 3
  17. 2025Flow-based generative models as iterative algorithms in probability spacecited by 2
  18. 2025Accelerated Markov Chain Monte Carlo Algorithms on Discrete Statescited by 2
  19. 2025Dale meets Langevin: A Multiplicative Denoising Diffusion Modelcited by 2
  20. 2026Learning Discrete Diffusion of Graphs via Free-Energy Gradient Flowscited by 2
  21. 2025Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexitycited by 1
  22. 2025Flow-based Generative Modeling of Potential Outcomes and Counterfactualscited by 1
  23. 2026Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flowcited by 1
  24. 2025Adversarial Transform Particle Filterscited by 0
  25. 2025Liouville PDE-based sliced-Wasserstein flowcited by 0
  26. 2026Analyzing and Improving Diffusion Models for Time-Series Data Imputation: A Proximal Recursion Perspectivecited by 0
  27. 2026Beyond Trajectory Matching: Reflow with Marginal Distribution Alignmentcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025ODEt(ODEl): Shortcutting the Time and the Length in Diffusion and Flow Models for Faster Sampling
  2. 2026From Euler to Dormand-Prince: ODE Solvers for Flow Matching Generative Models
  3. 2026Normalizing Flows with Iterative Denoising
  4. 2022VQ-Flows: Vector Quantized Local Normalizing Flows
  5. 2025Generative Modeling with Continuous Flows: Sample Complexity of Flow Matching
  6. 2024Markovian Flow Matching: Accelerating MCMC with Continuous Normalizing Flows
  7. 2024Convergence Analysis of Flow Matching in Latent Space with Transformers
  8. 2023Improving and generalizing flow-based generative models with minibatch optimal transport