A First-Principles Tutorial · Chapter 03

Autoencoders

Compress, reconstruct, generate — how bottleneck networks learn the hidden structure of data.

By Prashant Kulkarni · with help from Claude OpusCC BY-NC 4.0

1 The core idea

Copy your input — through a bottleneck

An autoencoder is a neural network trained to copy its input to its output — but forced to do so through a narrow bottleneck. The bottleneck prevents the identity function and instead forces the network to discover compact, meaningful structure in the data.

The autoencoder pipeline. A 784-dimensional MNIST image is compressed through the encoder to a $d$-dimensional latent code $z$ (the bottleneck), then reconstructed by the decoder. The loss is simply the squared distance between input and reconstruction. When $d \ll D$, the network must learn to keep only the information that matters. Original SVG

The bottleneck dimension $d$ controls how aggressively the network compresses. With $d = 2$ on MNIST, the network must squeeze a 784-pixel image into just two numbers — forcing it to discover that digits lie on a low-dimensional manifold. These two numbers encode the "type" and "style" of the digit in a continuous, learned coordinate system.

2 Vanilla autoencoders

Compression without generation

A vanilla AE minimizes reconstruction loss: $\mathcal{L} = \frac{1}{N}\sum_i \|x_i - g_\phi(f_\theta(x_i))\|^2$. After training, the encoder produces a compact representation, and the decoder can reconstruct from it. This gives you dimensionality reduction (like PCA but nonlinear) and learned features.

But there's a problem: you can't generate new data. The latent space of a vanilla AE has no structure — there are gaps, discontinuities, and dead regions. If you sample a random point $z$ from the latent space, the decoder will likely produce garbage. The encoder only learned to map training inputs to latent points; the space between those points is undefined.

"A vanilla autoencoder learns to compress — but only the data it has seen. The space between encoded points is a no-man's-land."

3 The VAE

From compression to generation

The Variational Autoencoder (Kingma & Welling, 2013) makes one decisive change: instead of encoding each input to a single point, the encoder outputs the parameters of a probability distribution — a mean $\mu$ and a variance $\sigma^2$ — from which the latent code is sampled.

The VAE architecture. The encoder outputs parameters ($\mu$, $\log\sigma^2$) of a Gaussian, not a single point. The reparameterization trick — $z = \mu + \sigma \odot \epsilon$ with $\epsilon \sim \mathcal{N}(0,I)$ — makes the sampling step differentiable, allowing gradients to flow back to the encoder. Original SVG · concept after Kingma & Welling (2013)

4 The objective

The ELBO — two forces in tension

The VAE maximizes the Evidence Lower Bound:

\text{ELBO} = \underbrace{\mathbb{E}_{q_\theta(z|x)}[\log p_\phi(x|z)]}_{\text{reconstruction: "make good outputs"}} \;-\; \underbrace{D_\text{KL}(q_\theta(z|x) \| p(z))}_{\text{regularization: "stay close to prior"}}

Two forces pulling in opposite directions:

Reconstruction term: pushes the model to encode enough information in $z$ to reconstruct $x$ accurately. Wants large $\sigma$ and high-information $\mu$ — the same pressure as a vanilla AE.
KL term: pulls every encoder posterior toward the standard normal prior $\mathcal{N}(0, I)$. Wants small $\mu$ (close to origin) and $\sigma \approx 1$ (not too peaked). This is the regularizer that gives the latent space its structure.

When both posteriors and priors are Gaussian, the KL has a closed form:

D_\text{KL} = -\frac{1}{2}\sum_{j=1}^{d}\left(1 + \log\sigma_j^2 - \mu_j^2 - \sigma_j^2\right)

Why KL regularization enables generation

Without KL, each input maps to an isolated point in latent space — the decoder memorizes those points, and the space between them is meaningless. The KL term forces encoded distributions to overlap near the origin, filling the gaps. Now you can sample $z \sim \mathcal{N}(0, I)$, decode it, and get a plausible output — because the decoder has seen training data everywhere near the origin.

5 The trick

Reparameterization — making sampling differentiable

The ELBO requires sampling $z \sim q_\theta(z|x)$, but sampling is not differentiable — gradients can't flow through a random node. The fix is elegant: instead of sampling from a $\theta$-dependent distribution, sample noise from a fixed distribution and transform it deterministically:

