Implicit Reparameterization Gradients

• Implicit reparameterization gradients are unbiased pathwise derivative estimators that use implicit differentiation of CDFs to compute low-variance stochastic gradients.
• They extend the classic reparameterization trick to support distributions like Gamma, Beta, Dirichlet, and von Mises that lack explicit invertible transformations.
• Empirical evaluations show these methods reduce gradient variance significantly and improve convergence in variational inference and reinforcement learning, despite higher computational costs in high dimensions.

Implicit reparameterization gradients are a class of unbiased pathwise derivatives that enable the computation of low-variance stochastic gradients for expectations with respect to continuous probability distributions, even when those distributions do not admit an explicit invertible and differentiable parameterization. This class of estimators generalizes the classic reparameterization trick by using implicit differentiation applied to standardization mappings, typically involving the cumulative distribution function (CDF) of the underlying distribution. By leveraging implicit differentiation, these gradients can be employed in variational inference, stochastic variational methods, and entropy-regularized actor-critic reinforcement learning for a wide range of distributions such as Gamma, Beta, Dirichlet, and von Mises—beyond the classic (Gaussian, location–scale) setting (Figurnov et al., 2018, Jankowiak et al., 2018, Libera, 2024).

1. Mathematical Foundation and Derivation

Consider the problem of differentiating the expectation of a function $f$ with respect to the parameters $\theta$ of a distribution $p(z;\theta)$: $$\nabla_{\theta}\,\mathbb{E}_{z\sim p(z;\theta)}\bigl[f(z)\bigr].$$ The reparameterization trick expresses $z$ as a deterministic function of a base noise variable $\varepsilon$ and the parameters $\theta$: $z = g(\varepsilon;\theta)$, allowing gradients to flow through sampling. However, for many distributions, such an explicit $g$ does not exist in closed form.

Implicit reparameterization replaces the reliance on the inverse transformation by differentiating the standardization mapping $S(z;\theta) = \varepsilon$, where $\varepsilon$ has a parameter-independent distribution. For univariate continuous distributions, taking $S(z;\theta) = F(z;\theta)$, the CDF, yields $F(z;\theta) \sim \mathrm{Uniform}(0,1)$. Differentiating the identity $F(z(\theta);\theta) = \varepsilon$ with respect to $\theta$, and applying the chain rule, gives: $$\frac{\partial F}{\partial z} \, \frac{dz}{d\theta} + \frac{\partial F}{\partial\theta} = 0 \implies \frac{dz}{d\theta} = -\,\frac{\partial_\theta F(z;\theta)}{p(z;\theta)}.$$ Thus, the gradient estimator is: $$\nabla_{\theta}\,\mathbb{E}[f(z)] = \mathbb{E}\Bigl[\,\nabla_{z}f(z)\;\bigl(-\,\tfrac{\partial_\theta F(z;\theta)}{p(z;\theta)}\bigr) \Bigr].$$ This formula generalizes to the multivariate setting using the continuity (transport) equation; any solution to the transport equation yields an unbiased pathwise estimator (Figurnov et al., 2018, Jankowiak et al., 2018).

2. Extension Beyond the Reparameterization Trick

The classic reparameterization trick applies only when the standardization transform (e.g., inverse CDF or location–scale relationship) is available in closed, differentiable form. Distributions such as Gamma, Beta, and Dirichlet lack such explicit standardizations. Implicit reparameterization gradients overcome this restriction by using implicit differentiation of the CDF or other standardization maps. This enables unbiased, stochastic gradients for any distribution where the CDF and its derivatives with respect to both $z$ and $\theta$ are available or can be numerically approximated (Figurnov et al., 2018, Jankowiak et al., 2018, Libera, 2024).

3. Application to Key Distributions

Gamma Distribution:

Let $z\sim\mathrm{Gamma}(\alpha,1)$. The regularized incomplete Gamma CDF is $$F(z;\alpha) = P(\alpha,z) = \frac{1}{\Gamma(\alpha)}\int_0^z t^{\alpha-1}e^{-t}dt$$. With

$$\partial_z F(z;\alpha) = \frac{z^{\alpha-1} e^{-z}}{\Gamma(\alpha)}, \quad \partial_\alpha F(z;\alpha) = \frac{\partial}{\partial \alpha}P(\alpha, z),$$

the implicit derivative is

$$\frac{dz}{d\alpha} = -\,\frac{\partial_\alpha F(z;\alpha)}{\partial_z F(z;\alpha)}.$$

No closed form exists for $\partial_\alpha P$, but efficient numerical, Taylor, and rational approximations attain relative errors $\sim 10^{-3}$ with 6–10 terms (Figurnov et al., 2018, Jankowiak et al., 2018).

Beta Distribution:

For $z\sim\mathrm{Beta}(\alpha,\beta)$, $F(z;\alpha,\beta) = I_z(\alpha,\beta)$. $$\partial_z F = z^{\alpha-1}(1-z)^{\beta-1}/B(\alpha,\beta)$$

$$\partial_\alpha F = \int_0^z t^{\alpha-1}(1-t)^{\beta-1} \ln t \, dt / B(\alpha,\beta) + \dots$$

and analogously for $\beta$. The gradients are again $- (\partial_\alpha F) / p(z;\alpha,\beta)$, and regions of ($z$, $\alpha$, $\beta$) are split among Taylor, Lugannani–Rice, and rational approximations (Figurnov et al., 2018, Jankowiak et al., 2018, Libera, 2024).

Dirichlet Distribution:

