Transport Score Climbing: Variational Inference Using Forward KL and Adaptive Neural Transport

Variational inference often minimizes the “reverse” Kullbeck-Leibler (KL) divergence, D_KL(q||p) from the approximate distribution q to the posterior p. Recent work instead studies the “forward” KL divergence, D_KL(p||q), which unlike reverse KL does not lead to variational approximations that underestimate uncertainty. To optimize the forward KL, these methods leveraged Markov chain Monte Carlo (MCMC) methods to evaluate the intractable expectation with respect to the posterior p. This paper introduces Transport Score Climbing (TSC), a method that optimizes D_KL(p||q) by using Hamiltonian Monte Carlo (HMC). For improved performance the HMC chain is run on a transformed, or warped, space. A function called the transport map performs the transformation by acting as a change-of-variable from the latent variable space. TSC uses HMC samples to dynamically train the transport map while optimizing D_KL(p||q). TSC leverages synergies, where better transport maps lead to better HMC sampling, which then leads to better transport maps. We demonstrate TSC on synthetic and real data, including using TSC to train variational auto-encoders. We find that TSC achieves competitive performance on the experiments.