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: The authors propose a technique to compute reparameterized gradients w.r.t. parameters of an unnormalized variational density based on slice sampling. Instead of performing slice sampling with the common "stepping out & shrinkage" procedure the authors use a numeric root finding procedure to determine the slice endpoints, which does not require accept/reject steps and is therefore differentiable.