Slice Sampling Reparameterization Gradients
Abstract
Slice sampling is a Markov chain Monte Carlo algorithm for simulating samples from probability distributions, with the convenient property that it is rejection-free.
When the slice endpoints are known, the sampling path is a deterministic function of noise variables since there are no accept-reject steps like those in Metropolis-Hastings algorithms.
Here we describe how to diﬀerentiate the slice sampling path to compute slice sampling reparameterization gradients .
Since slice sampling does not require a normalizing constant, this allows for computing reparameterization gradients of samples from potentially complicated multivariate distributions.
We apply the method in synthetic examples and to ﬁt a variational autoencoder with a conditional energy-based model approximate posterior
1. Introduction
Probabilistic objectives that take the form of an expectation of a function with respect to a base density p θ ( x ), i.e.,
[Equation]
are common in machine learning.
Typically, this expectation cannot be computed in closed form, so optimizing this objective with respect to the distributional parameters θ requires stochastic estimates of the gradient of the expectation.
Stochastic gradients estimators arise in a number of machine learning applications, including optimizing the evidence lower bound (ELBO) in variational inference (Blei et al., 2017), forming the policy gradient in reinforcement learning (Sutton et al., 1998), and optimizing the probability of improvement in experimental design (Wilson et al., 2018).
Two popular classes of gradient estimators are score function gradients and reparameterization (or pathwise) gradients; see Mohamed et al.
for a review.
Reparameterization gradients apply when samples from p θ ( x ) can be generated by a deterministic transformation f θ of samples from a base distribution p ( (cid:15) ).
That is, if ˜ (cid:15) ∼ p ( (cid:15) ) then ˜ x
= f θ (˜ (cid:15) ) ∼ p θ ( x ).
Applying this transformation, the gradient of the objective with respect to θ can then be expressed as an expectation with respect to a distribution that does not depend on θ : under mild regularity conditions,
[Equation]
Monte Carlo estimates of the gradient are then computed using samples from p ( (cid:15) ).
There are many examples of reparameterization gradients (Appendix A), but they typically are restricted to distributions with tractable normalization constants.
Recent work has extended reparameterization gradients to unnormalized probability distributions using reparameterized Markov chain Monte Carlo (MCMC) algorithms, including Gibbs samplers with reparameterizable Gibbs sampling steps (Vahdat et al., 2020) and dynamics-based MCMC samplers without accept/reject steps (e.g., Salimans et {+al. 
 ; Dai et al.+} 2019).
