Autoencoders and VAEs

An RNN can compress three ordered event records into one final state. An equipment-monitoring workflow also receives images: a close crop of a cracked panel, a loose connector, or a clean surface. These pixel arrays are much larger than the decision they support.

An autoencoder learns to reconstruct an input through a latent code. The undercomplete autoencoder used here makes that code smaller than the input, creating a visible compression bottleneck. A variational autoencoder (VAE) makes its encoder predict a distribution over latent codes and trains that distribution against a prior, so sampling becomes part of the model rather than an unsupported guess.^{[1]Reference 1Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114}

Compress one photo patch

Start with a tiny crop from an inspection image. Each grayscale number is between 0 (dark) and 1 (bright). A real photo has many more pixels, but a 5 by 5 crop lets us read every number.

An autoencoder has two learned functions:

$$z = f_\phi(x) \qquad \hat{x} = g_\theta(z)$$

The encoder $f_\phi$ maps the input $x$ to the latent code $z$. The decoder $g_\theta$ maps that code to a reconstruction $\hat{x}$. A narrow $z$ creates the bottleneck: the decoder can't copy all 25 pixels directly, so training must preserve information that helps reconstruction.

For normalized pixels, mean-squared error (MSE) is a simple reconstruction measure:

$$L_{\text{recon}} = \frac{1}{N}\sum_{i=1}^{N}(x_i-\hat{x}_i)^2$$

First run a hand-designed encoder and decoder. The weights aren't trained; they make the bottleneck calculation visible before we ask optimization to discover one.

The input shrinks from 25 numbers to 2. This crafted decoder does well on this one pattern because we built its weights around that pattern. A model becomes useful only when its weights are learned across many examples.

Learn the bottleneck

The next program makes noisy variations of the same inspection-image crop and trains an autoencoder in PyTorch. PyTorch is the right tool here because the weights must receive gradients and optimizer updates.

The decoder learned to reconstruct this small family of patches through two latent numbers. It wasn't trained to decode arbitrary random pairs of numbers. That's the generation problem: a deterministic autoencoder defines codes for observed inputs, but its loss doesn't make a simple random distribution over codes useful.

Why isn't a low reconstruction loss enough to make a deterministic autoencoder a generator?

Answer

The reconstruction loss evaluates codes produced by the encoder for training-like inputs. It doesn't require the decoder to behave sensibly on points sampled from a chosen prior such as a standard normal distribution.

A VAE learns distributions over codes

A VAE makes prior sampling useful through its training objective. Instead of producing one code, the encoder predicts the parameters of an approximate posterior, written $q_\phi(z \mid x)$. The word approximate matters: computing the exact posterior $p_\theta(z \mid x)$ is generally impractical, so the encoder learns a tractable stand-in. For the common diagonal-Gaussian case, it predicts a mean vector $\mu$ and a log-variance vector $\log \sigma^2$:

$$q_\phi(z \mid x) = \mathcal{N}\left(\mu(x), \operatorname{diag}(\sigma^2(x))\right)$$

The model also declares a simple prior, usually:

$$p(z) = \mathcal{N}(0, I)$$

The posterior describes codes for one observed patch. The prior is the distribution we want to sample from when no input patch is provided.

Why predict logvar rather than a raw standard deviation? A neural layer may output any real number, while a standard deviation must be positive. Exponentiating half the log variance converts any real output into a valid standard deviation.

Keep posteriors near a sampleable prior

If the encoder may put each patch distribution anywhere it likes, prior samples still won't reliably land where the decoder has trained. VAEs penalize mismatch between the posterior and prior with Kullback-Leibler (KL) divergence. KL isn't a symmetric distance metric: direction matters. For a diagonal Gaussian posterior and standard-normal prior, the VAE uses:

$$D_{\mathrm{KL}}\left(q_\phi(z \mid x) \,\|\, p(z)\right) = -\frac{1}{2} \sum_j \left(1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2\right)$$

The penalty is zero when $\mu=0$ and $\sigma=1$. Moving means away from zero or making variances very different from one increases it.

KL pressure isn't free. If it dominates reconstruction, all inputs can be encoded nearly like the prior and the decoder may ignore $z$. The next diagnostic catches that failure.

Sample while gradients still flow

Kingma and Welling's VAE estimator rewrites random sampling so gradients can reach encoder parameters.^{[1]Reference 1Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114} Draw independent noise $\epsilon$ from a standard normal and then transform it:

$$\epsilon \sim \mathcal{N}(0,I), \qquad z = \mu + \sigma \odot \epsilon$$

