Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference
We propose a simple and general variant of the standard reparameterized gradient estimator for the variational evidence lower bound. Specifically, we remove a part of the total derivative with respect to the variational parameters that corresponds to the score function. Removing this term produces an unbiased gradient estimator whose variance approaches zero as the approximate posterior approaches the exact posterior. We analyze the behavior of this gradient estimator theoretically and empirically, and generalize it to more complex variational distributions such as mixtures and importance-weighted posteriors.
Recent advances in variational inference have begun to make approximate inference practical in large-scale latent variable models. One of the main recent advances has been the development of variational autoencoders along with the reparameterization trick Welling, 2013, Rezende et al., 2014]. The reparameterization trick is applicable to most continuous latent-variable models, and usually provides lower-variance gradient estimates than the more general REINFORCE gradient estimator [Williams, 1992].
Intuitively, the reparameterization trick provides more informative gradients by exposing the dependence of sampled latent variables z on variational parameters φ. In contrast, the REINFORCE gradient estimate only depends on the relationship between the density function log q φ (z|x, φ) and its parameters.
Surprisingly, even the reparameterized gradient estimate contains the score function-a special case of the REIN-FORCE gradient estimator. We show that this term can easily be removed, and that doing so gives even lower-variance gradient estimates in many circumstances. In particular, as the variational posterior approaches the true posterior, this gradient estimator approaches zero variance faster, making stochastic gradient-based optimization converge and "stick" to the true variational parameters, as seen in figure 1.
The authors analyze the functional form of the reparameterization trick for obtaining the gradient of the ELBO, noting that one term depends on the gradient wrt. the sample z, which they call the path derivative, and the other term does not depend on the gradient wrt sample z, which they (and others) call the score function. The score function has zero expectation, but is not necessarily zero at particular samples of z. The authors propose to delete the score function term from the gradient, and show can be done with a one-line modification in popular autodiff software. They claim that this can reduce the variance—in particular, as the variational posterior approaches the true posterior, the variance in the stochastic gradient approaches zero. They also extent this idea nontrivially to more complicated models.
While this is a useful and easy-to-use trick, particularly in the way it is implemented in autodiff software, this does not appear to be a sufficiently novel contribution. The results about the score function are well-known, and eliminating or rescaling the score function seems like a natural experiment to do.
The discussion of control variates in lines 124-46 is confusing and seemingly contradictory. From my understanding, setting the control variate scale c = 1 corresponds to including the score function term, i.e. using the ordinary ELBO gradient. Their method proposes to not use a control variable, i.e. set c = 0. But lines 140-6 imply that they got better experimental results setting c = 1, i.e. the ordinary ELBO gradient, than with c = 0, i.e. the ELBO gradient without the score function term. This contradicts their experimental results. Perhaps I have misunderstood this section, but I’ve read it several times and still can’t come to a sensible interpretation.
The experimental results show only marginal improvements, though they are essentially free.