However, not all Gibbs sampling steps are reparameterizable using current methods, and dynamics-based MCMC samplers without accept/reject steps are approximate samplers.
Unfortunately, the sampling paths of MCMC algorithms with accept/reject steps are not deterministic and diﬀerentiable, precluding the use of reparameterization gradients. Instead, score function gradients or specialized methods for discrete variables or accept/reject steps (Naesseth et al., 2016) would need to be used to estimate gradients from such methods.
In this work, we develop reparameterization gradients for samples generated from slice sampling.
Slice sampling (Neal, 2003) is an MCMC algorithm for simulating samples from distributions p θ ( x ) =
π θ ( x ) /Z
( θ ) where the normalizing constant Z ( θ ) may be unknown.
It is often an appealing alternative to the Metropolis-Hastings algorithm, as slice sampling does not require an accept/reject step or sensitive tuning parameters.
The lack of an accept/reject step means that, for a ﬁxed pseudo-random sequence, the realized slice sampling Markov chain is diﬀerentiable with respect to the parameters θ .
Slice sampling reparameterization gradients apply to complicated multivariate distributions such as energy-based models (EBMs, LeCun et al., 2006).
However, the generated samples are correlated and the gradient estimates are biased because we simulate from a ﬁnite Markov chain.
In our experiments, we demonstrate the eﬃcacy of slice sampling reparameterization gradients and use them to ﬁt a conditional EBM as an approximate posterior in a variational autoencoder.
2. Slice sampling reparameterization gradients
Consider a distribution with density p ( x ) =
1 Z π ( x ) where the normalizing constant may not be known.
In slice sampling, samples are simulated from the distribution by uniformly sampling under the surface of the density (Neal, 2003), typically via two steps.
Starting from an initial point x n , a height y n +1 is sampled uniformly beneath the density at x n such that y n +1 ∼ U (0 , π ( x n )).
The height y n +1 deﬁnes a slice through the surface of the probability density given by S
= { x : y n +1 < π ( x ) } .
The next point x n +1 is then sampled uniformly from the set S .
We use random-direction slice sampling to sample from the set S (MacKay, 2002, Chapter 29.7).
A random direction d is generated from a uniform distribution over directions, and the direction deﬁnes a line segment through the slice with endpoints x − and x + .
Finally, a value x n +1 is sampled uniformly between the endpoints.
The most common procedure for ﬁnding the endpoints, as proposed by Neal (2003), is to use a reversible “stepping-out” procedure followed by “interval shrinking” to determine x n +1 ; we do not use these procedures and instead perform a direct search for the slice boundaries.
2.1. Random-direction slice sampling with numerical slice endpoints
Here we describe the steps of random-direction slice sampling in more detail.
We start from a point x n ∈ R d and a continuous, unnormalized density with parameters θ , π θ ( x ).
We sample three random quantities: two uniform random numbers u 1 , u 2 ∈ [0 , 1] and a uniform random unit vector d .
The value u 1 determines the height of the slice u 1 π θ ( x n ).
The direction d deﬁnes a line through the slice. This line intersects with the density at xxx 3 ... forward pass reverse-mode gradients of points where π θ ( x n + α d ) =
u 1 π θ ( x n ) for scalar values α .
These correspond to a set
A =
{ α : π θ ( x n +
[Equation]
d
[Equation]
= u [Equation] [Equation] Two speciﬁc members of this set, α + = min α> 0 A and α − = max α< 0 A , deﬁne the slice endpoints x − = x n + α − d and x + = x n + α + d .
These are the closest locations on the slice in the positive and negative direction where π θ ( x n + α d ) = u 1 π θ ( x n ).
We identify α +
and α − using numerical root ﬁnding (Appendix E).
Then, the next sample x n +1 is sampled uniformly between the endpoints x − and x + such that
[Equation]
We can view α − and α + as implicit functions of x n , d , u 1 , and θ , denoted α + ( x n , d , u 1 , θ ), and α − ( x n , d , u 1 , θ ).
Via equation (4), the location x n +1 is then a deterministic function of x n , θ , the random variables u 1 , u 2 , and d :
[Equation]
In Figure 1 and later sections, we use (cid:15) to refer to {+u 1 , u 2 , and d such that x n +
 = s ( x n , θ, [Equation]+}
2.2. Reverse mode gradients
The slice sampling procedure generates samples x 1: N .
After generating the samples, we evaluate the objective function and use reverse mode automatic diﬀerentiation (AD) to compute gradients (Griewank and Walther, 2008, Section 3.2; Baydin et al., 2017). Here we derive the reverse mode gradients.
Our derivation uses implicit diﬀerentiation to eﬃciently compute gradients of the slice endpoints, avoiding the need to compute gradients through the numerical root ﬁnding algorithm.
Each sample x n +1 depends on the previous samples x 1: n and the parameters θ , as shown by the computational graph (Figure 1).
To compute gradients with respect to the parameters (and initial condition) we therefore require the Jacobian of x n +1 with respect to x n and with respect to θ .
We denote these J x n ( x n +1 ) and J θ ( x n +1 ).
Using equation (4) for x n , these Jacobians can be computed indirectly via adjoints
[Equation]
We use implicit diﬀerentiation to compute ∇ x n [ α + , − ].
The values α are solutions to
[Equation]
Applying implicit diﬀerentiation (Griewank and Walther, 2008, Chap. 15) we get:
[Equation]
Let the loss be a sum over the samples x n such that L ( θ ) =
1 N (cid:80) Nn =1 (cid:96) n ( x n ).
The gradient of the loss with respect to θ is
[Equation]
The gradient ∇ θ L depends on the gradients of the loss with respect to each sample, ∇ x n L .
For the ﬁnal sample x N this gradient is simply ∇ x N L = 1 N ∇ x N (cid:96) n ( x N ).
For earlier samples x n where n = 1 , ..., N − 1, the loss gradients need to be propagated backwards via
[Equation]
After computing {∇ x 1 L, ... ∇ x N L } , the gradient is given by equation (11).
Critically, we compute ∇ θ L without ever fully representing either of the Jacobians by using vectorJacobian products (Appendix B).
We implemented the forward sampling and reverse mode AD in JAX (Bradbury et al., 2018).
2.3. Gradient of the ELBO
We use the reparameterized slice sampler to optimize the ELBO in variational inference with an unnormalized density q φ ( z )
= 1 Z ( φ ) e f φ ( z ) with parameters φ as the approximate posterior.
The reparameterized ELBO is
[Equation]
where z ( φ, (cid:15) ) is a sample generated from the reparameterized slice sampler with unnormalized density f φ ( z ) and base random variables (cid:15) .
Estimating the gradient appears to require the gradient of the normalizing constant Z ( φ ).
However, applying the total derivative (Roeder et al., 2017) yields a path derivative Monte Carlo estimator of the gradient for a sample (cid:15) that does not require the normalizing constant:
[Equation]
See Appendix C for a full derivation.
This estimator has favorable performance, and has zero variance
when p θ ( z | x )
=
q φ ( z | x ) (Roeder et al., 2017).
3. Experiments
[Figure]
Matching a target independent Gaussian
First, we optimize min φ D KL ( q φ || p ) where q and p are both multivariate Gaussians with diagonal covariances.
First, we show that the parameters φ optimized with slice sampling reparameterization gradients converge to the target values for a low-dimensional example with D = 5 (Figure 2).
Next, for D = 20 we compare performance with diﬀerent chain lengths S .
The gradient is computed by backpropagating the loss of the ﬁnal sample.
As S approaches D , slice sampling reparameterization gradients perform competitively with standard reparameterization gradients
[Figure]
Matching a target mixture of Gaussians
Next, we minimize D KL ( q φ || p ) with respect to φ where q and p are 1D mixtures of Gaussians (Figure 3).
We estimated gradients using N independent chains each of length S samples.
While performance is reasonable with S = 1 and N = 1, increasing the number of independent chains and the number of samples improved performance.
In conclusion, for multimodal distributions it may be advantageous to run multiple independent chains or increase the number of samples beyond the dimensionality of the space.
[Figure]
VAE with conditional EBM approximate posterior
We used slice sampling reparameterization gradients to ﬁt a variational autoencoder (VAE) (Kingma and Welling, 2013;
Rezende et al., 2014) with a conditional EBM as the approximate posterior: [Equation]
[Equation]
The observations are x , the latent variables are z ∈ R D , and log p
( z ) is a prior.
The functions f φ ( x ) and g φ ( z ) map to vectors with an embedding dimensionality D e .
This diﬀers from the standard VAE in that the encoder f does not map to the mean of the latent space, but rather a separate embedding that is nonlinearly combined by g ( z ) to give the energy.
The form of the conditional EBM has been studied previously in Khemakhem et al.
We ﬁt the VAE to MNIST data with a Bernoulli log likelihood and with D = 20 and D e = 50.
The gradients were estimated over mini batches of size 64 with one MCMC chain of S = 100 samples for each image.
We ran the optimization forepochs.
Unconditional samples from the generative model look reasonable (Figure 4).
We visualized the embedding of held out test images by the function f using t-SNE (Maaten and Hinton, 2008).
Interestingly, we see clusters in the embedding space corresponding to digits, with some overlap for similar digits.
4. Discussion
We presented slice sampling reparameterization gradients that apply to multivariate distributions without a normalizing constant, and demonstrated the method with synthetic examples and by ﬁtting a conditional EBM approximate posterior.
In future work we will extend our experimental evaluation of the method, investigate methods to improve VAE training, and explore additional applications of the method.
Acknowledgments
We thank Dougal Maclaurin for valuable discussions. This work was partially funded by NSF IIS-2007278.
D. Zoltowski was supported in part by NIH grant T32MH065214. D. Cai was supported in part by a Google Ph.D. Fellowship in Machine Learning.
We focus here on continuous distributions, for which many examples of reparameterization gradients exist. “One-liner” reparameterization gradients use a simple function to transform the base distribution into the desired parameterized distribution (Kingma and Welling, 2013; Rezende et al., 2014; Mohamed et al., 2019). Next, the inverse CDF can be used to transform uniform random variables into arbitrary 1D distributions, and gradients can be computed when evaluating and diﬀerentiating the inverse
CDF are numerically tractable. Additional examples of reparameterization gradients are implicit reparameterization gradients (Figurnov et al., 2018), doubly reparameterized gradients (Tucker et al., 2018), reparameterization of accept/reject sampling (Naesseth et al., 2016), and generalized reparameterization gradients (Ruiz et al., 2016). Other recent work has also described reparameterizing MCMC samplers for gradient estimation. Generative neural samplers are another approach for simulating reparameterized samples from a complicated distribution, although they are not an MCMC algorithm.
In this approach, random variables are passed through a neural network to generate samples from a more complicated distribution. This is the approach taken in the generator of a GAN (Goodfellow et al., 2014). However, this method does not directly apply to generating samples from an unnormalized distribution with a prescribed energy function, since there is not a closed form description of the distribution of the generated samples. Additionally, the parameters of the neural network must be tuned to match the desired distribution through a complicated optimization process.
In contrast, the reparameterized slice sampler directly samples from a distribution with a prescribed energy function without the need for an optimization process.
Variational autoencoders
Variational autoencoders (VAEs) are a class of deep generative models ﬁt with amortized variational inference (Kingma and Welling, 2013; Rezende et al., 2014).
The generative model is
[Equation]
where µ θ and σ θ are neural networks with parameters θ . The model has typically been ﬁt using variational inference with an amortized Gaussian approximate posterior q φ ( z | z ) = N ( µ φ ( x ) , σ φ ( x )) by maximizing the ELBO, although other inference methods and types of approximate posteriors have been studied. In VAEs, it has been observed that the approximate posterior may collapse to the prior distribution (Bowman et al., 2015; Hoﬀman and Johnson, 2016).
A number of papers have proposed methods to solve this problem and others related to VAE training.
One approach is to the alter the generative model.
For example, Dieng et al.
(2017) showed that adding skip connections to the generative model helps avoid posterior collapse and improves the latent variable representation.
We did not investigate if the VAE trained with a conditional EBM approximate posterior improved posterior collapse.
However, we did show interesting structure given by the embedding function f .
Similar to the encoder network of a standard VAE, which maps to the mean of a multivariate Guassian distributions, this function provides a low-dimensional representation of the input data.
It will be interesting to compare the representations of these two approaches in follow up experiments.
Energy-based models
Energy-based models are under increasingly active study.
Grathwohl et al.
(2019) relate classiﬁers with softmax outputs to EBMs, and propose a method for jointly training a classiﬁer and generative model.
Nijkamp et al.
(2019b,a) describe using
SGLD for learning deep EBMs for image data, and provide informative analysis on short run vs. long run MCMC for training EBMs.
We expect that reparamterization gradient methods for unnoramlized distributions will enable additional use cases for EBMs.
Appendix B. Eﬃcient computation
Importantly, we can compute ∇ θ L without ever fully representing either of the two Jacobians J x n ( x n +1 ) and J θ ( x n ).
In the reverse mode gradient accumulation, the Jacobian
J x n ( x n +1 ) is always transposed and then post-multiplied by the gradient vector ∇ x n +1 L .
Using the above formula for the Jacobian, this product is [ J x
[Equation]
Computing the products right to left only involves manipulating vectors.
Next, the transpose of the Jacobian J θ ( x n ) is also always post-multiplied by [ ∇ x n L ] T such that
[Equation]
Appendix C. ELBO Gradient Derivation
Here we discuss optimizing the ELBO with an unnormalized energy-based model (EBM).
Let the approximate posterior be given by an unnormalized density
[Equation]
The reparameterized ELBO is
[Equation]
where z ( φ, (cid:15) ) is a sample generated from the slice sampler with unnormalized density f φ ( z ) and the noise (cid:15) are the set of random variables u 1 , u 2 , and d used to generate the reparameterized samples.
The Monte Carlo estimate of the gradient given a sample (cid:15) is
[Equation]
This appears to require the gradient of the normalizing constant Z ( φ ).
We can estimate this gradient using samples generated from q φ ( z ) via slice sampling.
However, we take a diﬀerent approach.
Applying the total derivative (Roeder et al., 2017), we have
[Equation]
The components ∇ z log Z ( φ ) and ∇ φ log p θ ( x , z ) are zero.
Next, we can ignore ∇ φ log q φ ( z ) since it has mean zero.
With these terms removed, we have the path derivative estimator of the gradient that does not require the normalizing constant:
[Equation]
This estimator has zero variance when p θ ( z | x ) =
q φ ( z | x ) and favorable performance (Roeder et al., 2017).
It may be surprising that we can compute low-variance gradients without the normalization constant.
However, inspection of the pathwise gradient tells we can do just that since the path through z via slice sampling does not depend on the normalization constant.
Unfortunately, evaluation of the full ELBO does require the normalization constant of the approximate posterior, which is no longer available.
Appendix D. Additional experiment
We implemented the simulated inference network example from Rainforth et al.
Section 4.
In this example the generative model is
[Equation]
where x ∈ R D are the observations and z ∈ R D are the latent variables.
We ﬁt the model by optimizing the ELBO with approximate posterior q ( z | z ) =
N ( A x + b, 23 I ).
The parameters are the generative parameters µ ∈ R D and
the inference network parameters A ∈ R D × D and b ∈ R D .
For a dataset { x } 1: N , the optimal parameters (Rainforth et al., 2018) are
[Equation]
We simulated N = 1000 samples from this model and ﬁt the generative and inference network parameters using slice sampling reparameterization gradients.
The dimensionality was D = 20, the mini-batch size as 128, and we computed the slice sampling gradient using
S = 20 samples after a 30 sample burn in.
All parameters (true and ﬁt) were initialized from standard multivariate Gaussian distributions.
After optimizing for 1000 iterations, the ﬁt parameters were very close to the optimal values (Figure 5).
Appendix E. Implementation details
As mentioned in the main text, we implemented the forward slice sampling with numerical root ﬁnding and reverse mode automatic diﬀerentiation in JAX Bradbury et al.
Here we describe three primary components of our implementation.
network experiment (Rainforth et al., 2018)
Root ﬁnding
The forward sampling process for the reparameterized slice sampling algorithm requires numerical root ﬁnding to identify the slice endpoints.
This corresponds to ﬁnding the α +
and α − that satisfy
[Equation]
Importantly, as mentioned above, we are interested in two speciﬁc roots: the roots α − <and α + > 0 that are closest to zero.
We used a standard bisection algorithm to identify these scalar roots.
This corresponded to running two root ﬁnding operations; one bracketed below zero α − ∈ [ − a, δ ] and one bracketed above zero α + ∈ [ δ, a ] where δ is a small positive constant.
Notably, for multimodal distributions there may be multiple roots inside the brackets, and the bisection algorithm will only be guaranteed to ﬁnd one of these roots (rather than the root closest to zero).
Therefore we ﬁrst identiﬁed the bracket value a with a stepping out procedure.
We did this by stepping out a logarithmically until the bracketing condition was satisﬁed.
Once the condition was satisﬁed, we proceeded with the bisection algorithm.
We wrote the root ﬁnder using custom code in JAX such that it could be compiled using jit and batched using vmap .
These two features are critical for speed.
The jit compilation makes the root ﬁnding fast, and vmap makes it easy and fast to batch the root ﬁnding operation across multiple independent sampling chains.
Forward sampling
The forwards sampling process takes as input random noise variables, the current parameters, and a function that computes the log probability of the distribution of interest (up to an additive constant).
It proceeds with the slice sampling steps, using the root ﬁnding implementation described in the previous subsection to identify the roots.
We also implemented this function to be compiled using jit and batched across multiple chains using vmap .
Reverse mode gradients
We implemented the backwards pass using the eﬃcient computations described in B. Again, we implemented this function to be compiled using jit and batched across multiple chains using vmap .