During one backward pass, $\epsilon$ is a sampled constant. The multiplication and addition still expose how the loss changes when the encoder changes $\mu$ or $\sigma$.

The VAE objective

The VAE optimizes an evidence lower bound (ELBO) on the data log-likelihood.^{[1]Reference 1Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114} Written as a quantity to maximize:

$$\mathcal{L}_{\mathrm{ELBO}} = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - D_{\mathrm{KL}}(q_\phi(z \mid x)\,\|\,p(z))$$

The first term rewards a decoder that makes the observed patch likely after sampling its latent. The second discourages a posterior that moves too far from the prior.

In code we usually minimize the negative ELBO:

$$L_{\mathrm{VAE}} = L_{\mathrm{recon}} + \beta D_{\mathrm{KL}}$$

For $\beta=1$, this corresponds to the standard objective when L_recon is the selected negative log-likelihood. An MSE reconstruction term is proportional to a Gaussian decoder negative log-likelihood only after choosing a fixed observation variance, so treat the scale of MSE and KL as a modeling choice, not an identity.

The reconstruction term is an expectation over $q_\phi(z \mid x)$. Training code almost always draws one reparameterized sample of $z$ per example and uses that single-sample Monte Carlo estimate of the expectation. More samples would lower variance; one sample is the usual default because it's cheap and unbiased for the ELBO's reconstruction piece.

This ledger computes both terms for one patch using a fixed hypothetical reconstruction:

