Differentiating expectations: score-function and pathwise gradients

Before this lesson: lesson 04 (the chain rule), lesson 05 (what autodiff can and cannot see), lesson 10 (expectations as integrals, Monte Carlo, sampling paths such as $x = \mu + \sigma\epsilon$). It helps to know Monte Carlo estimates and their variance from the probability and statistics track; we repeat what we need. Time: about 75 minutes, plus the notebook.

---

Why this matters #

Many objectives in ML are expectations over a distribution that depends on the parameters:

• In a VAE, the encoder parameters define the distribution $q_\phi(z \mid x)$ of the latent code.
• In reinforcement learning, the policy parameters define the distribution of actions and trajectories. RL fine-tuning of language models samples tokens from the model itself.
• Variational inference, the tuning of a stochastic simulator, and every discrete choice (a token, an architecture, a molecule) have the same form.

$F$ almost never has a closed form, so we need an unbiased Monte Carlo estimate of $\nabla F$. The two basic estimators can differ in variance by orders of magnitude, and each one needs different conditions. This choice often decides whether training works at all.

Note the contrast with an ordinary training loss $\mathbb E_{x \sim \text{data}}[\ell(\theta; x)]$, or with the noise in diffusion-model training. There the distribution does not depend on $\theta$, and you just average $\nabla_\theta \ell$. The problem of this lesson starts when $\theta$ changes the distribution.

Notation. In this lesson $\epsilon$ is the random noise of a sampling path (for example $\epsilon \sim N(0, 1)$). It is not the machine epsilon of lesson 07.

1. Predict first #

2. The picture #

The function stays, the weights move. $F(\theta)$ is a weighted average of the values $f(x)$, with weights $p_\theta(x)$. The function $f$ does not depend on $\theta$. When $\theta$ changes, the density moves under $f$, and the average changes.

Interactive figure: loading…

Why you cannot backpropagate through a sample. Look at how a program computes an estimate of $F(\theta)$:

1. Draw $x_1, \dots, x_N$ from $p_\theta$. The program receives $N$ numbers.
2. Compute $\hat F = \frac1N \sum_n f(x_n)$.
3. Call autodiff on $\hat F$.

In step 1 the link between $\theta$ and the numbers $x_n$ is hidden inside the random number generator. Autodiff (lesson 05) differentiates only the operations it records. It sees $f(x_n)$ with constants $x_n$, so it returns $0$ (or an error). The gradient of $F$ is not $0$. We need another way to express how the samples depend on $\theta$. There are two ways.

Way 1: change the weights (the score function). The simplest case is a coin: $x \in \{0, 1\}$ with $P(x = 1) = \theta$. Then $F(\theta) = \theta f(1) + (1 - \theta) f(0)$ and $F'(\theta) = f(1) - f(0)$. The samples are 0 or 1. They cannot move; only their probabilities change. Write the change of each probability relative to the probability itself: $\partial_\theta \log p_\theta(x)$. It is $1/\theta$ for $x = 1$ and $-1/(1-\theta)$ for $x = 0$. Then

The factor $\partial_\theta \log p_\theta(x)$ is the score. A one-sample estimate is: draw $x$, report $f(x)$ times the score. It uses only the value $f(x)$.

Way 2: move the points (the pathwise gradient). The simplest case is a Gaussian location. You can produce $x \sim N(\mu, 1)$ as $x = \mu + \epsilon$ with $\epsilon \sim N(0, 1)$. Keep $\epsilon$ fixed: a change of $\mu$ moves the point $x$ by the same amount. So $F(\mu) = \mathbb E_\epsilon[f(\mu + \epsilon)]$ and $F'(\mu) = \mathbb E_\epsilon[f'(\mu + \epsilon)]$. A one-sample estimate is: draw $\epsilon$, report $f'(\mu + \epsilon)$. It uses the slope of $f$.

Interactive figure: loading…

Same mean, different noise. Take $f(x) = x$, so $F(\mu) = \mu$ and $F'(\mu) = 1$.

• Pathwise: $f'(x) = 1$ for every sample. The variance is $0$.
• Score function: the score of $N(\mu, 1)$ is $x - \mu = \epsilon$. The estimate is $x\epsilon = (\mu + \epsilon)\epsilon = \mu\epsilon + \epsilon^2$. Its mean is $0 + 1 = 1$, correct. Its variance is $\mu^2 \operatorname{Var}(\epsilon) + \operatorname{Var}(\epsilon^2) = \mu^2 + 2$ (the cross term $2\mu\,\mathbb E[\epsilon^3] = 0$).

The score-function estimate multiplies the value of $f$ by a random factor with mean zero. The larger $|f|$ is, the larger the noise, even when the slope of $f$ is the same everywhere.

3. Definitions and main results #

