# Stable Training of Normalizing Flows for High-dimensional Variational Inference

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 41 references that are in the map.
- First posted on arXiv 26 Feb 2024.

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

- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [Advances in Variational Inference](https://mlsanity.com/paper/1711.05597/) (2017)
- [Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro](https://mlsanity.com/paper/1912.11554/) (2019)
- [Automatic Differentiation Variational Inference](https://mlsanity.com/paper/1603.00788/) (2016)
- [Guided Image Generation with Conditional Invertible Neural Networks](https://mlsanity.com/paper/1907.02392/) (2019)
- [Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference](https://mlsanity.com/paper/1703.09194/) (2017)
- [Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2102.07501/) (2021)
- [Understanding and Mitigating Exploding Inverses in Invertible Neural Networks](https://mlsanity.com/paper/2006.09347/) (2020)
- [Augmented Normalizing Flows: Bridging the Gap Between Generative Flows and Latent Variable Models](https://mlsanity.com/paper/2002.07101/) (2020)
- [Continual Repeated Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2201.13117/) (2022)
- [normflows: A PyTorch Package for Normalizing Flows](https://mlsanity.com/paper/2302.12014/) (2023)
- [Can Push-forward Generative Models Fit Multimodal Distributions?](https://mlsanity.com/paper/2206.14476/) (2022)
- [Challenges and Opportunities in High-dimensional Variational Inference](https://mlsanity.com/paper/2103.01085/) (2021)
- [An invitation to sequential Monte Carlo samplers](https://mlsanity.com/paper/2007.11936/) (2020)
- [Stabilizing Invertible Neural Networks Using Mixture Models](https://mlsanity.com/paper/2009.02994/) (2020)
- [Universal approximation property of invertible neural networks](https://mlsanity.com/paper/2204.07415/) (2022)
- [Variational Refinement for Importance Sampling Using the Forward Kullback-Leibler Divergence](https://mlsanity.com/paper/2106.15980/) (2021)
- [Representational aspects of depth and conditioning in normalizing flows](https://mlsanity.com/paper/2010.01155/) (2020)
- [Gradients should stay on Path: Better Estimators of the Reverse- and Forward KL Divergence for Normalizing Flows](https://mlsanity.com/paper/2207.08219/) (2022)
- [Robust, Accurate Stochastic Optimization for Variational Inference](https://mlsanity.com/paper/2009.00666/) (2020)
- [Fat-Tailed Variational Inference with Anisotropic Tail Adaptive Flows](https://mlsanity.com/paper/2205.07918/) (2022)
- [Machine Learning and the Future of Bayesian Computation](https://mlsanity.com/paper/2304.11251/) (2023)

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

- [FALCON: Few-step Accurate Likelihoods for Continuous Flows](https://mlsanity.com/paper/2512.09914/) (2025)
- [Efficient Regression-Based Training of Normalizing Flows for Boltzmann Generators](https://mlsanity.com/paper/2506.01158/) (2025)
- [Disentangling impact of capacity, objective, batchsize, estimators, and step-size on flow VI](https://mlsanity.com/paper/2412.08824/) (2024)
- [Preferential Normalizing Flows](https://mlsanity.com/paper/2410.08710/) (2024)
- [Analytic Bijections for Smooth and Interpretable Normalizing Flows](https://mlsanity.com/paper/2601.10774/) (2026)
- [NFPO: Stabilized Policy Optimization of Normalizing Flow for Robotic Policy Learning](https://mlsanity.com/paper/2603.11470/) (2026)

## Abstract (arXiv)

Variational inference with normalizing flows (NFs) is an increasingly popular alternative to MCMC methods. In particular, NFs based on coupling layers (Real NVPs) are frequently used due to their good empirical performance. In theory, increasing the depth of normalizing flows should lead to more accurate posterior approximations. However, in practice, training deep normalizing flows for approximating high-dimensional posterior distributions is often infeasible due to the high variance of the stochastic gradients. In this work, we show that previous methods for stabilizing the variance of stochastic gradient descent can be insufficient to achieve stable training of Real NVPs. As the source of the problem, we identify that, during training, samples often exhibit unusual high values. As a remedy, we propose a combination of two methods: (1) soft-thresholding of the scale in Real NVPs, and (2) a bijective soft log transformation of the samples. We evaluate these and other previously proposed modification on several challenging target distributions, including a high-dimensional horseshoe logistic regression model. Our experiments show that with our modifications, stable training of Real NVPs for posteriors with several thousand dimensions is possible, allowing for more accurate marginal likelihood estimation via importance sampling. Moreover, we evaluate several common training techniques and architecture choices and provide practical advise for training NFs for high-dimensional variational inference.