Reparameterization Trick

The canonical application of the reparameterization trick is in training Variational Autoencoders (VAEs). It enables gradient-based optimization of the evidence lower bound (ELBO) by reparameterizing the latent variable z.

• Standard Formulation: Instead of sampling z ~ q_φ(z|x) = N(μ, σ²), the trick expresses it as z = μ + σ ⊙ ε, where ε ~ N(0, I).
• Gradient Flow: This deterministic transformation allows gradients with respect to the distribution parameters φ (mean μ and variance σ) to flow through the sampling operation via backpropagation.
• Impact: This makes it possible to jointly train the encoder (inference network) and decoder (generative network) end-to-end, which was previously intractable.

In differentiable rendering, the trick is used to estimate gradients through Monte Carlo integration of the rendering equation. It transforms the problem of differentiating an expectation into computing the expectation of a derivative.

• Core Mechanism: A rendered pixel's intensity I is an expectation: I = E_{p(ω)}[L(ω)]. Reparameterizing the sampled direction ω = T(θ, ε) makes the expectation differentiable w.r.t. scene parameters θ.
• Enables Inverse Rendering: This allows for gradient-based optimization of geometry, materials (BRDFs), and lighting from multi-view images by minimizing a photometric loss.
• Example: Optimizing a neural signed distance function (SDF) or a NeRF's density and color fields relies on this principle to compute gradients through volumetric sampling.

In reinforcement learning, the trick provides a lower-variance alternative to the REINFORCE estimator for training stochastic policies. It is central to algorithms like the reparameterization trick policy gradient.

• Problem: A policy π_θ(a|s) samples an action a. The score function estimator (REINFORCE) has high variance.
• Solution: If the policy can be reparameterized as a = f_θ(s, ε) with noise ε ~ p(ε), the expected return's gradient becomes a lower-variance estimate: ∇_θ E_{a~π_θ}[Q(s,a)] = E_{ε~p(ε)}[∇_a Q(s, a) ∇_θ f_θ(s, ε)].
• Use Case: This is employed in soft actor-critic (SAC) and other continuous control algorithms for more stable and efficient policy optimization.

The trick enables practical variational inference for Bayesian Neural Networks (BNNs), where weights are treated as probability distributions rather than point estimates.

• Challenge: Performing backpropagation through random weights w ~ q_φ(w) is non-differentiable.
• Application: Weights are reparameterized as w = μ + σ ⊙ ε, where ε is a fixed noise distribution (e.g., Gaussian). This allows gradients for the variational parameters φ = {μ, σ} to be computed.
• Benefit: It facilitates learning a distribution over network weights, quantifying predictive uncertainty—a critical feature for safety-critical applications like autonomous systems or medical diagnosis.

The reparameterization trick is a prerequisite for training normalizing flows, a class of generative models that construct complex distributions through a series of invertible, differentiable transformations.

• Foundation: A flow models a variable z as z = f_θ(x), where x is from a simple base distribution (e.g., Gaussian). The transformation f_θ must be bijective and differentiable.
• Gradient Computation: The trick is inherent because sampling is done by z = f_θ(x) with x ~ p(x). Gradients for the model parameters θ flow directly through the deterministic f_θ.
• Result: This enables exact log-likelihood computation and efficient sampling, used for high-fidelity image generation and density estimation.

The trick is extended to differentiable physics simulation, allowing gradients to propagate through stochastic dynamics for optimal control and system identification.

• Principle: Physical states evolve via s_{t+1} = f_θ(s_t, a_t) + ξ, where ξ is noise. Reparameterizing the noise enables gradient-based optimization of control actions a_t or system parameters θ.
• Applications:
  • Robotics: Training control policies via backpropagation through time with simulated dynamics.
  • Inverse Design: Optimizing material properties or shapes to achieve desired physical behaviors.
• Connection: This is analogous to its use in graphics, bridging the optimization of digital and physical systems.