Stable Training of Normalizing Flows for High-dimensional Variational Inference

Daniel Andrade
2024-02-26
Abstract: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.
Machine Learning
What problem does this paper attempt to address?
This paper focuses on the issue of training instability when using normalizing flows (NFs) in high-dimensional variational inference. Normalizing flows, especially real NVPs based on coupling layers, have become an alternative choice to Markov Chain Monte Carlo (MCMC) methods due to their theoretical performance and computational efficiency. However, training deep normalizing flows becomes impractical when dealing with high-dimensional posterior distributions due to the high variance of stochastic gradients. The authors identified that during the training process, the sample values often exhibit extremely high magnitudes, leading to unstable training. To address this issue, they propose two methods: (1) soft-thresholding the scaling in Real NVPs, and (2) applying a bijective soft-logarithm transform called LOFT to the samples. Experiments show that these modifications enable stable training even for posterior distributions with thousands of dimensions, thereby achieving more accurate estimation of marginal likelihood through importance sampling. In addition, the paper evaluates other architectural adjustments and training techniques, including path gradients and annealing, and provides practical recommendations for optimizing the performance of high-dimensional variational inference with NFs. With these improvements, they achieve more accurate marginal likelihood estimation compared to previous NFs and sequential Monte Carlo methods in high-dimensional logistic regression models. Overall, the paper aims to address the training instability problem when using normalizing flows in high-dimensional variational inference and proposes a method combining soft-thresholding and sample transformation to improve the model's performance in high-dimensional problems.