Variational Autoencoders

Jr Codex Deep Learning Notes

Level: Advanced Prerequisites: Chapter 1: Autoencoders Time to complete: ~25 minutes

The Key Change: a Probabilistic Latent Space
Encoding to a Distribution, Not a Point
The Reparameterization Trick
The VAE Loss Function
Generating New Data With a Trained VAE
A VAE in PyTorch
VAEs vs Plain Autoencoders
Summary & Next Steps

1. The Key Change: a Probabilistic Latent Space

Chapter 1, Section 7 identified the problem: a plain autoencoder's latent space has no guaranteed structure, making random sampling unreliable for generation. A Variational Autoencoder (VAE) fixes this with one key change: instead of encoding an input to a single fixed point in the latent space, it encodes to a probability distribution — forcing the latent space to be smooth and continuous.

2. Encoding to a Distribution, Not a Point

Plain Autoencoder vs VAE Encoder
─────────────────────────────────────────
  Plain autoencoder (Ch.1):     encoder outputs ONE fixed
                                    vector — a single POINT in
                                    latent space
  VAE:                              encoder outputs TWO vectors:
                                       a MEAN (μ) and a STANDARD
                                       DEVIATION (σ) — defining a
                                       DISTRIBUTION (typically
                                       Gaussian) for each input,
                                       rather than one fixed point
─────────────────────────────────────────

Why This Forces a WELL-STRUCTURED Latent Space
─────────────────────────────────────────
  During training, the VAE's loss (Section 4) explicitly
  penalizes these distributions for straying too far from a
  STANDARD normal distribution (mean 0, standard deviation 1)
  — this REGULARIZES the latent space, actively encouraging
  nearby points to represent SIMILAR, semantically related
  data, and discouraging "gaps" that correspond to nothing
  realistic. This is EXACTLY what a plain autoencoder lacked.
─────────────────────────────────────────

3. The Reparameterization Trick

A subtle but critical implementation detail: randomly sampling from the distribution described by (μ, σ) is not, by itself, a differentiable operation — and Module 2, Chapter 4's backpropagation requires every operation in the forward pass to be differentiable to compute gradients through it.

import torch
 
def reparameterize(mu, log_var):
    """
    The REPARAMETERIZATION TRICK: instead of sampling z directly
    from N(mu, sigma) — which isn't differentiable — sample
    NOISE from a FIXED, simple distribution, then use it to
    construct z through operations THAT ARE differentiable.
    """
    std = torch.exp(0.5 * log_var)
    epsilon = torch.randn_like(std)      # random noise, sampled from N(0, 1) — has NO learnable parameters
    z = mu + epsilon * std                    # z is now a DIFFERENTIABLE function of mu and std
    return z

