Reparameterization Gradients through Acceptance-Rejection Sampling Algorithms

Introduction

Variational inference (Hinton and van Camp, 1993; Waterhouse et al., 1996; Jordan et al., 1999) underlies many recent advances in large scale probabilistic modeling. It has enabled sophisticated modeling of complex domains such as images (Kingma and Welling, 2014) and text (Hoffman et al., 2013). By definition, the success of variational approaches depends on our ability to (i) formulate a flexible parametric family of distributions; and (ii) optimize the parameters to find the member of this family that most closely approximates the true posterior. These two criteria are at odds—the more flexible the family, the more challenging the optimization problem. In this paper, we present a novel method that enables more efficient optimization for a large class of variational distributions, namely, for distributions that we can efficiently simulate by acceptance-rejection sampling, or rejection sampling for short.

For complex models, the variational parameters can be optimized by stochastic gradient ascent on the evidence lower bound (elbo), a lower bound on the marginal likelihood of the data. There are two primary means of estimating the gradient of the elbo: the score function estimator (Paisley et al., 2012; Ranganath et al., 2014; Mnih and Gregor, 2014) and the reparameterization trick (Kingma and Welling, 2014; Rezende et al., 2014; Price, 1958; Bonnet, 1964), both of which rely on Monte Carlo sampling. While the reparameterization trick often yields lower variance estimates and therefore leads to more efficient optimization, this approach has been limited in scope to a few variational families (typically Gaussians). Indeed, some lines of research have already tried to address this limitation (Knowles, 2015; Ruiz et al., 2016).

There are two requirements to apply the reparameterization trick. The first is that the random variable can be obtained through a transformation of a simple random variable, such as a uniform or standard normal; the second is that the transformation be differentiable. In this paper, we observe that all random variables we simulate on our computers are ultimately transformations of uniforms, often followed by accept-reject steps. So if the transformations are differentiable then we can use these existing simulation algorithms to expand the scope of the reparameterization trick.

Thus, we show how to use existing rejection samplers to develop stochastic gradients of variational parameters. In short, each rejection sampler uses a highly-tuned transformation that is well-suited for its distribution. We can construct new reparameterization gradients by “removing the lid” from these black boxes, applying $65+$ years of research on transformations (von Neumann, 1951; Devroye, 1986) to variational inference. We demonstrate that this broadens the scope of variational models amenable to efficient inference and provides lower-variance estimates of the gradient compared to state-of-the-art approaches.

We first review variational inference, with a focus on stochastic gradient methods. We then present our key contribution, rejection sampling variational inference (rsvi), showing how to use efficient rejection samplers to produce low-variance stochastic gradients of the variational objective. We study two concrete examples, analyzing rejection samplers for the gamma and Dirichlet to produce new reparameterization gradients for their corresponding variational factors. Finally, we analyze two datasets with a deep exponential family (def) (Ranganath et al., 2015), comparing rsvi to the state of the art. We found that rsvi achieves a significant reduction in variance and faster convergence of the elbo. Code for all experiments is provided at github.com/blei-lab/ars-reparameterization.

Variational Inference

Let $p(x,z)$ be a probabilistic model, i.e., a joint probability distribution of data $x$ and latent (unobserved) variables $z$. In Bayesian inference, we are interested in the posterior distribution $p(z|x)=\frac{p(x,z)}{p(x)}$. For most models, the posterior distribution is analytically intractable and we have to use an approximation, such as Monte Carlo methods or variational inference. In this paper, we focus on variational inference.

In variational inference, we approximate the posterior with a variational family of distributions $q(z\,;\theta)$, parameterized by $\theta$. Typically, we choose the variational parameters $\theta$ that minimize the Kullback-Leibler (kl) divergence between $q(z\,;\theta)$ and $p(z|x)$. This minimization is equivalent to maximizing the elbo (Jordan et al., 1999), defined as

Score function estimator. The score function estimator, also known as the log-derivative trick or reinforce (Williams, 1992; Glynn, 1990), is a general way to estimate the gradient of the elbo (Paisley et al., 2012; Ranganath et al., 2014; Mnih and Gregor, 2014). The score function estimator expresses the gradient as an expectation with respect to $q(z\,;\theta)$:

