VAE Reparameterisation Trick

Draw a latent sample in a variational autoencoder with the reparameterisation trick.

The encoder of a VAE outputs a mean mu and a log variance log_var for each latent dimension. The decoder needs a sample z ~ N(mu, sigma^2). Sampling is not differentiable, so gradients could not reach the encoder. The trick moves the randomness into an input eps ~ N(0, 1) and writes the sample as a smooth function of mu and log_var:

sigma = exp(0.5 * log_var)
z     = mu + sigma * eps

Now z depends on mu and log_var through plain arithmetic, so backprop works. The encoder predicts the log variance rather than sigma because the log can be any real number, while sigma must stay positive.

Function signature

def reparameterize(mu, log_var, eps):

Example

Input

mu      = np.array([0.0, 1.0])
log_var = np.array([0.0, np.log(4.0)])
eps     = np.array([1.0, -0.5])
reparameterize(mu, log_var, eps)

Output

array([1., 0.])

Reasoning

The first dimension has variance exp(0) = 1, so sigma = 1 and z = 0 + 1 * 1 = 1. The second has variance 4, so sigma = 2 and z = 1 + 2 * (-0.5) = 0.

Constraints and notes

• All three inputs have the same shape, which can be (d,) or (batch, d).
• eps is given, so the output is deterministic.
• Use exp(0.5 * log_var), not sqrt(exp(log_var)). They are equal, but the first is cheaper and does not overflow as early.