Medium End-to-End

Train VAE End-to-End

Implement train_vae_step: take one SGD step for a small variational autoencoder (VAE), computing both weight gradients by hand.

The encoder produces a mean and log variance for a latent Gaussian. Sample a latent vector from that Gaussian, decode it, and train the reconstruction together with a KL penalty toward a standard normal distribution.

Inputs and return value

x: (N, d), the input batch and reconstruction targets.
w_enc: (d, 2*d_z), the encoder matrix.
w_dec: (d_z, d), the decoder matrix.
lr: a finite, nonnegative learning rate.
seed: an integer controlling the latent noise draw.

N, d and d_z are positive. Shapes are consistent, and floating-point tensors are finite and share a dtype and device. The encoder and decoder have no biases, hidden layers or output activations. Leave all inputs unchanged.

Return the updated encoder matrix flattened in row-major order, followed by the updated decoder matrix flattened the same way. The result is a vector of length d2d_z + d_z*d, preserving dtype and device. Return updated weights, not loss, latent vectors or reconstructed inputs. With lr=0, return the original weights concatenated in this order.

Encode, sample and decode

h = x @ w_enc
mu = h[:, :d_z]
log_var = h[:, d_z:]
sigma = exp(0.5 * log_var)
z = mu + sigma * eps
x_hat = z @ w_dec

log_var is the natural logarithm of the variance. Multiplication by 0.5 before exponentiating converts it to a standard deviation. Draw one noise array eps of shape (N, d_z), with independent standard normal entries. Use the same sampled eps throughout this forward and backward pass.

In PyTorch

Use a fresh CPU generator. Draw in float32, then convert to the input dtype and device:

gen = torch.Generator(device="cpu").manual_seed(int(seed))
eps = torch.randn(N, d_z, generator=gen, dtype=torch.float32,
                  device="cpu").to(device=x.device, dtype=x.dtype)

In JAX

The runtime uses Threefry with partitionable mode enabled. Draw in float32 before converting to the input dtype:

key = jax.random.PRNGKey(int(seed))
eps = jax.random.normal(key, (N, d_z), dtype=jnp.float32).astype(x.dtype)

The frameworks have separate expected outputs because their random generators produce different samples from the same seed. Follow the corresponding recipe; do not change global random state. Numerical results are checked with tolerance.

Loss and reductions

Use exactly this reconstruction-plus-KL objective:

recon_loss = sum((x_hat - x)**2) / N
kl_loss = 0.5 * sum(mu**2 + exp(log_var) - 1 - log_var) / N
loss = recon_loss + kl_loss

Each sum covers all elements of its array. This averages over examples while summing over input features or latent coordinates. Do not divide the reconstruction term by d or the KL term by d_z, and do not add a factor of one half to the reconstruction term.

The KL term is the analytic divergence from the encoder's diagonal Gaussian to a standard normal. It does not use the sampled noise. Reconstruction uses the single sampled latent vector for each example.

Manual backward pass

Treat eps as fixed data when differentiating; z still depends on mu and log_var. Start at the decoder:

g_x_hat = 2 * (x_hat - x) / N
grad_w_dec = z.T @ g_x_hat
g_z = g_x_hat @ w_dec.T

Add the reconstruction and KL contributions for each encoder output:

g_mu = g_z + mu / N
g_log_var = g_z * 0.5 * sigma * eps + (exp(log_var) - 1) / (2*N)
g_h = concat(g_mu, g_log_var, axis=-1)
grad_w_enc = x.T @ g_h

Use the original decoder matrix in g_z. Compute both gradients before updating either matrix, then apply new_weight = weight - lr * gradient to both. The batch averaging is already present in these formulas; do not divide the weight gradients by N again.

Use basic tensor operations and manual derivatives in either framework. Do not use automatic differentiation, backward(), jax.grad or optimizer helpers.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Encode a distribution over latent values

A variational autoencoder (VAE) usually predicts a Gaussian distribution for each input. For a diagonal Gaussian, the encoder outputs a mean mu and log variance ell for each latent coordinate. The standard deviation is exp(ell / 2). Using exp(ell) would give the variance instead.

To sample a latent vector, draw independent standard normal noise epsilon and compute

$$z=\mu+\sigma\epsilon,\qquad\sigma=\exp(\ell/2).$$

This is the reparameterization trick. For a fixed noise draw, gradients can flow through the mean and scale: dz/dmu = 1 and dz/dell = 0.5 * sigma * epsilon. The random draw itself does not depend on the encoder parameters.

Balance reconstruction with a prior

The decoder tries to reconstruct the input from the sample. A second term keeps the encoder's distribution close to a chosen prior, commonly a standard normal distribution. Training usually minimizes a reconstruction loss plus this KL-divergence penalty, equivalent to minimizing the negative evidence lower bound when the likelihood and weighting match.

For a diagonal Gaussian and standard normal prior, the KL penalty is

$$\frac12\sum_j\left(\mu_j^2+e^{\ell_j}-1-\ell_j\right).$$

At mu = 0 and ell = 0, the distributions match and the penalty is zero. Squared-error reconstruction corresponds to a Gaussian likelihood with an appropriate scale. Follow the task's specified loss and reduction, since these set the balance between reconstruction and the KL term.

Sum latent coordinates and average examples as requested. Keep the same noise sample for a forward/backward pair; resampling would change the computation being differentiated. Very large log variances can overflow when exponentiated.

Detecting runtime…