LOFT - Stable Training of Normalizing Flows for Variational Inference

Daniel Andrade

Keywords: Normalizing Flows, Bayesian Inference, Stochastic Gradients

TL;DR: Stable training of normalizing flows to handle high-dimensional / fat-tail posterior distributions.

Abstract:

Variational Inference (VI) with Normalizing Flows (NFs) is an increasingly popular alternative to MCMC methods.
However, despite recent progress on stabilizing the variance of stochastic gradient descent during training, we observe that convergence is still difficult to achieve in practice.
In particular, if the target distribution's dimension is high or exhibits fat tails, convergence of NFs fail and only the much simpler Gaussian mean-field VI converges.
As a remedy, we introduce the log soft extension (LOFT) layer, which can effectively restrain the samples of NFs to lie in a reasonable range.
For various different target distributions with high-dimensions or fat tails, we observe that LOFT enables successful training of NFs that was previously not possible.
Moreover, the computational overhead of the LOFT layer is only marginal. Therefore, we expect that LOFT becomes a new standard tool for training deep NFs for Bayesian inference.

Submission Number: 21