Setup. $\theta \in \R^k$ are the parameters. $p_\theta$ is a density (or a probability mass function) on $\R^d$. $f: \R^d \to \R$ does not depend on $\theta$. $F(\theta) = \mathbb E_{p_\theta}[f(x)]$.

A gradient estimator is a random vector $\hat g$ computed from samples. It is unbiased if $\mathbb E[\hat g] = \nabla F(\theta)$. The average of $N$ independent copies has the same mean and $1/N$ times the covariance. So, at equal cost per sample, the number to compare is the per-sample variance: to reach a standard error $s$ you need about $N = \text{variance}/s^2$ samples.

The score is $s_\theta(x) = \nabla_\theta \log p_\theta(x) \in \R^k$. It is a gradient with respect to the parameters.

The score has mean zero #

Lemma 1. If $p_\theta(x)$ is differentiable in $\theta$ and $\nabla_\theta$ can be moved inside the integral, then $\mathbb E_{p_\theta}[s_\theta(x)] = 0$.

Proof.

1. Where $p_\theta(x) > 0$, the derivative of the log gives $\nabla_\theta p_\theta(x) = p_\theta(x)\, \nabla_\theta \log p_\theta(x)$. This is the log-derivative trick.
2. So $\mathbb E[s_\theta] = \int p_\theta\, \nabla_\theta \log p_\theta\, dx = \int \nabla_\theta p_\theta\, dx$.
3. Swap $\nabla$ and $\int$: this is $\nabla_\theta \int p_\theta\, dx = \nabla_\theta 1 = 0$. $\blacksquare$

In words: the total probability is always 1, so the changes of the probabilities cancel on average. For the coin: $\theta \cdot \frac1\theta + (1-\theta) \cdot \frac{-1}{1-\theta} = 1 - 1 = 0$. The covariance of $s_\theta$ is the Fisher information $I(\theta)$.

The score-function identity #

Theorem 1 (score-function identity; also "likelihood ratio" or "REINFORCE"). Assume

• (i) $p_\theta(x)$ is differentiable in $\theta$ for every $x$;
• (ii) the support $\{x : p_\theta(x) > 0\}$ does not depend on $\theta$;
• (iii) $\nabla_\theta$ and $\int$ can be swapped (for example, some integrable $G(x)$ satisfies $|f(x)|\, \|\nabla_\theta p_\theta(x)\| \le G(x)$ for all $\theta$ near the current value).

Then

Proof.

1. $\nabla_\theta F$ $= \nabla_\theta \int p_\theta(x) f(x)\, dx$ $= \int f(x)\, \nabla_\theta p_\theta(x)\, dx$, by (iii) and because $f$ does not depend on $\theta$.
2. By (ii), every $x$ with $\nabla_\theta p_\theta(x) \ne 0$ has $p_\theta(x) > 0$. There the log-derivative trick gives $\nabla_\theta p_\theta = p_\theta\, \nabla_\theta \log p_\theta$.
3. Hence $\nabla_\theta F$ $= \int p_\theta(x)\, f(x)\, s_\theta(x)\, dx$ $= \mathbb E[f(x)\, s_\theta(x)]$. $\blacksquare$

The estimator is $\hat g = \frac1N \sum_n f(x_n)\, s_\theta(x_n)$ with $x_n \sim p_\theta$. For discrete $x$, replace the integral by a sum; the proof is the same.

What it does not need: $f$ can be discontinuous, a black box, or a reward from a simulator, and $x$ can be discrete. What it needs: samples from $p_\theta$, and the gradient of the normalized log-density. If you only know $\tilde p_\theta = Z(\theta)\, p_\theta$, then $\nabla \log \tilde p_\theta$ is off by $\nabla \log Z(\theta)$.

Why condition (ii) matters. Take $x \sim \mathrm{Uniform}(0, \theta)$ and $f(x) = x$. Then $F = \theta/2$ and $F' = 1/2$. On the support, $\log p_\theta = -\log \theta$, so $s_\theta = -1/\theta$ and $\mathbb E[f s_\theta] = -(\theta/2)/\theta = -1/2$. This has the wrong sign. The Leibniz rule shows what is missing. Write $F(\theta) = \int_0^\theta x \cdot \frac1\theta\, dx$. Its derivative has two parts:

The boundary term comes from the support that moves with $\theta$. The score only sees the interior term.

Interactive figure: loading…

Baselines #

Theorem 2 (baselines keep the estimator unbiased). Let $b$ not depend on $x$: a constant, a function of $\theta$, or an average of other, independent samples. Then $\mathbb E[(f(x) - b)\, s_\theta(x)] = \nabla_\theta F$.

Proof. $\mathbb E[b\, s_\theta(x)] = \mathbb E[b]\, \mathbb E[s_\theta(x)] = 0$, by independence and Lemma 1. $\blacksquare$

So we may subtract any such $b$ from $f$. The mean stays the same; the variance changes. A good $b$ makes $f(x) - b$ small, so the random factor $s_\theta(x)$ is multiplied by a smaller number.

