Skip to content

← JAX Stochasticity step 16 of 25

Medium Primitives

Reparameterization Gradient

Jump to the task

Why this matters

A Monte Carlo estimate replaces an expectation with an average over samples. Reparameterization writes each Gaussian sample as z = mu + exp(log_sigma) * eps, with standard normal noise independent of the parameters. Autodiff can then differentiate the sample average while holding that noise fixed.

This exercise returns a gradient estimate from 1,000 samples. It does not return the exact gradient of the expectation. Pathwise estimates can have lower variance than score-function estimates, but this is not a universal guarantee for every objective or estimator.

Worked mini-example

For fixed noise values, the loss is mean(z ** 2). Its derivatives are:

  • With respect to mu: 2 * mean(z).
  • With respect to log_sigma: 2 * mean(z * sigma * eps).
import jax
import jax.numpy as jnp

eps = jnp.array([-1.0, 1.0])  # Simple illustrative noise, not the test draw.

def loss(mu, log_sigma):
    z = mu + jnp.exp(log_sigma) * eps
    return jnp.mean(z ** 2)

gradients = jax.grad(loss, argnums=(0, 1))(2.0, 0.0)
# (4.0, 2.0): z = [1, 3], sigma = 1.

For the full Gaussian distribution, E[z ** 2] = mu ** 2 + sigma ** 2. Its exact gradients are 2 * mu and 2 * sigma ** 2. A finite noise sample usually gives different results. In particular, its mean need not be zero and its mean square need not be one.

Problem

Implement reparam_grad(seed, mu, log_sigma).

Create jax.random.PRNGKey(int(seed)) and draw exactly 1,000 standard normal values with jax.random.normal(key, (1000,)). Using that noise, differentiate mean((mu + exp(log_sigma) * eps) ** 2) with respect to both parameters. Return a float32 array of shape (2,) containing [grad_mu, grad_log_sigma] in that order.

seed is a nonnegative integer-valued scalar supplied as a number. mu and log_sigma are finite floating-point scalars. Assume the scale, loss and gradients remain representable with the default float32 random configuration. The expected outputs use the specified noise draw, so substituting the exact expectation's gradient will not pass.

Common pitfalls

  • Drawing noise outside the loss makes its role clear, but is not required for autodiff. Drawing it inside with the same fixed, parameter-independent key also works. No derivative through the key is needed.
  • Use jax.grad(loss, argnums=(0, 1)) to differentiate both parameters. The default only differentiates the first argument.
  • The chain rule contributes d sigma/d log_sigma = sigma. The finite-sample derivative includes a term involving mu * eps; it is not generally just twice the variance.
  • If wrapping the function in jax.jit, mark seed static because the required Python int(seed) conversion needs a concrete value.