\epsilon \sim \mathcal{N}(0, I), \qquad z = \mu_\theta + \sigma_\theta \odot \epsilon

Now $z$ is a differentiable function of $\mu$ and $\sigma$ (given fixed $\epsilon$). Gradients flow through to the encoder. The randomness comes from $\epsilon$, which doesn't depend on any learned parameters.

In code, the encoder outputs $\mu$ and $\log\sigma^2$ (log-variance for numerical stability). The sampling step is one line: z = mu + torch.exp(0.5 * logvar) * eps.

6 Latent space

What makes a VAE latent space special

The KL regularizer gives VAE latent spaces three properties that vanilla AEs lack:

Smoothness. Nearby points in latent space decode to similar outputs. There are no sharp discontinuities — the space is a continuous manifold.
Completeness. Every point in the "inhabited" region of latent space decodes to something meaningful. No dead zones.
Interpolation. Walking in a straight line between two encoded points produces a smooth morph between their corresponding outputs — a "3" gradually becomes an "8" through plausible intermediate forms.

Vanilla AE latent space

Clusters are tight and separated. Gaps between clusters decode to garbage. No principled way to sample new points. Good for compression, bad for generation.

VAE latent space

Clusters overlap and fill the space. Any point near the origin decodes to something plausible. Interpolation is smooth. Generation works by sampling $z \sim \mathcal{N}(0,I)$ and decoding.

7 VQ-VAE

Discrete codes and codebooks

The Vector Quantized VAE (van den Oord et al., 2017) takes a different path: instead of a continuous Gaussian latent space, it uses a discrete codebook of $K$ learned embedding vectors $\{e_1, \ldots, e_K\}$.

The encoder produces a continuous vector $z_e$, which is then replaced by its nearest codebook entry:

z_q = e_k \quad\text{where}\quad k = \arg\min_j \|z_e - e_j\|_2

The result: the latent representation is a grid of discrete indices into the codebook — exactly the format that autoregressive models (GPT, PixelCNN) can model. This is the bridge between vision and language: images become sequences of codebook tokens.

VQ-VAE replaces the continuous Gaussian with a discrete nearest-neighbor lookup into a learned codebook. The encoder produces continuous $z_e$; quantization snaps it to the nearest codebook vector $e_k$; the decoder reconstructs from the quantized code $z_q$. The straight-through estimator copies gradients from $z_q$ directly back to $z_e$ during backpropagation. Original SVG · concept after van den Oord et al. (2017)

The VQ-VAE loss has three terms: reconstruction loss (same as vanilla AE), a codebook loss that moves codebook entries toward encoder outputs, and a commitment loss that moves encoder outputs toward codebook entries. The $\arg\min$ is non-differentiable; the straight-through estimator copies gradients directly from decoder input to encoder output, bypassing quantization.

Why VQ-VAE matters for modern generative AI

Stable Diffusion's latent space uses a VAE (continuous). DALL-E 1 used a dVAE (discrete, VQ-VAE variant). Modern audio codecs (EnCodec, SoundStream) use residual VQ-VAE to tokenize audio for language-model-style generation. The VQGAN combines VQ-VAE with adversarial training for higher-quality codebook learning. Everywhere you see "tokenize continuous data for a transformer," a VQ-VAE variant is doing the work.

8 Applications

Where autoencoders appear in generative AI

Stable Diffusion's VAE. A convolutional VAE compresses 512×512×3 images to 64×64×4 latents (48× reduction). The diffusion process runs entirely in this latent space. The VAE is trained first, then frozen; the diffusion model only sees latents.
DALL-E 1's dVAE. A discrete VAE (VQ-VAE variant with Gumbel-softmax relaxation) tokenizes images into a grid of 32×32 discrete codes from a vocabulary of 8192. A transformer then models the joint distribution of text tokens and image tokens autoregressively.
Audio tokenization. EnCodec and SoundStream use residual VQ-VAE to compress audio waveforms into sequences of discrete tokens at various bitrates. These tokens are then modeled by language-model-style transformers (MusicGen, AudioLM) for generation.
Representation learning. Even when not generating, autoencoders provide pretrained features. The encoder's latent space captures semantically meaningful structure that downstream tasks can exploit.

9 Field notes

AE vs VAE vs VQ-VAE on MNIST

