# Taming Hyperparameter Tuning in Continuous Normalizing Flows Using the JKO Scheme

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

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 78 references that are in the map.
- First posted on arXiv 30 Nov 2022.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [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)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Improving Variational Inference with Inverse Autoregressive Flow](https://mlsanity.com/paper/1606.04934/) (2016)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Neural Autoregressive Flows](https://mlsanity.com/paper/1804.00779/) (2018)
- [OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport](https://mlsanity.com/paper/2006.00104/) (2020)
- [Improving GANs Using Optimal Transport](https://mlsanity.com/paper/1803.05573/) (2018)
- [An Introduction to Deep Generative Modeling](https://mlsanity.com/paper/2103.05180/) (2021)
- [A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems](https://mlsanity.com/paper/1912.01825/) (2019)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Unconstrained Monotonic Neural Networks](https://mlsanity.com/paper/1908.05164/) (2019)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [On the Convergence and Robustness of Training GANs with Regularized Optimal Transport](https://mlsanity.com/paper/1802.08249/) (2018)
- [Discretize-Optimize vs. Optimize-Discretize for Time-Series Regression and Continuous Normalizing Flows](https://mlsanity.com/paper/2005.13420/) (2020)
- [Alternating the Population and Control Neural Networks to Solve High-Dimensional Stochastic Mean-Field Games](https://mlsanity.com/paper/2002.10113/) (2020)
- [Normalizing field flows: Solving forward and inverse stochastic differential equations using physics-informed flow models](https://mlsanity.com/paper/2108.12956/) (2021)
- [Monge-Amp`ere Flow for Generative Modeling](https://mlsanity.com/paper/1809.10188/) (2018)

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

- [Normalizing flow neural networks by JKO scheme](https://mlsanity.com/paper/2212.14424/) (2022)
- [Convergence of flow-based generative models via proximal gradient descent in Wasserstein space](https://mlsanity.com/paper/2310.17582/) (2023)
- [Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport](https://mlsanity.com/paper/2402.05443/) (2024)
- [Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approach](https://mlsanity.com/paper/2506.03979/) (2025)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency](https://mlsanity.com/paper/2505.13499/) (2025)
- [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)
- [OT-Transformer: A Continuous-time Transformer Architecture with Optimal Transport Regularization](https://mlsanity.com/paper/2501.18793/) (2025)
- [Local Flow Matching Generative Models](https://mlsanity.com/paper/2410.02548/) (2024)
- [End-to-End Training of High-Dimensional Optimal Control with Implicit Hamiltonians via Jacobian-Free Backpropagation](https://mlsanity.com/paper/2510.00359/) (2025)
- [Flow-based generative models as iterative algorithms in probability space](https://mlsanity.com/paper/2502.13394/) (2025)
- [Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow](https://mlsanity.com/paper/2601.05583/) (2026)
- [On the Convergence of Jacobian-Free Backpropagation for Optimal Control Problems with Implicit Hamiltonians](https://mlsanity.com/paper/2602.00921/) (2026)
- [StyleAutoEncoder for manipulating image attributes using pre-trained StyleGAN](https://mlsanity.com/paper/2412.20164/) (2024)
- [Dynamical Implicit Neural Representations](https://mlsanity.com/paper/2511.21787/) (2025)
- [End-to-End Learning of Safe Optimal Feedback Control in High Dimensions with Control Barrier Function Layers](https://mlsanity.com/paper/2607.20674/) (2026)

## Abstract (arXiv)

A normalizing flow (NF) is a mapping that transforms a chosen probability distribution to a normal distribution. Such flows are a common technique used for data generation and density estimation in machine learning and data science. The density estimate obtained with a NF requires a change of variables formula that involves the computation of the Jacobian determinant of the NF transformation. In order to tractably compute this determinant, continuous normalizing flows (CNF) estimate the mapping and its Jacobian determinant using a neural ODE. Optimal transport (OT) theory has been successfully used to assist in finding CNFs by formulating them as OT problems with a soft penalty for enforcing the standard normal distribution as a target measure. A drawback of OT-based CNFs is the addition of a hyperparameter, α, that controls the strength of the soft penalty and requires significant tuning. We present JKO-Flow, an algorithm to solve OT-based CNF without the need of tuning α. This is achieved by integrating the OT CNF framework into a Wasserstein gradient flow framework, also known as the JKO scheme. Instead of tuning α, we repeatedly solve the optimization problem for a fixed α effectively performing a JKO update with a time-step α. Hence we obtain a "divide and conquer" algorithm by repeatedly solving simpler problems instead of solving a potentially harder problem with large α.