We then form Monte Carlo estimates by approximating the expectation with independent samples from the variational distribution. Though it is very general, the score function estimator typically suffers from high variance. In practice we also need to apply variance reduction techniques such as Rao-Blackwellization (Casella and Robert, 1996) and control variates (Robert and Casella, 2004).

Reparameterization trick. The reparameterization trick (Salimans and Knowles, 2013; Kingma and Welling, 2014; Price, 1958; Bonnet, 1964) results in a lower variance estimator compared to the score function, but it is not as generally applicable. It requires that: (i) the latent variables $z$ are continuous; and (ii) we can simulate from $q(z\,;\theta)$ as follows,

Here, $s(\varepsilon)$ is a distribution that does not depend on the variational parameters; it is typically a standard normal or a standard uniform. Further, $h(\varepsilon,\theta)$ must be differentiable with respect to $\theta$. In statistics, this is known as a non-central parameterization and has been shown to be helpful in, e.g., Markov chain Monte Carlo methods (Papaspiliopoulos et al., 2003).

Using (2), we can move the derivative inside the expectation and rewrite the gradient of the elbo as

Empirically, the reparameterization trick has been shown to be beneficial over direct Monte Carlo estimation of the gradient using the score fuction estimator (Salimans and Knowles, 2013; Kingma and Welling, 2014; Titsias and Lázaro-Gredilla, 2014; Fan et al., 2015). Unfortunately, many distributions commonly used in variational inference, such as gamma or Dirichlet, are not amenable to standard reparameterization because samples are generated using a rejection sampler (von Neumann, 1951; Robert and Casella, 2004), introducing discontinuities to the mapping. We next show that taking a novel view of the acceptance-rejection sampler lets us perform exact reparameterization.

Reparameterizing the Acceptance-Rejection Sampler

The basic idea behind reparameterization is to rewrite simulation from a complex distribution as a deterministic mapping of its parameters and a set of simpler random variables. We can view the rejection sampler as a complicated deterministic mapping of a (random) number of simple random variables such as uniforms and normals. This makes it tempting to take the standard reparameterization approach when we consider random variables generated by rejection samplers. However, this mapping is in general not continuous, and thus moving the derivative inside the expectation and using direct automatic differentiation would not necessarily give the correct answer.

Our insight is that we can overcome this problem by instead considering only the marginal over the accepted sample, analytically integrating out the accept-reject variable. Thus, the mapping comes from the proposal step. This is continuous under mild assumptions, enabling us to greatly extend the class of variational families amenable to reparameterization.

We first review rejection sampling and present the reparameterized rejection sampler. Next we show how to use it to calculate low-variance gradients of the elbo. Finally, we present the complete stochastic optimization for variational inference, rsvi.

Acceptance-Rejection sampling is a powerful way of simulating random variables from complex distributions whose inverse cumulative distribution functions are not available or are too expensive to evaluate (Devroye, 1986; Robert and Casella, 2004). We consider an alternative view of rejection sampling in which we explicitly make use of the reparameterization trick. This view of the rejection sampler enables our variational inference algorithm in Section 3.2.

To generate samples from a distribution $q(z\,;\theta)$ using rejection sampling, we first sample from a proposal distribution $r(z\,;\theta)$ such that $q(z\,;\theta)\leq M_{\theta}r(z\,;\theta)$ for some $M_{\theta}<\infty$. In our version of the rejection sampler, we assume that the proposal distribution is reparameterizable, i.e., that generating $z\sim r(z\,;\theta)$ is equivalent to generating $\varepsilon\sim s(\varepsilon)$ (where $s(\varepsilon)$ does not depend on $\theta$) and then setting $z=h(\varepsilon,\theta)$ for a differentiable function $h(\varepsilon,\theta)$. We then accept the sample with probability $\min\left\{1,\frac{q\left(h(\varepsilon,\theta)\,;\theta\right)}{M_{\theta}r\left(h(\varepsilon,\theta)\,;\theta\right)}\right\}$; otherwise, we reject the sample and repeat the process. We illustrate this in Figure 1 and provide a summary of the method in Algorithm 1, where we consider the output to be the (accepted) variable $\varepsilon$, instead of $z$.

