Slice Sampling Reparameterization Gradients

1 Introduction . Probabilistic modeling is a powerful approach to inferring latent structure in complex real-world processes , but often presents computational challenges for inference and further downstream tasks . In many modern probabilistic models , inference is recast as optimization of a probabilistic objective that takes the form of an expectation of a loss function with respect to some distribution . A salient example is variational inference [ 27 , 3 ] , where a simpler distribution is optimized to approximate a posterior distribution by minimizing the Kullback-Leibler ( KL ) divergence to the truth . In particular , Monte Carlo methods for estimating the gradient of the KL with respect to the parameters have extended the use of variational inference to non-conjugate Bayesian models [ 43 , 44 ] and neural networks [ 31 ] . Such probabilistic objectives also appear in a number of other applications in machine learning and computational science , including the generator loss in generative adversarial networks [ 15 ] , computing the sensitivity of expectations under the posterior distribution to prior 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . hyperparameters [ 21 , 19 , 12 ] , computing the Black-Scholes delta in computational finance [ 13 ] , and choosing a design that maximizes the probability of improvement in an experiment [ 56 ] . Typically , the probabilistic objectives are comprised of expectations that can not be computed in closed form , so gradients are often estimated via Monte Carlo sampling [ 39 ] . Two popular classes of Monte Carlo gradient estimators are the score function estimator [ 32 , 14 , 55 , 43 , 44 ] and the pathwise ( or “ reparameterization gradient ” ) estimator [ 23 , 24 , 45 , 31 , 49 ] . Score function gradient estimators are general-purpose , as they apply when the underlying density is differentiable and can be sampled from , even if the cost function is not differentiable . However , they often have high variance and are used with variance reduction techniques ( e.g. , Ranganath et al . [ 44 ] ) . In contrast , reparameterization gradients apply when the loss function is differentiable and samples from the underlying density can be generated by a known deterministic , differentiable transformation of samples from a simpler distribution that does not depend on the model parameters . Typically , reparameterization gradients have lower variance than score function gradients , provided the loss function is sufficiently smooth [ 39 ] . Therefore , developing effective reparameterized gradient estimators has been an active area of research in sensitivity and perturbation analysis [ 23 , 24 , 13 ] and stochastic backpropagation [ 31 , 45 ] . However , reparameterization gradients are primarily limited to distributions with tractable normalizing constants . This precludes their use for complex models of interest such as energy-based models and non-conjugate Bayesian models ( although alternative training methods exist , see e.g . Lawson et al . [ 33 ] ) . Thus , generalizing reparameterization with MCMC to unnormalized distributions is an important direction for developing effective gradient estimators for complicated models . Indeed , some recent work has focused on reparameterization gradients for unnormalized distributions in specialized problems , including Gibbs samplers with reparameterizable conditional sampling steps [ 54 ] and dynamics-based MCMC without accept/reject steps ( e.g. , Salimans et al . [ 50 ] , Dai et al . [ 8 ] ) . However , developing estimators for general unnormalized distributions using reparameterized gradients and MCMC is not straightforward with existing approaches . For instance , not all Gibbs sampling steps are reparameterizable using current methods , due to , e.g. , rejection sampling [ 11 ] or “ Metropoliswithin-Gibbs ” sampling [ 7 ] . Next , dynamics-based MCMC samplers without accept/reject steps are approximate samplers with asymptotic bias . Finally , a key obstacle in applying MCMC methods with accept/reject steps , such as the Metropolis-Hastings ( MH ) [ 22 ] , Metropolis-Adjusted Langevin ( MALA ) [ 16 ] , and Hamiltonian Monte Carlo ( HMC ) [ 42 ] algorithms , is that they do not have differentiable sample paths . An appealing alternative to MCMC methods currently used for reparameterized gradients is slice sampling [ 41 ] , an auxiliary-variable MCMC method that can be applied to unnormalized probability distributions and does not require an accept/reject step or sensitive tuning parameters . Crucially , the lack of an accept/reject step leads to the key observation that for a fixed pseudo-random sequence , the realized slice sampling Markov chain is differentiable with respect to the model parameters . In this work , we develop reparameterization gradients for samples generated from slice sampling that apply to distributions known only up to a normalizing constant . Slice sampling reparameterization gradients are broadly applicable to complicated multivariate distributions , such as energy-based models ( EBMs ) [ 35 ] and non-conjugate Bayesian models . While the generated samples are correlated and the gradient estimates are biased because we simulate from a finite Markov chain , we demonstrate the efficacy of slice sampling reparameterization gradients in simulations , investigating the bias and variance properties of reparameterized slice sampling in comparison with existing gradient estimators . We then show applications of slice sampling reparameterized gradients to a number of problems in machine learning and statistics in the areas of deep generative modeling , approximate inference , and Bayesian sensitivity analysis . 2 Background . In this work we are concerned with probabilistic objectives that take the form of an expectation of a real-valued function ` with respect to a base density pθ ( x ) , i.e. , L ( θ ) = Epθ ( x ) [ ` ( x ) ] . ( 1 ) Typically the expectation can not be computed in closed form , so computing the gradient of this objective with respect to the distributional parameters θ requires stochastic estimates of the gradient of the expectation . Two popular classes of stochastic gradient estimators are score function gradients and reparameterization ( or pathwise ) gradients ; see Mohamed et al . [ 39 ] for a review . 2.1 Continuous reparameterization gradients . Reparameterization gradients are used for problems with a differentiable loss function ` and continuous density pθ . They apply when samples from pθ can be generated by a deterministic transformation fθ of samples from a base distribution p ( ) that does not depend on θ . That is , if ∼ p ( ) then x = fθ ( ) ∼ 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 θ . Assuming that the derivative and integral can be interchanged [ 39 , Section 4.3.1 ] : ∇θL ( θ ) = ∇θEp ( ) [ ` ( fθ ( ) ) ] = Ep ( ) [ ∇θ ` ( fθ ( ) ) ] . ( 2 ) Monte Carlo estimates of the gradient are then computed using samples from p ( ) . Many examples of reparameterization gradients for continuous distributions exist . “ One-liner ” reparameterization gradients use a simple function to transform the base distribution into the desired parameterized distribution [ 31 , 45 , 39 ] . The inverse CDF ( quantile function ) can be used to transform uniform random variables into arbitrary 1D distributions , and gradients can be computed when evaluating and differentiating the inverse CDF are numerically tractable . Additional examples are implicit reparameterization gradients [ 10 ] , doubly reparameterized gradients [ 53 ] , reparameterization of accept/reject sampling [ 40 ] , and generalized reparameterization gradients [ 48 ] . 2.2 Markov chain Monte Carlo via slice sampling . For complicated target distributions , direct sampling via , e.g. , the inverse CDF , is often not feasible and moreover it is often the case that the target density function can only be evaluated up to an unknown normalization constant ; this commonly occurs in Bayesian posterior inference where the marginal likelihood is unavailable . Markov chain Monte Carlo ( MCMC ) algorithms address both of these challenges . An MCMC algorithm is a recipe for constructing a Markov chain that is easy to simulate and that converges in distribution to the target , gaining computational tractability at the cost of initialization bias and correlation between the samples . The most common recipe for constructing an MCMC transition operator is the Metropolis–Hastings algorithm [ 38 ] , which combines an easy-to-simulate proposal distribution with a randomized accept/reject decision that depends only on a ratio of densities , thereby avoiding the need to know the normalization constant . Variations on this theme include Gibbs sampling , in which the proposal is chosen to be the conditional distribution for a subset of the variables , and Hamiltonian Monte Carlo ( HMC ) [ 42 ] , in which the proposal is constructed from a fictive dynamical system . Slice sampling [ 41 ] is an alternative approach to MCMC based on the observation that samples taken uniformly from the volume beneath the target PDF have the correct marginal distribution . Slice sampling defines a Markov transition operator that leaves this uniform distribution invariant , and is example of an auxiliary variable MCMC method ( along with HMC , Swendsen-Wang [ 51 ] , pseudo-marginal MCMC [ 1 ] , and others ) . It is appealing because such transition operators can be constructed without having to identify a proposal distribution and the updates never get “ stuck ” as can happen with Metropolis–Hastings ; it is also generally robust to choices of tuning parameters . Consider a distribution with density pθ ( x ) = 1Z ( θ ) πθ ( x ) where the normalizing constant may not be known . Slice sampling typically proceeds in two steps . Starting from a point xn , a height yn+1 is sampled uniformly beneath the density at xn such that yn+1 ∼ U ( 0 , πθ ( xn ) ) . The height yn+1 defines a “ slice ” through the probability density given by S = { x : yn+1 < πθ ( x ) } . The next point xn+1 is then sampled uniformly from the set S. There are different methods for sampling uniformly from S , including component-wise updates and construction of hyper-rectangles ; here we use random-direction slice sampling to sample from the set S ( 37 , Chapter 29.7 ) . A random direction d is generated from a uniform distribution over directions , and the direction defines a line segment through the slice with endpoints x− and x+ . Finally , a value xn+1 is sampled uniformly between the endpoints . The most common procedure for finding the endpoints , as proposed by Neal [ 41 ] , is to use a reversible “ stepping-out ” procedure followed by “ interval shrinking ” to determine xn+1 . In this work we will not use this stepping out procedure and will instead perform a direct search for the slice boundaries . This direct search is key to the present work , as it eliminates the non-differentiable accept/reject step that otherwise appears in most implementations of slice sampling . 3 Slice sampling reparameterization gradients . 3.1 Random-direction slice sampling with numerical slice endpoints . Here we describe the steps of random-direction slice sampling in more detail . We consider a continuous , unnormalized density πθ ( x ) with parameters θ and start from a point xn ∈ Rd . We sample three random quantities : two uniform random numbers u1 , u2 ∈ [ 0 , 1 ] and a uniform random unit vector d ∈ Rd . The value u1 determines the height of the slice u1πθ ( xn ) . The direction d induces a line through the slice . This line intersects with the density at points where πθ ( xn+αd ) = u1πθ ( xn ) for scalar values α . The intersecting points on the slice S induce the set of α values A : = { α ∈ R : πθ ( xn + αd ) = u1πθ ( xn ) } . ( 3 ) Define the slice endpoints representing the closest intersecting locations on the slice S in the positive and negative directions where πθ ( xn + αd ) = u1πθ ( xn ) as x+ : = xn + α +d , x− : = xn + α −d , where α+ : = min α > 0 A , α− : = max α < 0 A . ( 4 ) The points α+ and α− are identified using numerical root finding ( Appendix A ) . The next sample xn+1 is then taken uniformly between the endpoints x− and x+ : xn+1 = x − + u2 ( x + − x− ) = xn + u2α+d + ( 1− u2 ) α−d . ( 5 ) Importantly , α− and α+ are implicit functions of xn , d , u1 , and θ , denoted by α+ ( xn , d , u1 , θ ) , and α− ( xn , d , u1 , θ ) . Via Equation ( 5 ) , the next sample iterate xn+1 is then a deterministic function of the previous iterate xn , the parameters θ , and the random variables u1 , u2 , and d : xn+1 = xn + u2 α + ( xn , d , u1 , θ ) d + ( 1− u2 ) α− ( xn , d , u1 , θ ) d . ( 6 ) In Figure 1 and later sections , we use to refer to the random draws u1 , u2 , and d. Notably , the periteration cost of identifying the slice endpoints is similar for one-dimensional and higher-dimensional distributions , as the algorithm searches for slice endpoints along a single direction per-iteration regardless of the dimensionality of the sampling space .

Summary: This paper introduces an implementation of slice sampling that is compatible with automatic differentiation. This allows differentiation through slice sampling trajectories in probabilistic loss functions when learning generative models. The authors implement their method for different generative learning tasks to show its utility.