[Paper Review] The Thermodynamic Variational Objective

Vaden Masrani, Tuan Anh Le+1 more|arXiv (Cornell University)|Jun 28, 2019

Gaussian Processes and Bayesian Inference39 references4 citations

TL;DR

The paper introduces the Thermodynamic Variational Objective (TVO), a tighter lower bound on the log marginal likelihood that unifies and improves upon existing variational inference methods. By leveraging a thermodynamic integration path between the true posterior and variational approximation, TVO enables low-variance, reparameterization-free gradient estimation for both continuous and discrete latent variables, achieving state-of-the-art performance in model and inference network learning.

ABSTRACT

We introduce the thermodynamic variational objective (TVO) for learning in both continuous and discrete deep generative models. The TVO arises from a key connection between variational inference and thermodynamic integration that results in a tighter lower bound to the log marginal likelihood than the standard variational variational evidence lower bound (ELBO) while remaining as broadly applicable. We provide a computationally efficient gradient estimator for the TVO that applies to continuous, discrete, and non-reparameterizable distributions and show that the objective functions used in variational inference, variational autoencoders, wake sleep, and inference compilation are all special cases of the TVO. We use the TVO to learn both discrete and continuous deep generative models and empirically demonstrate state of the art model and inference network learning.

Motivation & Objective

• To address the challenge of low-variance gradient estimation in deep generative models with discrete and non-reparameterizable latent variables.
• To tighten the variational lower bound on the log marginal likelihood without sacrificing broad applicability.
• To unify and generalize existing objectives like ELBO, IWAE, VAE, and Wake-Sleep under a single principled framework.
• To improve both model and inference network learning by avoiding the degradation in performance seen with tighter bounds in standard methods.
• To provide a computationally efficient, reparameterization-free gradient estimator applicable across diverse latent variable types.

Proposed method

• The TVO is derived from the thermodynamic variational identity, which expresses the log marginal likelihood as a 1D integral of the instantaneous ELBO along a geometric path between the model and inference distributions.
• The path is defined by a family of unnormalized densities $\tilde{\pi}_{\beta}(\mathbf{z}) = p_{\theta}(\mathbf{x},\mathbf{z})^{\beta} q_{\phi}(\mathbf{z}|\mathbf{x})^{1-\beta}$, parameterized by $\beta \in [0,1]$.
• The TVO is formed as a Riemann sum approximation of this integral, with $K$ equally spaced points, resulting in a tighter lower bound than the standard ELBO.
• A novel gradient estimator is derived that avoids the reparameterization trick and has lower variance than REINFORCE, enabling effective training for non-reparameterizable distributions.
• The method generalizes existing objectives: the ELBO corresponds to $K=1$, while the importance-weighted autoencoder (IWAE) arises as a special case when $K>1$.
• The approach is applied to both continuous and discrete latent variable models, with empirical evaluation on MNIST and other benchmarks.

Experimental results

Research questions

• RQ1 Can a unified objective be derived that tightens the variational lower bound while remaining applicable to both continuous and discrete latent variables?
• RQ2 How can low-variance gradient estimation be achieved for non-reparameterizable distributions in deep generative models?
• RQ3 Does optimizing a tighter bound like TVO avoid the performance degradation in inference networks observed with IWAE?
• RQ4 Can the connection between thermodynamic integration and variational inference be leveraged to improve model learning and generalization?
• RQ5 Is the TVO capable of achieving state-of-the-art performance in both model and inference network training across diverse architectures?

Key findings

• The TVO provides a tighter lower bound on the log marginal likelihood than the standard ELBO, with the bound becoming arbitrarily tight as the number of integration points $K$ increases.
• The proposed gradient estimator for TVO has empirically lower variance than the REINFORCE estimator and does not require reparameterization, enabling training of discrete and non-reparameterizable models.
• Optimizing the TVO leads to better inference network performance than optimizing the IWAE objective, resolving the known issue where tighter bounds degrade inference network quality.
• The TVO subsumes and generalizes existing objectives: ELBO, VAE, Wake-Sleep, and inference compilation are all special cases of the TVO.
• Empirical results on MNIST show that the TVO achieves state-of-the-art performance in both model likelihood and inference network quality for both continuous and discrete latent variable models.
• The method is effective across diverse architectures, including deep VAEs with normal and Bernoulli likelihoods, and demonstrates robustness in both discrete and continuous settings.