Best constant baseline (scalar $\theta$). The per-sample variance is $V(b) = \mathbb E[(f - b)^2 s_\theta^2] - (F')^2$. It is a parabola in $b$. Set its derivative to zero, $-2\,\mathbb E[(f - b)\, s_\theta^2] = 0$:

In practice $b \approx \mathbb E[f]$ (a running mean, or the mean of the other samples) is close to $b^*$ when $f$ and $s_\theta^2$ are only weakly correlated. Each coordinate of $\theta$ may have its own $b$.

A baseline built from the same samples

With $K$ samples, a common choice is the mean $\bar f$ of all $K$ values as the baseline of sample $i$. This mean contains $f_i$ itself, so it depends on $x_i$. A short computation gives $f_i - \bar f = \frac{K-1}{K}\big(f_i - \bar f_{-i}\big)$, where $\bar f_{-i}$ is the mean of the other $K - 1$ values. So this estimator is $\frac{K-1}{K}$ times the leave-one-out estimator, and its mean is $\frac{K-1}{K}\nabla F$: the direction is right, the length is not. The leave-one-out mean $\bar f_{-i}$ keeps the estimator unbiased (this is the RLOO estimator in RL).

The pathwise gradient #

Theorem 3 (pathwise gradient; the reparameterization trick). Assume

• (i) there is a base distribution $p(\epsilon)$ that does not depend on $\theta$, and a map $g(\epsilon, \theta)$, such that $x = g(\epsilon, \theta)$ with $\epsilon \sim p(\epsilon)$ has the distribution $p_\theta$ (a sampling path);
• (ii) $\theta \mapsto g(\epsilon, \theta)$ is differentiable for each $\epsilon$, and $f$ is differentiable (continuously differentiable is enough; a Lipschitz $f$ also works when $x$ has a density, because such an $f$ is then differentiable at $x$ with probability 1);
• (iii) $\nabla_\theta$ and $\mathbb E_\epsilon$ can be swapped.

Then

where $\partial g/\partial \theta$ is the $d \times k$ Jacobian of $g$ with respect to $\theta$.

Proof.

1. By (i), $x$ and $g(\epsilon, \theta)$ have the same distribution, so $F(\theta) = \mathbb E_\epsilon[f(g(\epsilon, \theta))]$ (the expectation of a function of a random variable, lesson 10).
2. The distribution of $\epsilon$ does not depend on $\theta$. Now $\theta$ appears only inside the integrand. By (iii), $\nabla_\theta F = \mathbb E_\epsilon[\nabla_\theta f(g(\epsilon, \theta))]$.
3. The chain rule (lesson 04) gives $\nabla_\theta f(g(\epsilon, \theta)) = (\partial g/\partial \theta)^\top \nabla_x f(x)$ at $x = g(\epsilon, \theta)$. $\blacksquare$

The trick moves the randomness into $\epsilon$, which does not depend on $\theta$. Then the sample is an ordinary differentiable function of $\theta$, and autodiff can follow it.

Gaussian case. $x = \mu + \sigma \odot \epsilon$ with $\epsilon \sim N(0, I_d)$ and $\theta = (\mu, \sigma)$: $\nabla_\mu = \nabla_x f(x)$ and $\nabla_\sigma = \nabla_x f(x) \odot \epsilon$. With a full covariance $\Sigma = LL^\top$ and $x = \mu + L\epsilon$: $\nabla_L = \nabla_x f(x)\, \epsilon^\top$. These are the sampling maps of lesson 10, but here $g$ need not be invertible.

Why condition (ii) matters. Take $x \sim \mathrm{Exp}(\lambda)$ (density $\lambda e^{-\lambda x}$ for $x > 0$) and the step $f(x) = \mathbf 1[x > 1]$. Then $F(\lambda) = P(x > 1) = e^{-\lambda}$ and $F'(\lambda) = -e^{-\lambda}$. The path is $x = \epsilon/\lambda$ with $\epsilon \sim \mathrm{Exp}(1)$. But $f'(x) = 0$ for every $x \ne 1$, so the pathwise estimate is $0$ with probability 1: it is biased. The jump at $x = 1$ carries the whole gradient, and the chain rule does not see it. The score function works here: $s_\lambda = 1/\lambda - x$, and

The conditions side by side.

	Score function	Pathwise
Needs from $f$	values only	the gradient $\nabla_x f$; no jumps
Needs from $p_\theta$	$\nabla_\theta \log p_\theta$ (normalized); support fixed	a differentiable path $x = g(\epsilon, \theta)$
Type of $x$	discrete or continuous	continuous
Cost per sample	one evaluation of $f$	$f$ plus one backward pass through $f$
Variance	grows with $\lvert f \rvert$ and with the dimension; needs baselines	controlled by the slope of $f$

Variance: why pathwise is usually better, and when it is not #

Pathwise. For $\nabla_\mu$ the estimator is $\nabla_x f(x)$. If $f$ is $L$-Lipschitz, then $\|\nabla f\| \le L$, so the total variance (the trace of the covariance) is at most $\mathbb E\|\nabla f(x)\|^2 \le L^2$. This bound does not depend on the dimension $d$, and adding a constant to $f$ does not change it.

Score function. Two effects make it noisy.

• The size of $f$. For a constant $f = c$ the true gradient is $0$, but the estimator $c\, s_\theta(x)$ has covariance $c^2 I(\theta)$: pure noise that grows like $c^2$. The baseline $b = c$ removes it exactly.
• The dimension. Take $x \sim N(\mu, I_d)$ at $\mu = 0$ and $f(x) = \sum_j x_j$, so $\partial F/\partial \mu_1 = 1$. The score estimate for $\mu_1$ is $f(x)\,\epsilon_1 = \epsilon_1^2 + \sum_{j \ge 2} \epsilon_j \epsilon_1$, with variance $d + 1$. Every coordinate of $x$ adds noise to every coordinate of the gradient. The pathwise estimate is exactly $1$.

When pathwise is worse. The pathwise variance follows the slope of $f$. Take $f(x) = \cos(kx)$ with $x \sim N(\mu, 1)$. Then $F'(\mu) = -k \sin(k\mu)\, e^{-k^2/2}$, which is tiny for large $k$. The pathwise estimate $-k\sin(kx)$ has a variance close to $k^2/2$ ($12.5$ for $k = 5$). The score estimate $\cos(kx)(x - \mu)$ has a variance of at most $\mathbb E[(x - \mu)^2] = 1$ (about $0.5$ for $k = 5$). Large slopes appear in chaotic simulators and in long unrolled computations with exploding gradients. And the pathwise estimator is not usable at all for discrete $x$, for $f$ with jumps, or without a differentiable path.

Why the two estimators agree for a Gaussian: Stein's lemma

For $x \sim N(\mu, \sigma^2)$ the score for $\mu$ is $(x - \mu)/\sigma^2$. Integration by parts (lesson 10, Stein's lemma) gives $\mathbb E\big[f(x)\,(x - \mu)/\sigma^2\big] = \mathbb E[f'(x)]$ when $f\,p_\theta \to 0$ at $\pm\infty$. The left side is the mean of the score-function estimator; the right side is the mean of the pathwise estimator. So the two estimators have the same mean, as they must. Integration by parts moves the derivative from the density to $f$. It does not keep the variance: the two random variables $f(x)(x - \mu)/\sigma^2$ and $f'(x)$ are different.

Answers to section 1. (1) No useful gradient. sample() returns numbers that are not connected to $\mu$, so backward() fails with "element 0 of tensors does not require grad". torch.normal(mu, 1.0) gives the gradient $0$ for the mean. $F$ depends on $\mu$ through the distribution, and this computation never differentiates the distribution. Use rsample() (pathwise) or a score-function surrogate (section 5). (2) The single-sample value changes by $c\, s_\theta(x)$. The mean does not change, because $\mathbb E[c\, s_\theta] = 0$. The variance changes by $2c\, \mathbb E[f s_\theta^2] + c^2\, \mathbb E[s_\theta^2]$, a quadratic in $c$ that is large for large $|c|$.

4. Worked example #

Take $x \sim N(\mu, \sigma^2)$, $f(x) = x^2$, at $\mu = 2$, $\sigma = 1$. The path is $x = \mu + \sigma\epsilon$ with $\epsilon \sim N(0, 1)$. We use the Gaussian moments $\mathbb E[\epsilon] = \mathbb E[\epsilon^3] = \mathbb E[\epsilon^5] = 0$, $\mathbb E[\epsilon^2] = 1$, $\mathbb E[\epsilon^4] = 3$, $\mathbb E[\epsilon^6] = 15$.

Step 1: exact gradient. $F = \mathbb E[x^2] = \mu^2 + \sigma^2$, so $\partial F/\partial\mu = 2\mu = 4$ and $\partial F/\partial\sigma = 2\sigma = 2$.

Step 2: pathwise. $\partial x/\partial\mu = 1$ and $\partial x/\partial\sigma = \epsilon$, and $f'(x) = 2x$.

• For $\mu$: the estimate is $2x = 4 + 2\epsilon$. Mean $4$. Variance $4\,\mathbb E[\epsilon^2] = 4$.
• For $\sigma$: the estimate is $2x\epsilon = 4\epsilon + 2\epsilon^2$. Mean $0 + 2 = 2$. Second moment $16\,\mathbb E[\epsilon^2] + 16\,\mathbb E[\epsilon^3] + 4\,\mathbb E[\epsilon^4]$ $= 16 + 0 + 12$ $= 28$. Variance $28 - 2^2 = 24$.

Step 3: score function. The log-density is $\log p = -\frac{(x - \mu)^2}{2\sigma^2} - \log\sigma - \frac12\log 2\pi$. So $s_\mu = (x - \mu)/\sigma^2 = \epsilon$ and $s_\sigma = (x-\mu)^2/\sigma^3 - 1/\sigma = \epsilon^2 - 1$ (at $\sigma = 1$).

• For $\mu$: the estimate is $x^2 \epsilon = (2 + \epsilon)^2 \epsilon = 4\epsilon + 4\epsilon^2 + \epsilon^3$. Mean $0 + 4 + 0 = 4$.
• Its square is $16\epsilon^2 + 16\epsilon^4 + \epsilon^6 + 32\epsilon^3 + 8\epsilon^4 + 8\epsilon^5$. The odd powers have mean $0$, so the second moment is $16 \cdot 1 + 16 \cdot 3 + 15 + 8 \cdot 3 = 103$. Variance $103 - 16 = 87$.
• For $\sigma$: the estimate is $x^2(\epsilon^2 - 1)$. Mean $2$, variance $346$ (same method, more terms).

Step 4: baselines for $\mu$. With $b = \mathbb E[f] = 5$: $(x^2 - 5)\epsilon = \epsilon^3 + 4\epsilon^2 - \epsilon$. Mean $4$. Second moment $\mathbb E[\epsilon^6] + 16\,\mathbb E[\epsilon^4] + \mathbb E[\epsilon^2] - 2\,\mathbb E[\epsilon^4] = 15 + 48 + 1 - 6 = 58$ (the other cross terms are odd). Variance $58 - 16 = 42$. The best constant is $b^*$ $= \mathbb E[x^2 \epsilon^2]/\mathbb E[\epsilon^2]$ $= \mathbb E[(2 + \epsilon)^2\epsilon^2]$ $= 4 + 3$ $= 7$, with variance $38$.

Step 5: what the numbers mean. For a standard error of $0.1$ on $\partial F/\partial\mu$ you need $N = \text{variance}/0.01$ samples: $400$ (pathwise), $8{,}700$ (score), $4{,}200$ (score with $b = 5$). For other $\mu$ the score variance is $\mu^4 + 14\mu^2 + 15$: it grows like $\mu^4$ (the size of $f$, squared). The pathwise variance stays $4\sigma^2$, independent of $\mu$.

Estimator of $\partial F/\partial\mu$	Mean	Per-sample variance	$N$ for standard error $0.1$
pathwise $2x$	4	4	400
score $x^2(x - \mu)$	4	87	8,700
score, $b = 5$	4	42	4,200
score, $b^* = 7$	4	38	3,800

Interactive figure: loading…

5. In code #

NumPy checks every number of the worked example: exact moments by quadrature, then Monte Carlo with $400{,}000$ samples, then the moving-support example.

import numpy as np
from scipy import integrate

# (a) exact mean and per-sample variance by quadrature over eps ~ N(0, 1)
phi = lambda e: np.exp(-e**2 / 2) / np.sqrt(2 * np.pi)
def mean_var(g):
    m = integrate.quad(lambda e: g(e) * phi(e), -np.inf, np.inf)[0]
    s = integrate.quad(lambda e: g(e)**2 * phi(e), -np.inf, np.inf)[0]
    return round(m, 4), round(s - m**2, 4)

mu = 2.0                                    # x = mu + eps, f(x) = x^2, dF/dmu = 4
for name, g in [("pathwise   ", lambda e: 2 * (mu + e)),
                ("score      ", lambda e: (mu + e)**2 * e),
                ("score, b=5 ", lambda e: ((mu + e)**2 - 5) * e),
                ("score, b=7 ", lambda e: ((mu + e)**2 - 7) * e)]:
    print(name, *mean_var(g))
# pathwise    4.0 4.0
# score       4.0 87.0
# score, b=5  4.0 42.0
# score, b=7  4.0 38.0

# (b) the same estimators from samples
rng = np.random.default_rng(0)
eps = rng.standard_normal(400_000)
x = mu + eps
g_path = 2 * x                  # f'(x) * dx/dmu, with dx/dmu = 1
g_score = x**2 * eps            # f(x) * score, score = (x - mu) / sigma^2 = eps
print(round(g_path.mean(), 2), round(g_path.var(), 1))     # 4.0 4.0
print(round(g_score.mean(), 2), round(g_score.var(), 1))   # 4.01 87.0

# (c) a support that moves: x ~ Uniform(0, theta), f(x) = x, true gradient 1/2
theta = 1.0
u = rng.uniform(size=400_000)
xu = theta * u
print(round(np.mean(xu * (-1 / theta)), 3))   # -0.5: score function, d/dtheta log(1/theta) = -1/theta
print(round(np.mean(1.0 * u), 3))             # 0.5: pathwise, f'(x) * dx/dtheta = 1 * u

In PyTorch, rsample() keeps the path $x = g(\epsilon, \theta)$ in the graph; sample() does not. The score-function estimator is the gradient of a surrogate $f(x)\log p_\theta(x)$, where $f(x)$ enters as a constant:

import torch
from torch.distributions import Normal

torch.manual_seed(0)
n = 400_000
mu = torch.full((n,), 2.0, requires_grad=True)   # one copy of mu per sample
dist = Normal(mu, 1.0)

(dist.rsample() ** 2).sum().backward()         # rsample: x = mu + eps, a differentiable path
g_path = mu.grad.clone()
mu.grad = None

xs = dist.sample()                              # sample: no path, xs carries no gradient
(xs ** 2 * dist.log_prob(xs)).sum().backward()  # surrogate: f(x) enters as a constant weight
g_score = mu.grad.clone()

print(round(g_path.mean().item(), 2), round(g_path.var().item(), 1))    # 3.99 4.0
print(round(g_score.mean().item(), 2), round(g_score.var().item(), 1))  # 3.99 86.2

Two practical points. First, the surrogate $f(x)\log p_\theta(x)$ is not a loss you can read; only its gradient means something. This is the line loss = -m.log_prob(action) * reward in the REINFORCE example of the PyTorch documentation. Second, rsample() exists only for distributions with a differentiable path (has_rsample; true for Normal, Uniform, Gamma, Beta, Dirichlet, false for Bernoulli and Categorical). If $f$ also depends on $\theta$, add the ordinary term: $\nabla_\theta \mathbb E_{p_\theta}[f_\theta(x)] = \mathbb E[\nabla_\theta f_\theta(x)] + \mathbb E[f_\theta(x)\, s_\theta(x)]$.

6. Where it shows up #

• VAEs. The ELBO is $\mathbb E_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \mathrm{KL}(q_\phi(z \mid x)\,\|\,p(z))$. The decoder parameters $\theta$ appear only in $f$: an ordinary gradient. The encoder parameters $\phi$ define the distribution: pathwise, with $z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon$. The Gaussian KL term has a closed form and needs no estimator. See the generative models track.
• Reinforcement learning, and RL for language models. The policy gradient is the score-function estimator on trajectories: $\nabla_\theta \log p_\theta(\tau) = \sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)$. The unknown dynamics of the environment do not depend on $\theta$, so they drop out. This is why RL uses the score function. REINFORCE uses a baseline; RLOO uses the leave-one-out mean (unbiased); GRPO subtracts the group mean, which includes the sample itself, and also divides by the group standard deviation. See the reinforcement learning track.
• Discrete latent variables and discrete choices (tokens, architectures, molecules). There is no differentiable path. Options: the score function with control variates (leave-one-out baselines, VIMCO), or a continuous relaxation (Gumbel-softmax) or the straight-through estimator. The last two have lower variance but are biased.
• Black-box simulators in science. To tune a climate-model parameterization or a seismic source model whose code has no adjoint, use Gaussian smoothing (evolution strategies): $\nabla_\mu \mathbb E_\epsilon[f(\mu + \sigma\epsilon)] = \mathbb E[f(\mu + \sigma\epsilon)\,\epsilon]/\sigma$. This is the score-function estimator of $N(\mu, \sigma^2 I)$. The antithetic form $[f(\mu + \sigma\epsilon) - f(\mu - \sigma\epsilon)]\,\epsilon/(2\sigma)$ acts like a baseline. With a differentiable simulator or emulator, use pathwise gradients and adjoints (ML for science track).
• Probabilistic forecasts. An ensemble model that maps noise to members and is trained with the CRPS (a proper scoring rule) gets its gradient through the noise path: a pathwise estimator.
• Gamma, Beta, Dirichlet. These have no simple path. Implicit reparameterization uses the CDF $C(x; \theta)$ as the standardizing map: from $\epsilon = C(x; \theta)$ with $\epsilon \sim \mathrm{Uniform}(0, 1)$, implicit differentiation (lesson 09) gives $\nabla_\theta x = -\nabla_\theta C(x; \theta)/p_\theta(x)$. PyTorch's Gamma, Beta and Dirichlet have rsample.