Why This Works
─────────────────────────────────────────
  The RANDOMNESS is isolated entirely in `epsilon`, which has
  NOTHING to learn (it's just standard Gaussian noise) — the
  actual LEARNABLE parameters (mu, std, and everything that
  produced them) are combined with epsilon through simple
  ADDITION and MULTIPLICATION, both of which are perfectly
  ordinary, differentiable operations that Module 2, Chapter
  4's backpropagation handles without any special treatment.
─────────────────────────────────────────

4. The VAE Loss Function

A VAE's loss combines two terms, balancing reconstruction quality against latent space structure.

The Two-Term Loss
─────────────────────────────────────────
  Reconstruction loss:      how well does the decoder
                               reconstruct the ORIGINAL input
                               from the sampled z? (Chapter 1,
                               Section 4's MSE, or often binary
                               cross-entropy for image pixel
                               values — Module 2, Ch.3)
  KL divergence:                 how far is the encoded
                                     distribution (mu, sigma)
                                     from a STANDARD normal
                                     distribution N(0, 1)? —
                                     THIS is what enforces
                                     Section 2's well-structured
                                     latent space
─────────────────────────────────────────

import torch
import torch.nn.functional as F
 
def vae_loss(reconstructed, original, mu, log_var):
    reconstruction_loss = F.binary_cross_entropy(reconstructed, original, reduction="sum")
 
    # KL divergence between the encoded distribution N(mu, sigma) and N(0, 1) — has a closed-form solution
    kl_divergence = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
 
    return reconstruction_loss + kl_divergence
 
# The TWO terms are in TENSION with each other:
#   - Reconstruction loss ALONE would push the model toward Chapter 1's
#       plain-autoencoder behavior (encode however is MOST convenient
#       for reconstruction, ignoring latent structure)
#   - KL divergence ALONE would push EVERY input toward the SAME
#       distribution N(0,1), destroying any useful information
# Balancing BOTH is what produces a latent space that is BOTH
# informative AND well-structured

5. Generating New Data With a Trained VAE

Once trained, generating new data is straightforward: sample a random point directly from the standard normal distribution the latent space was regularized toward, and run it through the decoder alone.

import torch
 
def generate_new_samples(decoder, num_samples, latent_dim):
    with torch.no_grad():
        z = torch.randn(num_samples, latent_dim)      # sample DIRECTLY from N(0, 1) — no encoder needed
        generated = decoder(z)
    return generated
 
new_images = generate_new_samples(trained_decoder, num_samples=16, latent_dim=32)

This is precisely what a plain autoencoder could not reliably do (Chapter 1, Section 7) — because the VAE's training explicitly shaped the latent space to match N(0,1), randomly sampled points now correspond to realistic, plausible outputs, not arbitrary noise.

6. A VAE in PyTorch

import torch
import torch.nn as nn
 
class VAE(nn.Module):
    def __init__(self, input_dim=784, latent_dim=32):
        super().__init__()
        self.encoder_shared = nn.Sequential(
            nn.Linear(input_dim, 256), nn.ReLU(),
        )
        self.fc_mu = nn.Linear(256, latent_dim)          # produces μ (Section 2)
        self.fc_log_var = nn.Linear(256, latent_dim)     # produces log(σ²), for numerical stability
 
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 256), nn.ReLU(),
            nn.Linear(256, input_dim), nn.Sigmoid(),
        )
 
    def encode(self, x):
        hidden = self.encoder_shared(x)
        return self.fc_mu(hidden), self.fc_log_var(hidden)
 
    def reparameterize(self, mu, log_var):      # Section 3
        std = torch.exp(0.5 * log_var)
        epsilon = torch.randn_like(std)
        return mu + epsilon * std
 
    def forward(self, x):
        mu, log_var = self.encode(x)
        z = self.reparameterize(mu, log_var)
        reconstructed = self.decoder(z)
        return reconstructed, mu, log_var
 
model = VAE()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
for images, _ in train_loader:
    images_flat = images.view(images.size(0), -1)
    reconstructed, mu, log_var = model(images_flat)
    loss = vae_loss(reconstructed, images_flat, mu, log_var)      # Section 4
 
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

7. VAEs vs Plain Autoencoders

Choosing Between Them
─────────────────────────────────────────
  Plain autoencoder (Ch.1):     use for dimensionality
                                    reduction, anomaly detection,
                                    or denoising — tasks that
                                    only need RECONSTRUCTION,
                                    not GENERATION of new data
  VAE:                              use when you actually need
                                        to GENERATE new, plausible
                                        data by sampling — at the
                                        cost of somewhat blurrier
                                        reconstructions than a
                                        plain autoencoder typically
                                        produces (a known,
                                        documented VAE limitation
                                        that GANs, Chapter 3,
                                        address differently)
─────────────────────────────────────────

8. Summary & Next Steps

Key Takeaways

VAEs encode inputs to a probability distribution (mean and standard deviation) rather than a single point, explicitly regularizing the latent space to be smooth and well-structured.
The reparameterization trick isolates randomness into a fixed, parameter-free noise term, keeping the sampling operation differentiable for backpropagation.
A VAE's loss combines reconstruction quality with KL divergence (pulling the encoded distribution toward a standard normal), balancing informativeness against latent space structure.
Once trained, new data can be generated by sampling directly from a standard normal distribution and passing it through the decoder alone — something a plain autoencoder cannot reliably do.

Concept Check

What specific problem does encoding to a distribution (rather than a point) solve, compared to a plain autoencoder?
Why is the reparameterization trick necessary for training a VAE with backpropagation?
What would happen to the latent space if the VAE loss used only reconstruction loss, without the KL divergence term?

Next Chapter

→ Chapter 3: GANs

Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to DL Index