Higgins et al. studied the $\beta$-VAE objective, where increasing $\beta$ applies stronger latent-capacity and independence pressure while trading against reconstruction accuracy.^{[2]Reference 2beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Frameworkhttps://openreview.net/forum?id=Sy2fzU9gl} That trade-off must be measured on your task; a higher value isn't automatically better.

These rows are an accounting exercise, not a reported experiment. They show why you should log reconstruction and KL separately instead of reading only the total.

Train a tiny VAE

Now train a real two-dimensional VAE on two small families of noisy patches: one with bright center cells and one with bright lower cells. The sample batch at the end demonstrates the generative interface. It doesn't claim photographic quality from this miniature dataset.

Each object in the program now has a job: mu and logvar define each approximate posterior, torch.randn_like supplies $\epsilon$, decoder consumes sampled $z$, and the logged losses reveal the compromise.

Watch for posterior collapse

If a decoder can reconstruct without using its latent input, KL pressure can push every posterior close to the prior. Then the latent code stops carrying information. This failure is called posterior collapse. Van den Oord et al. describe it as a common issue for VAEs paired with high-capacity autoregressive decoders.^{[3]Reference 3Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937}

A first diagnostic is to track KL per latent dimension. The illustrative monitor below flags a run whose KL has nearly vanished in every dimension:

Low KL alone isn't a proof of failure. Pair it with reconstruction quality, generated samples, and checks that predictions change when $z$ changes. Still, a rapid collapse toward zero is worth investigating.

If the latent has collapsed, common levers are:

• KL annealing: start with a small $\beta$ (or zero) so reconstruction can learn to use $z$, then ramp $\beta$ up.
• Free bits: floor per-dimension KL so the optimizer isn't rewarded for driving every unit all the way to the prior.
• Weaker decoder early: a powerful decoder can reconstruct from $x$ side channels and ignore $z$; capacity limits or simpler early decoders force dependence on the code.
• Lower $\beta$ overall: if reconstruction is fine but KL is near zero, reduce the KL weight and recheck $z$-sensitivity.

None of these is a free lunch; they trade how hard you push the prior against how much information the code must carry.

VQ-VAE: replace continuous samples with codes

A Vector Quantised Variational Autoencoder (VQ-VAE) takes a different route. Its encoder output is mapped to the nearest vector in a learned codebook, so its latent representation is discrete rather than Gaussian. The original VQ-VAE paper also learns a prior over those discrete codes.^{[3]Reference 3Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937}

For one encoder vector, quantization is a nearest-neighbor lookup:

$$k = \arg\min_j \lVert z_e(x)-e_j\rVert_2^2, \qquad z_q(x)=e_k$$

Nearest-code selection isn't differentiable. VQ-VAE uses a straight-through gradient estimator for the encoder path plus codebook and commitment losses so encoder outputs stay close to useful entries.^{[3]Reference 3Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937} The important distinction is interface:

Latent models reduce the expensive space

Rombach et al. train diffusion models in the latent space of pretrained autoencoders rather than applying all denoising operations directly to pixels.^{[4]Reference 4High-Resolution Image Synthesis with Latent Diffusion Models.https://arxiv.org/abs/2112.10752} This makes the autoencoder's role precise: compress image data before expensive generation steps, then decode a final latent back to pixels.

The size reduction itself is simple arithmetic. A downsampling factor of 8 reduces each spatial axis by 8, so it reduces spatial positions by 64:

Channels and decoder quality still matter, and individual systems choose their own latent representation. The durable idea is that a learned bottleneck can make later modeling cheaper while retaining useful structure.

Debugging checklist

• Plot or print reconstruction loss and KL loss separately. Total loss hides which side dominates.
• Check that sigma = exp(0.5 * logvar), not exp(logvar), before reparameterized sampling. If samples or KL become NaN or Inf, inspect logvar for unstable extremes before adding an explicit clamp.
• Keep posterior and prior names straight: the encoder supplies $q_\phi(z \mid x)$; generation samples from $p(z)$.
• Inspect whether reconstructions or predictions change when $z$ changes. If not, the decoder may ignore its latent input.
• For VQ-VAE, log code usage. A codebook where only a few IDs are selected wastes capacity.

---

Mastery check

Key concepts

• An undercomplete autoencoder reconstructs inputs through a narrow latent bottleneck.
• Reconstruction loss trains codes for observed inputs but doesn't define a sampleable prior.
• A VAE encoder predicts a Gaussian approximate posterior using mu and logvar.
• KL divergence keeps each posterior near the chosen prior.
• Reparameterization isolates random noise while preserving gradients to encoder outputs.
• The negative ELBO balances reconstruction likelihood and KL regularization.
• Posterior collapse means the decoder stops using its latent code.
• VQ-VAE exposes discrete codebook selections rather than continuous Gaussian samples.
• Latent diffusion runs expensive generation steps in a compressed autoencoder space.

Evaluation rubric

• Foundational: Explains why a bottleneck forces compression and computes an MSE reconstruction loss.
• Foundational: Distinguishes reconstructing encoded data from generating by sampling a prior.
• Intermediate: Converts logvar to sigma, calculates Gaussian KL, and produces a reparameterized sample.
• Intermediate: Identifies reconstruction and KL terms in a PyTorch VAE training loop.
• Advanced: Diagnoses posterior collapse with multiple signals rather than KL alone.
• Advanced: Compares continuous VAE latents with VQ-VAE code IDs and latent diffusion's compressed space.

Common pitfalls

• Symptom: Random deterministic-autoencoder codes are treated as valid samples. Cause: Reconstruction training was mistaken for prior training. Fix: State how new codes are generated and whether that source appeared in the objective.
• Symptom: Sampling or KL loss returns NaN or Inf. Cause: Scale parameterization is wrong or logvar became extreme before exponentiation. Fix: Compute sigma = exp(0.5 * logvar), inspect the logvar range, and stabilize optimization before adding an explicit clamp.
• Symptom: Reconstruction looks adequate while latent dimensions carry no information. Cause: Posterior collapse. Fix: Inspect KL per dimension and test sensitivity to changing z. If collapsed, try KL annealing, free bits, a weaker early decoder, or a lower $\beta$.
• Symptom: MSE and ELBO are described as identical in every VAE. Cause: Decoder likelihood choice was skipped. Fix: Tie reconstruction loss to the observation model and its scaling.
• Symptom: VQ lookup is presented as ordinary differentiable code. Cause: The discrete selection step was omitted. Fix: Explain the straight-through estimator and commitment/codebook losses.

Follow-up questions

What changes when an autoencoder becomes a VAE?

Answer

The encoder no longer emits only one deterministic latent vector. It predicts parameters of a posterior distribution, the model samples a latent through reparameterization, and training adds a KL term that relates that posterior to a prior used for generation.

Why does the VAE use z = mu + sigma * epsilon?

Answer

The random draw is moved into epsilon, which is independent of the encoder parameters. For a fixed sampled epsilon, z remains a differentiable function of mu and sigma, so reconstruction gradients reach the encoder.

Why should you log KL separately from reconstruction loss?

Answer

A small total can hide failure. Very weak KL can leave poor prior sampling, while KL collapsing near zero can indicate that the decoder ignores z. The two terms explain different behaviors.

How does this lesson prepare you for Transformer architecture?

Answer

You now recognize an encoder as a machine that converts raw input into a representation another component can use. Transformers build contextual representations for tokens and connect them through attention rather than through an image reconstruction bottleneck.