7. Common mistakes #

• Sampling without a path. sample() gives no gradient, and torch.normal(mu, ...) gives a zero gradient for the mean. Use rsample() or an explicit $x = \mu + \sigma\epsilon$.
• A baseline that depends on the sample. $b = f(x_i)$, a group mean that includes $x_i$, or an action-dependent baseline in RL changes the mean of the estimator. Use constants, running means or leave-one-out means.
• Two different scores. $\nabla_\theta \log p_\theta(x)$ (parameters; this lesson) is not $\nabla_x \log p(x)$ (data; score matching and diffusion).
• Forgetting the direct term. If $f$ depends on $\theta$ too, $\mathbb E[\nabla_\theta f_\theta]$ must be added.
• Score function with a moving support (uniform bounds, truncation points). It misses the boundary term and is biased.
• Pathwise through a jump (thresholds, argmax, rounding, an indicator such as "rain above 10 mm"). The gradient is zero almost everywhere, not the true gradient.
• An unnormalized $\log p_\theta$ in the score. It misses $-\nabla_\theta \log Z(\theta)$.

8. Check yourself #

1. $x \sim N(\mu, \sigma^2)$ and $f(x) = x$. Give the pathwise and score-function estimates of $\partial F/\partial\mu$ and their variances. What does the baseline $b = \mu$ do?
   Answer

   Pathwise: $f'(x) = 1$, variance $0$. Score: $x(x - \mu)/\sigma^2 = \mu\epsilon/\sigma + \epsilon^2$ (with $x = \mu + \sigma\epsilon$), mean $1$, variance $\mu^2/\sigma^2 + 2$. With $b = \mu$: $(x - \mu)^2/\sigma^2 = \epsilon^2$, mean $1$, variance $2$ for every $\mu$.
