Practical build 7

Lab: build a variational autoencoder

Prerequisites

Complete Lessons 1 to 4 and the earlier practical builds. You need tensor arithmetic, gradients, Adam, linear layers, and a basic idea of probability. We use a model class following the Lesson 3 bridge. __init__ creates reusable layers and stores them on self. forward describes one sampled training pass. Extra methods named encode and decode make the two directions explicit.

Plain-language start: reconstruct, but keep the middle organized

An ordinary autoencoder compresses an input into one fixed latent vector and decodes it. If we give it point (2.1, -1.9), its encoder may return one latent point such as (4.7, 0.2). Nothing forces nearby empty regions of latent space to decode sensibly. Random sampling can therefore produce poor outputs.

A variational autoencoder, or VAE, makes the encoder describe a probability distribution. For each input, it returns a mean vector and a log-variance vector. We sample a latent vector from that distribution, then decode it. A KL penalty encourages the encoded distributions to stay near a simple standard-normal prior. This makes prior sampling possible, although a small KL weight and small dataset do not guarantee perfect samples.

The generated task

The dataset contains points near four centres: (-2,-2), (-2,2), (2,-2), and (2,2). Gaussian noise with standard deviation 0.30 moves each point slightly around its centre. Training has 1,024 points made with seed 50. Held-out evaluation has 256 points made with seed 60. There is no network access and no hidden file.

points, labels = make_dataset(count=1024, seed=50)
assert points.shape == (1024, 2)
assert points.dtype == torch.float32
assert labels.shape == (1024,)

An input is one row x = (x_1, x_2). Its target is itself because reconstruction tries to recover the input. The encoder returns mean mu and log-variance logvar, each shape (B, 2). A sampled latent tensor z is (B, 2). The decoder returns reconstruction x_hat with shape (B, 2).

The simple baseline always predicts the mean training point. Because the four clusters are balanced around zero, this is close to (0,0). We calculate its mean squared error on held-out points. A useful VAE reconstruction should be much lower.

From mean and variance to a sampled latent vector

Calculate one sample: if mean = 2, standard deviation is 0.5, and the sampled noise is -1, then z = 2 + 0.5 × (-1) = 1.5. The noise is random. The mean and scale come from the encoder.

The encoder represents each latent feature as a Gaussian distribution. For feature j, it predicts mean mu_j and variance sigma_j^2. Variance must be positive. Instead of predicting it directly, the network predicts log(sigma_j^2), which may be any real number. We recover standard deviation with:

sigma = exp(0.5 * log_variance)

The factor 0.5 appears because log(sigma^2) = 2 log(sigma). Direct random sampling from a distribution whose parameters came from the network hides the gradient path. The reparameterization trick moves randomness into a separate standard-normal variable:

epsilon ~ Normal(0, I)
z = mu + sigma * epsilon

Now epsilon is random but independent of the encoder parameters. For a fixed sample of epsilon, z is an ordinary differentiable calculation from mean and standard deviation. Gradients can reach both encoder heads.

def reparameterize(mean, log_variance, *, generator=None, epsilon=None):
    standard_deviation = torch.exp(0.5 * log_variance)
    if epsilon is None:
        epsilon = torch.randn(
            standard_deviation.shape,
            generator=generator,
        )
    return mean + standard_deviation * epsilon

The optional epsilon is useful for a precise test. If mean is (1,-1), log-variance is (0,0), and epsilon is (0.5,2), standard deviation is (1,1) and the result is exactly (1.5,1).

The VAE class and all shapes

class TinyVAE(nn.Module):
    def __init__(self, hidden_size=32, latent_size=2):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(2, hidden_size),
            nn.Tanh(),
        )
        self.mean_head = nn.Linear(hidden_size, latent_size)
        self.log_variance_head = nn.Linear(hidden_size, latent_size)
        self.decoder = nn.Sequential(
            nn.Linear(latent_size, hidden_size),
            nn.Tanh(),
            nn.Linear(hidden_size, 2),
        )

    def encode(self, x):
        hidden = self.encoder(x)                 # (B, 32)
        return self.mean_head(hidden), self.log_variance_head(hidden)

    def decode(self, z):
        return self.decoder(z)                   # (B, 2)

    def forward(self, x, *, generator=None):
        mean, log_variance = self.encode(x)      # each (B, 2)
        z = reparameterize(mean, log_variance, generator=generator)
        return self.decode(z), mean, log_variance

nn.Sequential applies its stored modules from left to right. It is suitable when there is one straight path. The encoder then branches, so the two heads are stored separately. One head cannot replace the other: the mean locates the distribution while log-variance describes its spread.

Tanh provides a nonlinear hidden transformation. Without a nonlinear activation, stacked linear layers collapse into one linear map and cannot bend the representation around the cluster structure. The final decoder layer has no activation because coordinates may be any real value.

Loss part 1: reconstruction

For one two-dimensional example, squared reconstruction error is:

L_reconstruction = (x_hat_1 - x_1)^2 + (x_hat_2 - x_2)^2

The implementation sums the two coordinate errors for each example, then averages over the batch. This tells the decoder and encoder to preserve information needed to rebuild each point. Reconstruction alone would create an ordinary stochastic autoencoder with no reason to match the chosen prior.

Loss part 2: KL divergence