The ability to simulate from $r(z\,;\theta)$ by a reparameterization through a differentiable $h(\varepsilon,\theta)$ is not needed for the rejection sampler to be valid. However, this is indeed the case for the rejection sampler of many common distributions.

2 The Reparameterized Rejection Sampler in Variational Inference

We now use reparameterized rejection sampling to develop a novel Monte Carlo estimator of the gradient of the elbo. We first rewrite the elbo in (1) as an expectation in terms of the transformed variable $\varepsilon$,

In this expectation, $\pi(\varepsilon\,;\theta)$ is the distribution of the accepted sample $\varepsilon$ in Algorithm 1. We construct it by marginalizing over the auxiliary uniform variable $u$,

where $\mathds{1}[x\in A]$ is the indicator function, and recall that $M_{\theta}$ is a constant used in the rejection sampler. This can be seen by the algorithmic definition of the rejection sampler, where we propose values $\varepsilon\sim s(\varepsilon)$ and $u\sim\mathcal{U}$ until acceptance, i.e., until $u<\frac{q\left(h(\varepsilon,\theta)\,;\theta\right)}{M_{\theta}r\left(h(\varepsilon,\theta)\,;\theta\right)}$. Eq. 3 follows intuitively, but we formalize it in Proposition 1.

Let $f$ be any measurable function, and $\varepsilon\sim\pi(\varepsilon\,;\theta)$, defined by (4) (and implicitly by Algorithm 1). Then

Using the definition of $\pi(\varepsilon\,;\theta)$,

where the second to last equality follows because ${h(\varepsilon,\theta),\varepsilon\sim s(\varepsilon)}$ is a reparameterization of $r(z\,;\theta)$.∎

where we have used the log-derivative trick and rewritten the integrals as expectations with respect to $\pi(\varepsilon\,;\theta)$ (see the supplement for all details.) We define $g_{\text{rep}}$ as the reparameterization term, which takes advantage of gradients with respect to the model and its latent variables; we define $g_{\text{cor}}$ as a correction term that accounts for not using $r(z\,;\theta)\equiv q(z\,;\theta)$.

Using (5), the gradient of the elbo in (1) can be written as

and thus we can build an unbiased one-sample Monte Carlo estimator $\hat{g}\approx\nabla_{\theta}\mathcal{L}(\theta)$ as

where $\varepsilon$ is a sample generated using Algorithm 1. Of course, one could generate more samples of $\varepsilon$ and average, but we have found a single sample to suffice in practice.

Note if $h(\varepsilon,\theta)$ is invertible in $\varepsilon$ then we can simplify the evaluation of the gradient of the log-ratio in $g_{\text{cor}}$,

See the supplementary material for details.

Alternatively, we could rewrite the gradient as an expectation with respect to $s(\varepsilon)$ (this is an intermediate step in the derivation shown in the supplement),

and build an importance sampling-based Monte Carlo estimator, in which the importance weights would be $q\left(h(\varepsilon,\theta)\,;\theta\right)/r\left(h(\varepsilon,\theta)\,;\theta\right)$. However, we would expect this approach to be beneficial for low-dimensional problems only, since for high-dimensional $z$ the variance of the importance weights would be too high.

3 Full Algorithm

We now describe the full variational algorithm based on reparameterizing the rejection sampler. In Section 5 we give concrete examples of how to reparameterize common variational families.

We make use of Eq. 6 to obtain a Monte Carlo estimator of the gradient of the elbo. We use this estimate to take stochastic gradient steps. We use the step-size sequence $\rho^{n}$ proposed by Kucukelbir et al. (2016) (also used by Ruiz et al. (2016)), which combines rmsprop (Tieleman and Hinton, 2012) and Adagrad (Duchi et al., 2011). It is

where $n$ is the iteration number. We set $\delta=10^{-16}$ and $t=0.1$, and we try different values for $\eta$. (When $\theta$ is a vector, the operations above are element-wise.)

We summarize the full method in Algorithm 2. We refer to our method as rsvi.

Related Work