A stick-breaking construction using independent Beta variables yields implicit gradients via chain rule, efficiently computable in $O(D)$ for $D$ dimensions. The full velocity field satisfies the transport equation and delivers pathwise gradients for stick-breaking representations (Jankowiak et al., 2018, Figurnov et al., 2018).

von Mises Distribution:

For circular distributions such as von Mises, the gradient construction is analogous, with the CDF and its derivatives available in Fourier or series expansions (Figurnov et al., 2018).

4. Variance Reduction and Theoretical Properties

Implicit reparameterization gradients typically exhibit significantly lower variance than score-function (REINFORCE) estimators and match or outperform partial-reparameterization methods such as RSVI and G-Rep at practical augmentation settings. For example, in Beta($\alpha,\alpha$) with $f(z)=z^3$, optimal-transport-motivated pathwise gradients achieve variance 2–5× lower than RSVI for $\alpha \in [0.5,10]$, with negligible finite-sample bias ($<10^{-4}$). For Dirichlet models with dimension $D\approx2000$, OMT-based gradients exhibit uniformly lower variance relative to RSVI over all tested augmentation parameters (Jankowiak et al., 2018, Figurnov et al., 2018).

For the multivariate Normal, optimal-transport (OT)-derived gradients reduce variance by about 50% compared to the standard reparameterization trick in typical regions, and up to 80% in highly correlated cases. In Gaussian Process regression, OMT gradients reach comparable evidence lower bound (ELBO) within 1.8× fewer iterations, though with higher per-iteration computational cost due to $O(D^3)$ matrix operations (Jankowiak et al., 2018).

5. Algorithmic Implementation

Achieving implicit reparameterization gradients in practice involves the following workflow:

1. Sample $z^{(i)}$ from $p(z;\theta)$.
2. Compute $f(z^{(i)})$ and $\nabla_z f(z^{(i)})$.
3. Evaluate $\partial_\theta F(z^{(i)};\theta)$ and $\partial_z F(z^{(i)};\theta)$, often via automatic differentiation on numerical CDFs or closed-form expansions.
4. Compute $$\frac{dz}{d\theta} = -(\partial_z F)^{-1}\,\partial_\theta F$$.
5. Accumulate batch gradients as $$\nabla_\theta L = \sum_i \nabla_z f(z^{(i)})\,\frac{dz}{d\theta}$$.

For multivariate distributions, this may require solving triangular or block-diagonal linear systems. Automatic differentiation of special function CDF routines is more precise than finite-difference approximations and facilitates stability in tail regions by operating in log-space and using double-precision arithmetic (Figurnov et al., 2018, Jankowiak et al., 2018, Libera, 2024).

6. Empirical Results and Benchmarks

Empirical studies demonstrate the efficacy and generality of implicit reparameterization gradients:

• Variance and Speed: Gradient variance is 2–3× lower than RSVI at $B=0$, converging only as $B\gg 1$. Implicit methods are 2–3× faster per sample (Figurnov et al., 2018).
• Latent Dirichlet Allocation: Achieved test perplexities $876 \pm 7$ and $896 \pm 6$, outperforming both RSVI ($B \leq 10$) and classic stochastic variational inference ($\sim$964, $\sim$1330). Direct Dirichlet priors support sparsity in topic weights, unlike Logistic-Normal surrogates (Figurnov et al., 2018).
• Variational Autoencoders: For VAEs on MNIST, Gamma, Beta, and von Mises priors/posteriors trained with implicit gradients either matched or outperformed Normal baselines at low dimensions, with diminishing differences at higher dimensions (Figurnov et al., 2018).
• Soft Actor-Critic with Beta Policy: Implicit reparameterization enables Beta policies in SAC for bounded action spaces. On MuJoCo tasks, SAC-Beta (using both automatic-differentiated and OT gradients) matches or exceeds the performance of squashed Gaussian policies, while pure Gaussian policies often fail due to unboundedness. For Ant-v4, Walker2d-v4, Beta-OMT and Beta-AD policies reach returns of 5456±260 and 4523±409, comparable or superior to Tanh-Normal baselines (Libera, 2024).

7. Practical Considerations and Limitations

General implementation requires only the ability to sample from the target distribution and to compute (analytically or numerically) the CDF and its derivatives. For univariate distributions, the additional computational cost is minor, typically constant-time per sample with high accuracy. For multivariate cases with large $D$, especially the full covariance Normal, the method's $O(D^3)$ scaling—due to SVD or solving Sylvester equations—can become significant, and exploiting structure (Kronecker, low-rank) or reverting to classic reparameterization may be necessary (Jankowiak et al., 2018, Libera, 2024).

Numerical stability is critical in the tails; clamping parameters ($\log\alpha,\log\beta$ to $[-20,2]$), clipping $z$ away from boundary points, and using double-precision arithmetic help mitigate underflow/overflow. No extra score-function terms, surrogate objectives, or auxiliary random variables are required, and implementation in autodiff frameworks is direct (“plug-and-play”) (Jankowiak et al., 2018, Libera, 2024).

The primary limitation is the requirement for a tractable and differentiable CDF. Highly singular, truncated, or mixture distributions may require alternative or approximate strategies (e.g., generalized reparameterization with a score function correction) (Figurnov et al., 2018, Jankowiak et al., 2018).

---

References:

• "Implicit Reparameterization Gradients" (Figurnov et al., 2018)
• "Pathwise Derivatives Beyond the Reparameterization Trick" (Jankowiak et al., 2018)
• "Soft Actor-Critic with Beta Policy via Implicit Reparameterization Gradients" (Libera, 2024)