PyTorch Reparameterization

Reparameterization is a way to draw a sample that keeps its gradient path. rsample builds the draw from separate noise, so gradients still reach the parameters.

sample cuts the path

sample draws a value from the distribution. The draw arrives cut off from the mean, with no gradient path back.

torch.manual_seed(0)
mean = torch.tensor(0.0, requires_grad=True)
dist = Normal(mean, 0.5)
draw = dist.sample()
print("Value:", round(draw.item(), 4))
print("requires_grad:", draw.requires_grad)
print("grad_fn:", draw.grad_fn)

rsample keeps the path

rsample redraws with the seed reset, so the noise matches the first draw. This time the draw stays connected to the mean.

torch.manual_seed(0)
draw2 = dist.rsample()
print("Value:", round(draw2.item(), 4))
print("requires_grad:", draw2.requires_grad)
print("Has grad_fn:", draw2.grad_fn is not None)

The gradient reaches the mean

A loss built on the rsample draw can send a gradient home. backward fills the gradient of the mean.

loss = (draw2 - 3.0) ** 2
loss.backward()
print("Mean gradient:", round(mean.grad.item(), 4))

Example

import torch
from torch.distributions import Normal

torch.manual_seed(0)
mean = torch.zeros(1, requires_grad=True)
optimizer = torch.optim.Adam([mean], lr=0.05)
target = 3.0

print("Start mean:", round(mean.item(), 4))
for step in range(150):
    dist = Normal(mean, 0.5)
    draws = dist.rsample((32,))
    loss = ((draws - target) ** 2).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print("Final mean:", round(mean.item(), 3))
print("Target:", target)

Key points

sample: a draw that carries no gradient back to the parameters
Noise: random numbers drawn separately from the parameters
rsample: a draw built from fixed noise plus the parameters, with the path kept
Gradients reach the mean and the scale through an rsample draw
Reparameterization moves the randomness aside, so the draw follows the parameters smoothly

In this example, one mean starts at 0.0 and the target value is 3.0. Each step draws 32 values with rsample from a Normal centred on the current mean. The loss is the mean squared distance from the draws to the target. Adam moves the mean for 150 steps at a learning rate of 0.05. Because the draws come from rsample, every loss carries a gradient back to the mean. The final mean prints 3.027, close to the target of 3.0.

Note

sample still serves evaluation, where no gradient is needed

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch.distributions import Normal

torch.manual_seed(0)
mean = torch.zeros(1, requires_grad=True)
optimizer = torch.optim.Adam([mean], lr=0.05)
target = 3.0

print("Start mean:", round(mean.item(), 4))
for step in range(150):
    dist = Normal(mean, 0.5)
    draws = dist.rsample((32,))
    loss = ((draws - target) ** 2).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print("Final mean:", round(mean.item(), 3))
print("Target:", target)

Exercise

A loss must train a distribution's mean through a random draw. Which draw keeps the gradient path?

sample rsample detach argmax