2. $x \sim \mathrm{Bernoulli}(\theta)$. Show that $\mathbb E[s_\theta] = 0$ and that $\mathbb E[f(x)\, s_\theta(x)] = f(1) - f(0)$. Why is there no pathwise estimator?
   Answer

   $s_\theta(1) = 1/\theta$ and $s_\theta(0) = -1/(1-\theta)$. $\mathbb E[s] = \theta/\theta - (1-\theta)/(1-\theta) = 0$. $\mathbb E[f s]$ $= \theta f(1)/\theta - (1-\theta) f(0)/(1-\theta)$ $= f(1) - f(0)$. Any path $g(\epsilon, \theta) \in \{0, 1\}$ (for example $\mathbf 1[u < \theta]$ with $u \sim \mathrm{Uniform}(0,1)$) is piecewise constant in $\theta$. Where it is differentiable, its derivative is $0$, so the chain rule gives $0$ instead of $f(1) - f(0)$.
3. $x \sim \mathrm{Exp}(\lambda)$ (density $\lambda e^{-\lambda x}$, $x > 0$) and $f(x) = x$. Derive both estimates of $dF/d\lambda$ and their variances.
   Answer

   $F = 1/\lambda$, $F' = -1/\lambda^2$. Pathwise: $x = \epsilon/\lambda$ with $\epsilon \sim \mathrm{Exp}(1)$, so $\partial x/\partial\lambda = -x/\lambda$. Mean $-1/\lambda^2$, variance $\operatorname{Var}(x)/\lambda^2 = 1/\lambda^4$. Score: $s = 1/\lambda - x$, estimate $x(1/\lambda - x)$. Using $\mathbb E[x^k] = k!/\lambda^k$: mean $1/\lambda^2 - 2/\lambda^2 = -1/\lambda^2$; second moment $\mathbb E[x^2]/\lambda^2 - 2\mathbb E[x^3]/\lambda + \mathbb E[x^4]$ $= (2 - 12 + 24)/\lambda^4$ $= 14/\lambda^4$; variance $13/\lambda^4$, thirteen times larger.