Predict the trade-off before changing its weight. A very small KL weight can give sharp reconstructions but a disorganized latent space. A very large weight can force tidy latent values while making reconstructions poor. Report both loss parts separately.

The prior is a two-dimensional standard normal: mean zero, variance one, and independent features. For a diagonal Gaussian predicted by the encoder, its KL divergence from that prior has a closed form:

L_KL = -0.5 * sum_j [1 + logvar_j - mu_j^2 - exp(logvar_j)]

The sum runs across latent features. This value is non-negative apart from tiny floating-point error. It is zero when every mean is zero and every variance is one. A large mean or a variance far from one increases the penalty.

The total objective uses a weight called beta:

L_total = L_reconstruction + beta * L_KL

This build uses beta = 0.05. A larger beta pushes distributions closer to the prior but may discard information and worsen reconstruction. A smaller beta improves reconstruction but can create a poorly organized latent space. There is no universally correct value.

def vae_loss(reconstruction, target, mean, log_variance, beta=0.05):
    reconstruction_loss = ((reconstruction - target) ** 2).sum(dim=1).mean()
    kl_loss = -0.5 * (
        1 + log_variance - mean.square() - log_variance.exp()
    ).sum(dim=1).mean()
    total = reconstruction_loss + beta * kl_loss
    return total, reconstruction_loss, kl_loss

Training with controlled randomness

model = TinyVAE()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

for step in range(500):
    indices = torch.randint(0, len(points), (128,), generator=batch_generator)
    batch = points[indices]

    optimizer.zero_grad()
    reconstruction, mean, log_variance = model(
        batch, generator=noise_generator
    )
    loss, reconstruction_loss, kl_loss = vae_loss(
        reconstruction, batch, mean, log_variance, beta=0.05
    )
    loss.backward()
    optimizer.step()

Separate seeded generators control batch sampling and latent noise. torch.manual_seed controls parameter initialization. This makes the educational run repeatable. It does not mean all GPU kernels or every future PyTorch version are bit-identical, but this script runs on CPU with one thread.

Evaluation and generation answer different questions

For held-out reconstruction, the program uses the encoder mean instead of a random sample. That removes sampling noise from the measurement:

with torch.no_grad():
    mean, log_variance = model.encode(validation_points)
    reconstruction = model.decode(mean)
    validation_mse = F.mse_loss(reconstruction, validation_points)

This mean squared error averages across both examples and coordinates, while the training reconstruction term sums coordinates before averaging examples. Their numeric scales differ by a factor of two. State the reduction whenever comparing losses.

Generation starts with z sampled from Normal(0,I) and calls only the decoder. It does not encode a source point. Generated samples show that the pipeline works, but eight printed coordinates are not a quality metric. A serious generative evaluation would inspect distribution coverage, cluster proportions, and distances with many samples.

Run the complete program

Download build_vae.py. It includes the data generator, reparameterization, model, loss, optimizer, held-out evaluation, mean baseline, prior sampling, assertions, and CLI controls.

python examples/pytorch/build_vae.py
python examples/pytorch/build_vae.py --steps 200
python examples/pytorch/build_vae.py --smoke-test

The smoke test runs no more than three optimizer updates and skips quality thresholds. It checks shapes and execution, not convergence. On the tested CPU runtime, the full 500-step command produced:

train and validation shapes: (1024, 2) (256, 2)
mean-predictor baseline MSE: 4.0845
held-out deterministic reconstruction MSE: 0.0144
held-out KL per example: 4.2619
sample shape: (8, 2)
first three generated points:
  (+1.152, -1.993)
  (-0.151, -1.539)
  (-0.109, -2.408)

The reconstruction clearly beats the baseline. The second and third generated points sit between chosen cluster centres. That is an honest reminder that prior samples are not guaranteed to land exactly in a data cluster, especially with a low beta and a tiny model.

Failure drills

Drill 1: sample without reparameterization. Convert mean and variance to detached values before sampling. The decoder may still get gradients, but reconstruction gradients cannot train the encoder correctly. Fix the path by expressing z as mean plus standard deviation times independent noise.

Drill 2: use variance as if it were log-variance. Replacing exp(0.5 * log_variance) with log_variance allows negative “standard deviations” and breaks the KL formula. Keep one clear representation and name it accurately.

Drill 3: remove KL completely. Set beta to zero. Reconstruction may improve, but decoded samples from a standard-normal prior can become less meaningful because the encoder had no reason to place codes near that prior.

Drill 4: make beta very large. Try beta 10. The encoder may make every distribution close to the prior and ignore the input. Reconstructions then approach an average. This is posterior collapse in a simple form.

Drill 5: evaluate with random samples only once. The reported reconstruction changes with noise and can hide regressions. Use the latent mean for a stable reconstruction metric, then evaluate sampling separately.

Solved exercises

Limits and what to try next

This complete build demonstrates the defining VAE pieces: distribution parameters, reparameterized sampling, reconstruction plus KL, held-out evaluation, and prior generation. It does not establish image quality, calibrated likelihood, cluster coverage, or useful disentanglement. The latent dimension equals the input dimension, and the generated distribution is intentionally simple.

Useful extensions are to sweep beta, plot encoded means, measure how often prior samples fall near each centre, or reduce the latent size to one and observe the reconstruction trade-off. Change one factor at a time, preserve the baseline, and report held-out measurements rather than judging only a few attractive samples.