The reparameterization trick has also been used in automatic differentiation variational inference (advi) (Kucukelbir et al., 2015, 2016). advi applies a transformation to the random variables such that their support is on the reals and then places a Gaussian variational posterior approximation over the transformed variable $\varepsilon$. In this way, advi allows for standard reparameterization, but it cannot fit gamma or Dirichlet variational posteriors, for example. Thus, advi struggles to approximate probability densities with singularities, as noted by Ruiz et al. (2016). In contrast, our approach allows us to apply the reparameterization trick on a wider class of variational distributions, which may be more appropriate when the exact posterior exhibits sparsity.

In the literature, we can find other lines of research that focus on extending the reparameterization gradient to other distributions. For the gamma distribution, Knowles (2015) proposed a method based on approximations of the inverse cumulative density function; however, this approach is limited only to the gamma distribution and it involves expensive computations. For general expectations, Schulman et al. (2015) expressed the gradient as a sum of a reparameterization term and a correction term to automatically estimate the gradient in the context of stochastic computation graphs. However, it is not possible to directly apply it to variational inference with acceptance-rejection sampling. This is due to discontinuities in the accept–reject step and the fact that a rejection sampler produces a random number of random variables. Recently, another line of work has focused on applying reparameterization to discrete latent variable models (Maddison et al., 2017; Jang et al., 2017) through a continuous relaxation of the discrete space.

The generalized reparameterization (g-rep) method (Ruiz et al., 2016) exploits the decomposition of the gradient as $g_{\text{rep}}+g_{\text{cor}}$ by applying a transformation based on standardization of the sufficient statistics of $z$. Our approach differs from g-rep: instead of searching for a transformation of $z$ that makes the distribution of $\varepsilon$ weakly dependent on the variational parameters (namely, standardization), we do the opposite by choosing a transformation of a simple random variable $\varepsilon$ such that the distribution of $z=h(\varepsilon,\theta)$ is almost equal to $q(z\,;\theta)$. For that, we reuse the transformations typically used in rejection sampling. Rather than having to derive a new transformation for each variational distribution, we leverage decades of research on transformations in the rejection sampling literature (Devroye, 1986). In rejection sampling, these transformations (and the distributions of $\varepsilon$) are chosen so that they have high acceptance probability, which means we should expect to obtain $g_{\text{cor}}\approx 0$ with rsvi. In Sections 5 and 6 we compare rsvi with g-rep and show that it exhibits significantly lower variance, thus leading to faster convergence of the inference algorithm.

Finally, another line of research in non-conjugate variational inference aims at developing more expressive variational families (Salimans et al., 2015; Tran et al., 2016; Maaløe et al., 2016; Ranganath et al., 2016). rsvi can extend the reparameterization trick to these methods as well, whenever rejection sampling is used to generate the random variables.

Examples of Acceptance-Rejection Reparameterization

As two examples, we study rejection sampling and reparameterization of two well-known distributions: the gamma and Dirichlet. These have been widely used as variational families for approximate Bayesian inference. We emphasize that rsvi is not limited to these two cases, it applies to any variational family $q(z\,;\theta)$ for which a reparameterizable rejection sampler exists. We provide other examples in the supplement.

One of the most widely used rejection sampler is for the gamma distribution. Indeed, the gamma distribution is also used in practice to generate e.g. beta, Dirichlet, and Student’s t-distributed random variables. The gamma distribution, $\operatorname*{Gamma}(\alpha,\beta)$, is defined by its shape $\alpha$ and rate $\beta$.

For $\operatorname*{Gamma}(\alpha,1)$ with $\alpha\geq 1$, Marsaglia and Tsang (2000) developed an efficient rejection sampler. It uses a truncated version of the following reparameterization

When $\beta\neq 1$, we divide $z$ by the rate $\beta$ and obtain a sample distributed as $\operatorname*{Gamma}(\alpha,\beta)$. The acceptance probability is very high: it exceeds $0.95$ and $0.98$ for $\alpha=1$ and $\alpha=2$, respectively. In fact, as $\alpha\to\infty$ we have that $\pi(\varepsilon\,;\theta)\to s(\varepsilon)$, which means that the acceptance probability approaches $1$. Figure 1 illustrates the involved functions and distributions for shape $\alpha=2$.