4. $x \sim N(\mu, 1)$. Is the pathwise estimator of $\partial F/\partial\mu$ valid for $f(x) = |x|$? For $f(x) = \mathrm{round}(x)$ (the nearest integer)?
   Answer

   $|x|$: yes. It is Lipschitz, and its kink at $0$ has probability $0$. The estimate is $\mathrm{sign}(x)$, with mean $2\Phi(\mu) - 1 = \frac{d}{d\mu}\mathbb E|x|$ ($\Phi$ is the standard normal CDF). $\mathrm{round}(x)$: no. It jumps by 1 at every half-integer, so $f' = 0$ almost everywhere and the estimate is $0$. The true gradient is $\sum_k \varphi(k + \tfrac12 - \mu) \approx 1$ ($\varphi$ is the standard normal density; each jump adds the density at the jump). The score-function estimator is unbiased for both. Kinks are fine; jumps are not.
5. For each prompt you sample $K = 4$ answers and use as baseline the mean reward of all 4 answers, including the answer itself. Is the policy-gradient estimate unbiased?
   Answer

   No. $f_i - \bar f = \frac34(f_i - \bar f_{-i})$, so the estimate is $\frac34$ times the unbiased leave-one-out estimate: its mean is $\frac34 \nabla F$. The direction is right, but the length is off. The leave-one-out mean $\bar f_{-i}$ fixes it.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. State the score-function (REINFORCE) identity.

   Show answer

