Source-linked AI summary
Pathwise Derivatives Beyond the Reparameterization Trick
Martin Jankowiak, Fritz Obermeyer
TL;DR
High-variance stochastic gradients can make optimization impractically slow, while pathwise gradients are not straightforward for some commonly used distributions. The paper connects pathwise gradients to optimal transport and develops OMT gradients; it reports favorable Beta results and improved sample efficiency for Gaussian Process regression.
Problem
Exact gradients are often infeasible, and pathwise reparameterization is nontrivial for distributions such as Gamma and Beta.
Source: §1 · p. 1, §2.2 · p. 2Method
The paper connects pathwise gradients to optimal transport, defines OMT gradients, and uses approximate pathwise gradients for distributions not readily handled by reparameterization.
Source: §1 · p. 1, §2.2 · p. 2Results
The OMT estimator performs favorably across the tested Beta(α, α) parameter range, while in Gaussian Process regression it has superior sample efficiency due to lower variance.
Source: §7.1.1 · p. 6, §7.2.2 · p. 8Takeaways & Limitations
Optimal transport offers a perspective for formulating pathwise gradients in difficult cases and identifying a different notion of optimality for multivariate Normal gradients.
Source: §8 · p. 8The paper leaves open when pathwise gradients are preferable to score-function gradients and how the relevant notion of optimality should guide practical use.
Source: §8 · p. 8
Abstract
from arXiv · showhide
We observe that gradients computed via the reparameterization trick are in direct correspondence with solutions of the transport equation in the formalism of optimal transport. We use this perspective to compute (approximate) pathwise gradients for probability distributions not directly amenable to the reparameterization trick: Gamma, Beta, and Dirichlet. We further observe that when the reparameterization trick is applied to the Cholesky-factorized multivariate Normal distribution, the resulting gradients are suboptimal in the sense of optimal transport. We derive the optimal gradients and show that they have reduced variance in a Gaussian Process regression task. We demonstrate with a variety of synthetic experiments and stochastic variational inference tasks that our pathwise gradients are competitive with other methods.
1. Introduction
Because exact gradients are often infeasible, stochastic optimization needs estimators with manageable variance. The paper connects pathwise gradients to optimal transport, defines OMT gradients, and examines their behavior for multivariate Normals and other distributions.
1. Introduction
- High-variance stochastic gradient estimates can make optimization impractically slow, motivating estimators with reduced variance.Exact gradients with respect to the parameters are often infeasible.Source: §1 · p. 1
- The paper shows that the reparameterization trick is one member of a continuous family of gradient estimators and defines the optimal-transport choice as OMT gradients.The paper connects pathwise gradients to tangent fields in optimal transport.Source: §1 · p. 1
- For a Cholesky-parameterized bivariate Normal, the reparameterization velocity field has a large rotational component and is suboptimal in the optimal-transport sense.Figure 1 contrasts it with the optimal-transport field: both move the covariance ellipses alike, but particles flow differently.Source: §1 · p. 1, Figure 1 · p. 2
- The paper also sets out to derive approximate pathwise gradients for Gamma, Beta, and Dirichlet distributions and evaluate its estimators in synthetic and real-world experiments.The introduction identifies stochastic variational inference as one evaluation setting.Source: §1 · p. 1
2. Stochastic Gradient Variational Inference
Stochastic variational inference uses Monte Carlo gradients when expectations and their gradients are not analytically tractable. This section reviews score-function and pathwise estimators, including their applicability and variance trade-offs.
2. Stochastic Gradient Variational Inference
- Stochastic variational inference is a motivating use case because black-box methods need Monte Carlo estimators when conjugacy and other simplifying assumptions are unavailable.The paper uses SGVI as its main testbed in Section 7.Source: §2 · p. 2
2.1. Score Function Estimator
- The score-function estimator applies broadly, including to discrete variables, but typically has high variance that can be mitigated with variance-reduction techniques.The review names Rao-Blackwellization and control variates as examples.Source: §2.1 · p. 2
2.2. Pathwise Gradient Estimator
- Pathwise gradients generally have lower variance than score-function estimates but apply to continuous variables that admit a differentiable reparameterization.For Gamma and Beta distributions, the required shape transformations involve special functions, making the reparameterization trick nontrivial to apply.Source: §2.2 · p. 2
3. Univariate Pathwise Gradients
For a univariate distribution, pathwise gradients can be obtained by differentiating an expectation through its dependence on the distribution parameters. The paper uses implicit differentiation to avoid requiring an explicit inverse CDF, while retaining exact samples under its approximation scheme.
3. Univariate Pathwise Gradients
- The objective is an expectation of a test function under a parameterized distribution, and its pathwise gradient requires differentiating the sample with respect to the parameter.The section identifies the derivative of the CDF with respect to the parameter as the key ingredient.Source: §3 · pp. 2–3
- For Gamma and Beta distributions, the inverse CDF and its parameter derivative do not have simple analytic expressions, complicating direct reparameterization.The section seeks an approach that avoids explicitly introducing the inverse CDF.Source: §3 · p. 3
- Implicit differentiation of the CDF yields the paper’s master formula for univariate pathwise gradients without explicitly introducing the inverse CDF.The formula relies on the fact that the uniform random variable does not depend on the parameter.Source: §3 · p. 3
- Approximating the CDF derivative preserves exact samples from the target distribution; approximation fidelity affects only the accuracy of the resulting sample derivatives.By contrast, an approximate inverse CDF can mismatch the probability assigned to samples and their actual distribution.Source: §3 · p. 3
4. Multivariate Pathwise Gradients
Multivariate pathwise gradients can be viewed as velocity fields that move samples as a distribution changes, with valid fields constrained by the transport equation. Optimal transport selects a minimum-kinetic-energy field, though that choice is nearly variance-optimal only for suitable test functions.
4.1. The Transport Equation
- The transport equation describes how a velocity field displaces samples as parameters change while preserving the changing distribution.The distribution's rate of change is balanced by particle flow into and out of each infinitesimal volume.Source: §4.1 · p. 3
4.2. Gradient Estimator
- A solution to the transport equation yields an unbiased multivariate pathwise gradient estimator.The estimator generalizes the univariate gradient estimator.Source: §4.2 · p. 4
4.3. Tangent Fields
- Because the transport equation admits infinitely many solutions, the optimal-transport tangent field selects a unique solution by minimizing total kinetic energy.The field is the tangent field from optimal transport.Source: §4.3 · p. 4
4.4. Gradient variance
- The OMT gradient is nearly variance-optimal when the test function's gradient is approximately constant across most of the distribution's support.Its kinetic-energy objective weights velocity components equally, whereas gradient variance scales them by the test-function gradient.Source: §4.4 · p. 4
- For generic test functions with relatively flat, isotropic gradients—or little known gradient structure—the OMT field is expected to have lower variance than the reparameterization field.Some specific test functions can still yield a lower-variance estimator than OMT.Source: §4.4 · p. 4
4.5. The Multivariate Normal
- For a zero-mean multivariate Normal parameterized by a Cholesky factor, the reparameterization field is nonoptimal, while the OMT field can reduce gradient-estimator variance.Computing the OMT gradients requires an SVD of the covariance matrix and scales as O(D^3).Source: §4.5 · p. 4
5. Numerical Recipes
The section develops practical approximations for pathwise gradients in Gamma, Beta, and Dirichlet distributions, combining expansion-based methods with efficient approximations where analytic derivatives are unavailable. The Dirichlet construction gives a complete gradient recipe without requiring the same stick-breaking scheme for sampling.
5. Numerical Recipes
- Accurate, efficient approximations to the inverse-CDF derivative provide pathwise gradients for Gamma, Beta, and Dirichlet distributions.The recipes use elementary Taylor expansions, the Lugannani–Rice saddlepoint expansion, and rational polynomial approximations.Source: §5 · pp. 4–5
5.1. Gamma
- For the standard Gamma distribution, six terms in the small-z Taylor expansion are accurate for z < 0.8.The shape-parameter derivative of the incomplete gamma function lacks a simple analytic expression.Source: §5.1 · p. 5
5.2. Beta
- The Beta CDF is a regularized incomplete beta function whose derivatives lack simple analytic expressions, so the paper uses numerical approximations.The approximation details are in the supplementary materials.Source: §5.2 · p. 5
5.3. Dirichlet
- The Dirichlet gradient recipe represents its n components through n − 1 independent Beta variates and respects the simplex constraint.The resulting derivatives sum to zero, and the derivation's stick-breaking construction does not dictate the sampling scheme.Source: §5.3 · p. 5
5.4. Implementation
- Pathwise gradients are plug-and-play: an exact or approximate velocity field can be inserted into an automatic differentiation engine without constructing a surrogate objective.Samples and their test-function values are then differentiable with respect to the distribution parameters.Source: §5.4 · p. 5
6. Related Work
The related work situates pathwise gradients alongside control-variate and partially reparameterized estimators, distinguishing the paper's OMT approach by its transport perspective and deterministic Beta and Dirichlet derivatives. It also notes overlap with contemporaneous work while identifying differences in stochasticity and scope.
6. Related Work
- Control-variate approaches reduce variance through neural baselines, Rao-Blackwellization, adaptive variates, or Taylor expansions, while the paper's multivariate-Normal OMT gradient can also be understood as a control variate.Miller et al. (2017) use Taylor expansions of the test function; scale-gradient variance reduction in their approach can require a multisample estimator in high-dimensional models.Source: §6 · p. 5
- Graves (2016) uses the inverse-CDF derivative for mixture distributions, but the resulting estimator relies on a recursive computation that scales with sample-space dimension.Knowles (2015) uses somewhat imprecise inverse-CDF approximations for Gamma random variables.Source: §6 · pp. 5–6
- G-REP is partially reparameterized, whereas RSVI can drop its score-function-like term at the cost of small bias.For Beta and Dirichlet distributions, rejection sampling can introduce auxiliary randomness; the paper's pathwise derivatives are deterministic for a given z and θ.Source: §6 · pp. 5–6
- Figurnov et al. (2018) derive the inverse-CDF derivative and multivariate generalization, but their Beta and Dirichlet derivatives include additional stochasticity.They do not develop the transport-equation and optimal-transport connection or examine the multivariate Normal in detail.Source: §6 · p. 6
7. Experiments
The experiments validate pathwise-gradient estimators on synthetic distribution tests and real-world variational inference tasks, comparing them with RT, score-function, or RSVI estimators where appropriate. Results range from reduced variance in synthetic tests to nearly identical or improved ELBO performance on real-world tasks.
7.1.1. BETA DISTRIBUTION
- For Beta(α, α) with f(z) = z^3, OMT has lower variance than RSVI across the tested α range, especially at smaller α where the objective has greater curvature.The authors identify RSVI's auxiliary random variable as one reason for its higher variance.Source: §7.1.1 · p. 6
7.1.2. DIRICHLET DISTRIBUTION
- For a 1,995-component Dirichlet, the pathwise gradient performs favorably across tested α0 values; increasing RSVI's boost B brings its variance closer to the pathwise gradient.The ELBO gradients are evaluated at the exact posterior, where the exact gradient is zero.Source: §7.1.2 · pp. 6–7, Figure 4 · p. 7
7.1.3. MULTIVARIATE NORMAL
- The multivariate-Normal synthetic tests examine OMT-to-RT mean-variance ratios for off-diagonal Cholesky elements as their magnitude varies, illustrating variance reduction across test functions.The experiment uses D = 50, with qualitatively similar results reported for different dimensions.Source: Figure 5 · p. 7, §7.1.3 · p. 7
7.2.1. SPARSE GAMMA DEF
- On the Olivetti faces task, OMT and RSVI have nearly identical training ELBOs, suggesting gradient variance is not the limiting factor for this dataset.The comparison uses Sparse Gamma DEF variants with and without Beta random variables.Source: §7.2.1 · p. 7, Figure 6 · p. 8
7.2.2. GAUSSIAN PROCESS REGRESSION
- By iteration 270, OMT reaches the ELBO that RT attains at iteration 500; its iterations are about 1.9x slower, largely canceling the sample-efficiency gain in wall-clock time.Despite this cost, OMT obtains a higher ELBO than RT.Source: §7.2.2 · p. 8
8. Discussion and Future Work
The discussion presents optimal transport as a useful perspective for pathwise gradients and identifies open directions in estimator choice and distribution coverage. It proposes adaptive control-variate selection as a possible lower-cost route to variance reduction.
Discussion
- The paper frames optimal transport as a perspective that helps formulate pathwise gradients beyond assumed-practical cases and motivates a new multivariate-Normal estimator.Source: §8 · p. 8
Open questions
- The authors identify better understanding of optimality and when pathwise gradients are preferable to score-function gradients as an open need for practical method selection.Source: §8 · p. 8
Future work
- Any solution of the transport equation yields an unbiased estimator, so differences between such estimators can serve as control variates; adaptive selection may reduce variance at lower computational cost than OMT.For multivariate Normals, computing OMT costs O(D^3).Source: §8 · p. 8
- The explored multivariate Normal and Dirichlet cases cover only part of the multivariate-distribution space; mixtures and other richer families remain future targets, especially for SGVI.The motivation is to pair rich distributions for complex posterior approximation with families amenable to stochastic optimization.Source: §8 · p. 8
9. Supplementary Materials
The supplementary material derives pathwise gradients from CDF derivatives and illustrates how transport-equation solutions yield unbiased estimators. Its examples show how quantile sensitivity and estimator variance depend on the distribution and chosen velocity field.
9.1.1. EXAMPLE: TRUNCATED UNIT NORMAL
- For a truncated unit Normal, the derivative with respect to the interval endpoint is 0 at z = 0, 1 at z = κ, and interpolates smoothly between them.Moving κ affects tail quantiles more than quantiles in the bulk.Source: §9.1.1 · p. 11
9.1.2. EXAMPLE: UNIVARIATE MIXTURE DISTRIBUTIONS
- When component CDFs or their parameter derivatives are available or approximable, the univariate CDF formula gives pathwise derivatives for mixture-component parameters.The truncated unit Normal also admits the reparameterization trick when erf^-1 can be computed.Source: §9.1.2 · p. 11
- For a two-component univariate Normal mixture, the OMT gradient can have much lower variance than the score-function gradient, while the reparameterization trick is impractical in this setting.The example uses component means (0, 1) and standard deviations (1, 1).Source: §9.1.2 · p. 11
9.2. The Multivariate Case
- A velocity field satisfying the transport equation yields an unbiased multivariate gradient estimator when the relevant surface integral can be dropped.This condition is the multivariate counterpart of the boundary assumption used in the univariate derivation.Source: §9.2 · pp. 11–12
9.3.1. WHITENED COORDINATES
- In whitened coordinates, adding an infinitesimal rotation preserves the transport equation and leaves the gradient unchanged in expectation, but the whitened-coordinate OMT solution minimizes the wrong kinetic-energy functional for the natural-coordinate objective.The derivation also identifies the velocity field as the gradient of a scalar field under the stated symmetry condition.Source: §9.3.1 · p. 12
9.4. Natural Coordinates
In natural coordinates, the paper constructs the OMT velocity field by imposing the transport equation and OMT condition. For the bivariate Normal and a specific linear test-function case, the reported OMT gradient has lower variance than the reparameterization-trick gradient.
9.4. Natural Coordinates
- The natural-coordinate OMT field is determined by imposing the OMT condition on the antisymmetric part and the transport equation on the symmetric part.The symmetric matrix is unique, and the paper describes an SVD-based solution that avoids explicitly constructing a dimension-quartic tensor.Source: §9.4 · p. 13
9.4.1. BIVARIATE NORMAL DISTRIBUTION
- The OMT gradient has lower variance than the reparameterization-trick gradient throughout the tested parameter range for the bivariate Normal.The comparison is for the gradient with respect to the Cholesky factor's off-diagonal element L21, with diagonal elements (1, 1).Source: Figure 10 · p. 13, §9.4.1 · p. 13
9.5. Gradient Variance for Linear Test Functions
- For a unit Normal and linear test function, the OMT estimator's variance is exactly half the reparameterization-trick estimator's in expectation when κ_i are drawn from N(0, 1).For a generic prior over κ_i, the paper expects the variance to be about half; it reports this magnitude as typical beyond the specific example.Source: §9.5 · p. 13
9.6. The Lugannani-Rice Approximation
The paper uses the Lugannani–Rice saddlepoint approximation to approximate CDF derivatives for pathwise gradients where marginal densities are approximately Gaussian. Its Gamma and Beta recipes combine region-specific approximations and report unit-tested relative accuracy.
9.6. The Lugannani-Rice Approximation
- The Lugannani–Rice CDF approximation, differentiated with respect to distribution parameters, underlies approximate Gamma, Beta, and Dirichlet pathwise gradients in regions where marginal densities are approximately Gaussian.The saddlepoint is defined implicitly by K′(ŝ) = z, and numerical evaluation requires care near z = μ to avoid finite-precision instability.Source: §9.6 · p. 14
9.7. Gamma Distribution
- The Gamma recipe partitions (z, α) space among Taylor, Lugannani–Rice, and rational-polynomial approximations, and its dz/dα approximation has unit-tested relative accuracy of 0.0005 over a wide range of inputs.Near the singularity, the Taylor expansion is chosen for greater numerical stability; the Lugannani–Rice expression is more accurate away from it.Source: §9.7 · p. 14
9.8. Beta Distribution
- The Beta recipe divides (z, α, β) space into regions using Taylor-series, Lugannani–Rice, and rational-polynomial approximations, and its dz/dβ approximation has unit-tested relative accuracy of 0.001 over a wide range of inputs.The Lugannani–Rice approximations are described as accurate for α > 6 and β > 6; the boundary parameter ϵ = 0.1 balances numerical accuracy and stability.Source: §9.8 · pp. 14–15
9.9. Dirichlet Distribution
The Dirichlet velocity field satisfies the transport equation but is not optimal in the z−j coordinates; for Student’s t, a Normal–Gamma construction provides a pathwise gradient with a variance trade-off. In the baseball experiment, OMT’s advantage over RT appears when the covariance has strong correlations, while the experimental details specify the settings used across several evaluations.
9.9. Dirichlet Distribution
- Eqn. 66 satisfies the transport equation but is not OMT in z−j coordinates, and the coordinate system in which it is OMT is unclear.The verification checks the transport equation for each j = 1, ..., n.Source: §9.9 · pp. 15–16
9.10. Student’s t-Distribution
- A Student’s t pathwise gradient can be obtained by composing Normal and Gamma gradients, but the auxiliary random degree of freedom raises variance versus direct Eqn. 21.The added variance decreases as ν increases.Source: §9.10 · p. 16, §9.11 · p. 16
9.11. Baseball Experiment
- With L initialized far from identity, OMT learns faster and reaches a higher ELBO than RT; near identity, they perform nearly identically, with OMT costing about 25% more per iteration.The authors expect better sample efficiency when covariance matrices have strong correlations, as in the Gaussian Process experiment.Source: Figure 12 · p. 16, §9.11 · p. 16
9.12. Experimental Details
- The experiments use single-sample gradient estimators, and RSVI includes the score-function term unless otherwise noted.Source: §9.12 · p. 16
9.12.1. MULTIVARIATE NORMAL SYNTHETIC TEST FUNCTION EXPERIMENT
- The D = 50 synthetic multivariate Normal experiment varies off-diagonal Cholesky magnitude r and uses three test functions with analytically computable gradients.Analytic gradients support reliable variance estimates for the gradient estimators.Source: §9.12.1 · p. 17
9.12.2. SPARSE GAMMA DEF
- For sparse Gamma DEF, the authors select learning rates by short-run ELBO grid search, then train for 2,000 iterations with shape augmentation B = 4.They report similar results across random seeds; RSVI’s score-function-like term is included for one model variant but not the other.Source: §9.12.2 · p. 17
9.12.3. GAUSSIAN PROCESS REGRESSION
- For Gaussian Process regression, both OMT and RT use Adam with single-sample gradients; learning rate and β1 are selected by grid search before 500 training iterations.The authors report that other random seeds give similar results.Source: §9.12.3 · p. 17