Source

Experiments from the Module 03 notebook in the learn-generative-ai repo. All three models trained on MNIST with a 2D latent space (for AE and VAE) or 64-entry codebook (for VQ-VAE).

Reconstruction quality

All three models reconstruct recognizable digits, but with different tradeoffs. The vanilla AE produces the sharpest reconstructions — it puts all capacity into exact copying. The VAE produces slightly blurrier outputs because the KL term forces it to "waste" some capacity on regularizing the latent space. The VQ-VAE is sharp again — no KL smoothing — but occasionally shows quantization artifacts when the encoder output falls between codebook entries.

Latent space structure

The 2D latent space scatter plots (colored by digit class) reveal the key difference between AE and VAE. The vanilla AE produces tight, well-separated clusters with large gaps between them — great for classification, terrible for generation. The VAE produces overlapping, smoothly-distributed clusters centered near the origin — exactly what the KL loss enforces. Walking between the "3" region and the "8" region produces plausible intermediate forms.

Grid sampling from the VAE

Sampling a uniform grid of $z$ values from $[-3, 3]^2$ and decoding each point produces a beautiful visualization: digits smoothly morph across the grid, transitioning from one class to another through recognizable intermediate forms. This is the payoff of the KL regularization — the entire latent space is "covered" and every point produces something meaningful.

Interpolation

Encoding two different digits (e.g., a "3" and an "8"), linearly interpolating between their latent codes, and decoding the intermediate points shows a smooth morphing sequence. The VAE interpolation is smoother and more plausible than the vanilla AE interpolation, which can pass through distorted, unrealistic forms in the gaps between clusters.

Codebook utilization (VQ-VAE)

The notebook also checks how many of the 64 codebook entries are actually used. A well-trained VQ-VAE uses most of its codebook; a collapsed one concentrates on a handful of entries (dead codes). The trained model uses 50+ of 64 entries, with the distribution roughly following a Zipf-like pattern — some codes are more popular than others, but few are completely dead.

10 Playground

Explore a 2D latent space

Below is a simulated 2D latent space with 10 "digit classes" arranged in a ring. Click anywhere to see what the decoder would produce at that point — represented as a colored shape that smoothly morphs between classes. Drag to draw an interpolation path and watch the smooth transition. Toggle between "AE mode" (sharp clusters, dead gaps) and "VAE mode" (smooth, continuous coverage).

      
      click to sample · drag to interpolate

VAE mode · smooth transitions, full coverage

⁂ ⁂ ⁂

11 Closing

Takeaways & exercises

The six things to remember

An autoencoder learns to compress by reconstructing. The bottleneck forces it to keep only the information that matters.
Vanilla AEs can't generate because the latent space has gaps and dead zones. The decoder only knows the points it was trained on.
The VAE adds a KL regularizer that pulls every encoded distribution toward the prior $\mathcal{N}(0,I)$, filling the gaps and enabling generation by sampling.
The reparameterization trick ($z = \mu + \sigma \odot \epsilon$) makes sampling differentiable by moving randomness to a fixed external source.
VQ-VAE uses a discrete codebook instead of a continuous Gaussian, bridging vision and language by turning images into sequences of tokens that transformers can model.
Autoencoders are everywhere in modern generative AI: Stable Diffusion's latent space (VAE), DALL-E's image tokenizer (dVAE), audio codecs (VQ-VAE), and representation learning in general.

Exercises

1. Why does the VAE produce blurrier outputs than the vanilla AE? Which term in the ELBO is responsible, and what would happen if you set its weight to zero?

2. Derive the closed-form KL divergence between $\mathcal{N}(\mu, \sigma^2)$ and $\mathcal{N}(0, 1)$ for the scalar case ($d=1$). Verify that it is zero when $\mu=0, \sigma=1$ and positive otherwise.

3. In VQ-VAE, the $\arg\min$ is non-differentiable. Explain the straight-through estimator in your own words. What approximation does it make, and why is it acceptable?

4. Open the notebook. Train the VAE with $\beta = 0.1$ (weak KL) and $\beta = 10$ (strong KL). How does the latent space scatter plot change? How does reconstruction quality change?

5. In the playground above, toggle between AE and VAE mode. Click in a gap between clusters in AE mode — what happens? Now do the same in VAE mode. Why is the difference important for generation?