2. Why does the score $\nabla_\theta \log p_\theta(x)$ have mean zero?

   Show answer

3. Write the pathwise (reparameterization) gradient for $x = g(\epsilon, \theta)$.

   Show answer

4. How do you reparameterize $x \sim N(\mu, \sigma^2)$, and what are $\partial x/\partial\mu$ and $\partial x/\partial\sigma$?

   Show answer

5. Why does a baseline $b$ not change the mean of the score-function estimator?

   Show answer

6. What is the best constant baseline for a scalar parameter?

   Show answer

7. For $x \sim N(2, 1)$ and $f = x^2$, what are the per-sample variances of the pathwise and score-function estimators of $\partial F/\partial\mu$?

   Show answer

8. Why does the pathwise estimator fail for $f(x) = \mathbf 1[x > 1]$?

   Show answer

9. $x \sim \mathrm{Uniform}(0, \theta)$, $f(x) = x$. What does the score-function estimator give, and why?

   Show answer

10. Which estimator does a VAE encoder use, and which does a policy gradient use?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "What is the reparameterization trick, and why does it reduce variance?" Write $x = g(\epsilon, \theta)$ with $\epsilon$ independent of $\theta$ and differentiate through $f$. The variance follows the slope of $f$, not the value of $f$ times a zero-mean score.
• "Derive REINFORCE. Why does a baseline not add bias? Can it depend on the action?" The log-derivative trick gives $\nabla F = \mathbb E[f\,\nabla\log p]$; $\mathbb E[\nabla \log p] = 0$ makes a constant baseline harmless. An action-dependent baseline adds bias (unless you correct for it).
• "How do you backpropagate through sampling a discrete token?" Score function with a baseline (unbiased, noisy), or Gumbel-softmax and straight-through (biased, lower variance).
• "Your simulator is not differentiable. How do you tune its parameters with gradients?" Gaussian smoothing (evolution strategies, a score-function estimator), or train a differentiable emulator.
• "What is the difference between sample() and rsample()?" rsample() keeps the path $x = g(\epsilon, \theta)$ in the computational graph; sample() does not.