We now study the quality of the transformation in (10) for different values of the shape parameter $\alpha$. Since $\pi(\varepsilon\,;\theta)\to s(\varepsilon)$ as $\alpha\to\infty$, we should expect the correction term $g_{\text{cor}}$ to decrease with $\alpha$. We show that in Figure 2, where we plot the log-ratio (8) from the correction term as a function of $\varepsilon$ for four values of $\alpha$. We additionally show in Figure 3 that the distribution $\pi(\varepsilon\,;\theta)$ converges to $s(\varepsilon)$ (a standard normal) as $\alpha$ increases. For large $\alpha$, $\pi(\varepsilon\,;\theta)\approx s(\varepsilon)$ and the acceptance probability of the rejection sampler approaches $1$, which makes the correction term negligible. In Figure 3, we also show that $\pi(\varepsilon\,;\theta)$ converges faster to a standard normal than the standardization procedure used in g-rep. We exploit this property—that performance improves with $\alpha$—to artificially increase the shape for any gamma distribution. We now explain this trick, which we call shape augmentation.

2 Dirichlet Distribution

Thus, we make a change of variables to reduce the problem to that of simulating independent gamma distributed random variables,

To showcase this, we study a simple conjugate model where the exact gradient and posterior are available: a multinomial likelihood with Dirichlet prior and Dirichlet variational distribution. In Figure 4 we show the resulting variance of the first component of the gradient, based on simulated data from a Dirichlet distribution with $K=100$ components, uniform prior, and $N=100$ trials. We compare the variance of rsvi (for various shape augmentation settings) with the g-rep approach (Ruiz et al., 2016). rsvi performs better even without the augmentation trick, and significantly better with it.

Experiments

In Section 5 we compared rejection sampling variational inference (rsvi) with generalized reparameterization (g-rep) and found a substantial variance reduction on synthetic examples. Here we evaluate rsvi on a more challenging model, the sparse gamma deep exponential family (def) (Ranganath et al., 2015). On two real datasets, we compare rsvi with state-of-the-art methods: automatic differentiation variational inference (advi) (Kucukelbir et al., 2015, 2016), black-box variational inference (bbvi) (Ranganath et al., 2014), and g-rep (Ruiz et al., 2016).

Data. The datasets we consider are the Olivetti faceshttp://www.cl.cam.ac.uk/research/dtg/attarchive/facedatabase.html and Neural Information Processing Systems (nips) 2011 conference papers. The Olivetti faces dataset consists of $64\times 64$ gray-scale images of human faces in $8$ bits, i.e., the data is discrete and in the set $\{0,\ldots,255\}$. In the nips dataset we have documents in a bag-of-words format with an effective vocabulary of $5715$ words.

We set $\alpha_{z}=0.1$ in the experiments. All priors on the weights are set to $\operatorname*{Gamma}(0.1,0.3)$, and the top-layer local variables priors are set to $\operatorname*{Gamma}(0.1,0.1)$. We use $3$ layers, with $100$, $40$, and $15$ components in each. This is the same model that was studied by Ruiz et al. (2016), where g-rep was shown to outperform both bbvi (with control variates and Rao-Blackwellization), as well as advi. In the experiments we follow their approach and parameterize the variational approximating gamma distribution using the shape and mean. To avoid constrained optimization we use the transform $\theta=\log(1+\exp(\vartheta))$ for non-negative variational parameters $\theta$, and optimize $\vartheta$ in the unconstrained space.

Results. For the Olivetti faces we explore ${\eta\in\{0.75,1,2,5\}}$ and show the resulting elbo of the best one in Figure 5. We can see that rsvi has a significantly faster initial improvement than any of the other methods.The results of g-rep, advi and bbvi where reproduced with permission from Ruiz et al. (2016). The wall-clock time for rsvi is based on a Python implementation (average $1.5$s per iteration) using the automatic differentiation package autograd (Maclaurin et al., 2015). We found that rsvi is approximately two times faster than g-rep for comparable implementations. One reason for this is that the transformations based on rejection sampling are cheaper to evaluate. Indeed, the research literature on rejection sampling is heavily focused on finding cheap and efficient transformations.

For the nips dataset, we now compare the variance of the gradients between the two estimators, rsvi and g-rep, for different shape augmentation steps $B$. In Table 1 we show the minimum, median, and maximum values of the variance across all dimensions. We can see that rsvi again clearly outperforms g-rep in terms of variance. Moreover, increasing the number of augmentation steps $B$ provides even further improvements.

Conclusions

We introduced rejection sampling variational inference (rsvi), a method for deriving reparameterization gradients when simulation from the variational distribution is done using a acceptance-rejection sampler. In practice, rsvi leads to lower-variance gradients than other state-of-the-art methods. Further, it enables reparameterization gradients for a large class of variational distributions, taking advantage of the efficient transformations developed in the rejection sampling literature.

This work opens the door to other strategies that “remove the lid” from existing black-box samplers in the service of variational inference. As future work, we can consider more complicated simulation algorithms with accept-reject-like steps, such as adaptive rejection sampling, importance sampling, sequential Monte Carlo, or Markov chain Monte Carlo.

Christian A. Naesseth is supported by CADICS, a Linnaeus Center, funded by the Swedish Research Council (VR). Francisco J. R. Ruiz is supported by the EU H2020 programme (Marie Skłodowska-Curie grant agreement 706760). Scott W. Linderman is supported by the Simons Foundation SCGB-418011. This work is supported by NSF IIS-1247664, ONR N00014-11-1-0651, DARPA PPAML FA8750-14-2-0009, DARPA SIMPLEX N66001-15-C-4032, Adobe, and the Alfred P. Sloan Foundation. The authors would like to thank Alp Kucukelbir and Dustin Tran for helpful comments and discussion.

References

Appendix A Supplementary Material

Here we formalize the claim in the main manuscript regarding the distribution of the accepted variable $\varepsilon$ in the rejection sampler. Recall that ${z=h(\varepsilon,\theta),~{}\varepsilon\sim s(\varepsilon)}$ is equivalent to $z\sim r(z\,;\theta)$, and that $q(z\,;\theta)\leq M_{\theta}r(z\,;\theta)$. For simplicity we consider the univariate continuous case in the exposition below, but the result also holds for the discrete and multivariate settings. The cumulative distribution function for the accepted $\varepsilon$ is given by

Here, we have applied that ${z=h(\varepsilon,\theta),~{}\varepsilon\sim s(\varepsilon)}$ is a reparameterization of $z\sim r(z\,;\theta)$, and thus

The density is obtained by taking the derivative of the cumulative distribution function with respect to $\varepsilon$,

which is the expression from the main manuscript.

The motivation from the main manuscript is basically a standard “area-under-the-curve” or geometric argument for rejection sampling [Robert and Casella, 2004], but for $\varepsilon$ instead of $z$.

A.2 Derivation of the Gradient

We provide below details for the derivation of the gradient. We assume that $h$ is differentiable (almost everywhere) with respect to $\theta$, and that ${f(h(\varepsilon,\theta))\frac{q(h(\varepsilon,\theta)\,;\theta)}{r(h(\varepsilon,\theta)\,;\theta)}}$ is continuous in $\theta$ for all $\varepsilon$. Then, we have

where in the last step we have identified $\pi(\varepsilon\,;\theta)$ and made use of the log-derivative trick

For invertible reparameterizations we can simplify the evaluation of the gradient of the log-ratio in $g_{\text{cor}}$ as follows using standard results on transformation of a random variable

A.3 Examples of Reparameterizable Rejection Samplers

We show in Table 2 some examples of reparameterizable rejection samplers for three distributions, namely, the gamma, the truncated normal, and the von Misses distributions (for more examples, see Devroye ). We show the distribution $q(z\,;\theta)$, the transformation $h(\varepsilon,\theta)$, and the proposal $s(\varepsilon)$ used in the rejection sampler.

A.4 Reparameterizing the Gamma Distribution

We provide details on reparameterization of the gamma distribution. In the following we consider rate $\beta=1$. Note that this is not a restriction, we can always reparameterize the rate. The density of the gamma random variable is given by

where $\Gamma(\alpha)$ is the gamma function. We make use of the reparameterization defined by

Because $h$ is invertible we can make use of the simplified gradient of the log-ratio derived in Section A.2 above. The gradients of $\log q$ and $-\log r$ are given by

where $\psi(\alpha)$ is the digamma function and