10. Practice #

The notebook for this lesson runs in Colab. It has four exercises, each with a check cell and a folded solution, and a final demo:

1. Write pathwise_grad_samples and score_function_grad_samples for $N(\mu, \sigma^2)$ and check their means against the exact gradient of $\mathbb E[x^2]$ for $\mu$ and $\sigma$.
2. Measure the variance of the score-function estimator as a function of the baseline $b$, and find $b^*$ from samples and from the formula.
3. Discrete samples: the score of $\mathrm{Bernoulli}(\mathrm{sigmoid}(\theta))$ and the gradient of an expected reward, compared with the exact value.
4. The moving support of $\mathrm{Uniform}(0, \theta)$: measure the bias of the score-function estimator and recover the missing boundary term.
5. Demo: minimize $\mathbb E[(x - 3)^2]$ over the mean $\mu$ of a Gaussian with SGD, using the pathwise estimator, the score function, and the score function with a running-mean baseline, and compare the training curves.

11. Watch and read #

DeepMind x UCL | Deep Learning Lectures | 11/12 | Modern Latent Variable ModelsAndriy Mnih · Google DeepMind · 88 min

The part on variational inference compares the REINFORCE (likelihood-ratio) and reparameterization gradient estimators and explains which distributions can be reparameterized.

Variational AutoencodersArxiv Insights · 15 min

Why the sampling step of a VAE blocks backpropagation, and how the reparameterization trick moves the noise out of the way.

CS 285: Lecture 5, Part 1Sergey Levine · UC Berkeley CS 285 · 14 min

The policy gradient derived with the log-derivative trick; the unknown dynamics drop out of the gradient.

Policy Gradient Methods | Reinforcement Learning Part 6Mutual Information · 29 min

REINFORCE with a baseline, and why the baseline reduces the variance without adding bias.

• Mohamed, Rosca, Figurnov and Mnih, Monte Carlo Gradient Estimation in Machine Learning (JMLR 2020): §3 (an intuitive comparison), §4 (score function: derivation, conditions, the $\mathrm{Uniform}(0, \theta)$ example, variance), §5 (pathwise: sampling paths, implicit reparameterization, the Lipschitz variance bound), §7.4 (control variates and baselines), §10.1 (how to choose an estimator).
• Kingma and Welling, An Introduction to Variational Autoencoders: §2.3 (why the gradient with respect to the encoder is the hard part) and §2.4 (the reparameterization trick).
• Schulman, Heess, Weber and Abbeel, Gradient Estimation Using Stochastic Computation Graphs (NeurIPS 2015): how autodiff frameworks combine score-function and pathwise terms with surrogate losses.
• Sutton and Barto, Reinforcement Learning: An Introduction, 2nd edition, free PDF: §13.3 (REINFORCE) and §13.4 (REINFORCE with baseline).
• PyTorch documentation, torch.distributions: the sections "Score function" and "Pathwise derivative".