Autoencoders and VAEs

Sequence models compress an ordered history into a compact hidden state. An inspection vision pipeline makes the same move on image data: taking a high-dimensional screenshot crop and squeezing it down to a handful of numbers.

Can two scalar numbers preserve a bright 2 by 2 alert badge on a 5 by 5 grayscale screenshot? And if you hand the decoder two numbers drawn at random instead of an encoded crop, should it know what to draw? Compression and generation impose different constraints. An autoencoder learns to reconstruct its input from a latent code, which we make compact here.^{[1]Reference 1Reducing the Dimensionality of Data with Neural Networkshttps://doi.org/10.1126/science.1127647} A variational autoencoder (VAE) learns a probabilistic latent model with a prior we can sample for generation. Its objective connects reconstruction and prior sampling, but useful samples still depend on the model, data, and training.^{[2]Reference 2Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114}

Compress one screenshot crop

Start with a tiny crop. Each grayscale number is between 0 (dark) and 1 (bright). A real screenshot has many more pixels, but a 5 by 5 crop lets us inspect every single number. The bright 2 by 2 block is the highlighted region. Before looking at encoder maps, predict: which pixels should one latent coordinate watch, and which should it ignore?

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$. Meanwhile, the decoder $g_\theta$ maps that code to a reconstruction $\hat{x}$. Here $\phi$ and $\theta$ denote their learned weights. A narrow $z$ creates the hourglass bottleneck: there's no direct wire carrying all 25 pixels to the decoder. This forces the network to discover compact summaries, but doesn't prevent a large network from memorizing a tiny dataset.

For this crop, the first encoder map averages the four bright center pixels and the second averages two lower pixels:

$$z_1 = 0.25 \times (0.85 + 0.88 + 0.90 + 0.92) = 0.8875 \approx 0.888$$ $$z_2 = 0.50 \times (0.18 + 0.16) = 0.170$$

Those aren't learned weights. They're a hand-designed encoder, so the bottleneck arithmetic is visible before we ask optimization to discover one.

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$$

$N$ is the number of pixels, here 25. Each term is one pixel's squared error. The next snippet runs that same encoder, then a matching crafted decoder, so you can see the 25-to-2 shrink and the reconstruction error together.

The input shrinks from 25 numbers to 2. This crafted decoder does well on this one pattern because we built its weights around it. Notice what it loses: rearranging the four selected center pixels while preserving their average leaves $z_1$ unchanged. The decoder then produces the same reconstruction for different inputs. To handle a family of crops, learn the weights instead of selecting them by hand.

Learn the bottleneck

The next program makes noisy variations of the same crop and trains an autoencoder in PyTorch. The training loop uses gradients to update weights with an optimizer, as in the backpropagation lesson. Compare its MSE with a deliberately uninformative baseline that returns the average crop for every input.

The trained MSE is close to the constant-crop baseline. These nearly identical inputs don't require an informative code: a decoder can do well by remembering their shared pattern. Optimization works as expected, but this test alone doesn't prove the latent captures useful variation. A solid evaluation needs distinct patterns, held-out crops, and a verification that perturbing the latent changes the reconstruction.

Even an autoencoder that reconstructs many distinct patterns accurately can't decode arbitrary random pairs of numbers. Its reconstruction loss only penalizes codes produced by the encoder for actual training inputs:

$$L_{\text{AE}} = \|x - g_\theta(f_\phi(x))\|_2^2$$

The encoder is free to arrange those codes without matching a chosen prior. It might place center-highlight crops near $(-1.8, 1.2)$ and lower-strip crops near $(1.7, -1.3)$. A standard-normal draw can then fall in a region poorly represented by encoded training crops, where reconstruction loss doesn't directly assess the decoder. Its output there may be useful, blurry, or implausible; this objective doesn't tell us which. Neither separated clusters nor failed random samples are inevitable. A separate model could also learn a distribution over deterministic codes.

A VAE learns distributions over codes

A VAE starts with a generative story: draw a latent $z$ from a prior $p(z)$, then draw an image from a decoder distribution $p_\theta(x\mid z)$. For an observed crop, inference runs that story backward: which latent values could have produced it? Instead of returning one code, the encoder predicts an approximate posterior, written $q_\phi(z \mid x)$, over those possible values.

The word approximate matters: computing the exact posterior $p_\theta(z \mid x)$ is generally impractical, so the encoder learns a tractable stand-in. This approximation method is variational inference. Here it's amortized: one encoder predicts posterior parameters for every crop instead of solving a fresh optimization problem for each input.

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 crop. The prior is the distribution we sample from when no input crop is provided. In $\mathcal{N}(0,I)$, each coordinate has mean zero and variance one, and the coordinates are independent. A diagonal posterior similarly has independent coordinates conditional on this crop, though its means and variances depend on the input. Mixing the posteriors of different crops needn't preserve that independence in the aggregate codes.

Why predict logvar rather than a raw standard deviation? An unconstrained linear layer can output negative values, while a standard deviation must be positive. Mathematically, exponentiating half a finite log variance makes a positive scale. Floating-point exponentiation can still underflow to zero or overflow to infinity. A positive transform such as softplus is another parameterization; log variance is convenient, not mandatory.

Those two vectors define a Gaussian for one crop. Next we keep that Gaussian near a sampleable prior, then we'll sample from it without blocking gradients.

Keep posteriors near a sampleable prior

If the encoder may put each crop's distribution anywhere it likes, prior samples still won't reliably land where the decoder has trained. Before writing the penalty, predict what it should do: it should be zero when the posterior matches the prior, and grow when means or variances move away.

VAEs measure that mismatch with Kullback-Leibler (KL) divergence. KL isn't a symmetric distance: direction matters. For a diagonal Gaussian posterior and standard-normal prior, the closed form is:^{[2]Reference 2Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114}

$$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.

Use the same two-dimensional posterior we'll sample from next: $\mu = [0.40, -0.20]$ and $\log\sigma^2 = [-0.70, 0.10]$. One coordinate at a time:

$$\begin{aligned} j=1 &\colon && 1 - 0.70 - (0.40)^2 - e^{-0.70} \approx -0.3566 \\ j=2 &\colon && 1 + 0.10 - (0.20)^2 - e^{0.10} \approx -0.0452 \\ D_{\mathrm{KL}} &= && -\tfrac{1}{2}(-0.3566-0.0452) \approx 0.2009 \end{aligned}$$

The next snippet checks that closed form on three posteriors, including this one.

KL pressure isn't free. If it dominates reconstruction, all inputs can be encoded nearly like the prior and the decoder may ignore $z$. We'll catch that failure after the training loop. First we need a way to sample from $q_\phi(z \mid x)$ without blocking gradients to $\mu$ and $\sigma$.

Sample while gradients still flow

Kingma and Welling reparameterize a Gaussian draw so reconstruction gradients can reach encoder parameters.^{[2]Reference 2Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114} Writing $z\sim\mathcal{N}(\mu,\sigma^2)$ specifies a distribution, not how a sampling API connects to autograd. PyTorch's ordinary Normal.sample() supplies no pathwise gradient to its mean or scale. Reparameterization makes that path explicit. Other estimators, including score-function gradients, can train stochastic models without differentiating the sampled value directly.^{[3]Reference 3Probability Distributionshttps://docs.pytorch.org/docs/2.14/distributions.html}

The reparameterization trick moves all stochasticity into an independent noise variable $\epsilon$ drawn from a standard normal distribution, then applies an algebraic transformation:

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

During the backward pass, $\epsilon$ is treated as a fixed constant tensor. The symbol $\odot$ means elementwise multiplication. For each coordinate, the partial derivatives are straightforward:

$$\frac{\partial z}{\partial \mu} = 1, \qquad \frac{\partial z}{\partial \sigma} = \epsilon$$

Reconstruction gradients $\partial L_{\text{recon}}/\partial z$ can then reach $\mu$, $\sigma$, and the encoder weights $\phi$. PyTorch's torch.distributions.Normal.rsample() implements this draw. Using sample() removes this reconstruction path to the encoder, but the decoder can still train, and an analytic KL term can still update encoder parameters.^{[3]Reference 3Probability Distributionshttps://docs.pytorch.org/docs/2.14/distributions.html}

Check the framework contract directly. Predict which draw keeps an autograd path, then verify the derivative through logvar: $\partial z_j/\partial\log\sigma_j^2=\tfrac12\sigma_j\epsilon_j=\tfrac12(z_j-\mu_j)$.

That sampled $z$ is what the decoder sees during training. The full path is now:

Generation later skips the encoder. It draws $z$ from the prior $p(z)$ and runs only the decoder.

The VAE objective

The VAE performs variational inference by optimizing an evidence lower bound (ELBO) on the data log-likelihood.^{[2]Reference 2Auto-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 crop likely after sampling its latent. The second discourages a posterior that moves too far from the prior. That bound equals $\log p_\theta(x)$ only when $q_\phi(z \mid x)$ matches the true posterior $p_\theta(z \mid x)$. Otherwise the leftover gap is $D_{\mathrm{KL}}(q_\phi(z \mid x)\,\|\,p_\theta(z \mid x))$.

In code we minimize the negative ELBO, or a weighted variant:

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

For $\beta=1$, this is the negative ELBO when L_recon estimates the selected negative log-likelihood. To connect that term to the pixel errors you've already calculated, let the decoder output a mean image $\hat{x}$ and assume independent Gaussian pixel noise with fixed variance $\tau^2$:

$$-\log p_\theta(x\mid z)=\frac{1}{2\tau^2}\sum_{i=1}^{N}(x_i-\hat{x}_i)^2+\frac{N}{2}\log(2\pi\tau^2)$$

For fixed $\tau$, the last term doesn't change with the weights. The first is scaled summed squared error, not automatically the mean MSE. A smaller observation variance penalizes the same pixel errors more strongly. Changing the reconstruction reduction without adjusting the KL coefficient changes the optimization problem.

The reconstruction term averages over posterior samples. Our code uses one reparameterized draw per example, giving a Monte Carlo estimate of that expectation. For a fixed model and example with finite variance, averaging more independent draws reduces the estimator's variance in proportion to their count, but requires more decoder work. The Gaussian KL term is evaluated analytically.^{[2]Reference 2Auto-Encoding Variational Bayes.https://arxiv.org/abs/1312.6114}

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

Higgins et al. studied $\beta$-VAE, increasing the KL weight to restrict latent information and encourage disentangled representations.^{[4]Reference 4beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Frameworkhttps://openreview.net/forum?id=Sy2fzU9gl} A factorized prior doesn't guarantee that learned coordinates represent independent real-world factors. Unsupervised disentanglement requires assumptions about the model and data.^{[5]Reference 5Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representationshttps://proceedings.mlr.press/v97/locatello19a.html} Measure reconstruction, latent usefulness, and prior samples on your task; a higher weight isn't automatically better.

These rows are an accounting exercise, not a reported experiment. Each row uses a different objective, so ranking the totals doesn't identify the best model. Log reconstruction and KL separately and inspect samples at each setting.

Train a tiny VAE

Now train a two-dimensional VAE on two families of synthetic crops: one with a bright center highlight and one with a bright lower strip. A constant image can't fit both well. After training, compare reconstruction from each posterior mean with reconstruction from the other family's means. Swapping codes should hurt if they carry pattern information. Finally, decode three prior draws and inspect their center and lower brightness rather than just their tensor shape.

The decoder returns Gaussian mean images in this MSE-based example. We display those means, without adding observation noise. The loss below is mean pixel error plus a chosen KL weight; it doesn't assume that this weight is the standard ELBO coefficient.

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.

The VAE's shape contract is small enough to write down. Here, B is batch size:

The reductions matter too. F.mse_loss averages pixel errors over the batch, while the KL term sums its two latent dimensions before averaging over the batch. That makes beta = 0.02 a scale choice for this setup, not a universal VAE setting. If you change latent width or loss reduction, retune it and keep logging both terms.

Swapping the families' codes raises MSE from 0.0005 to 0.1276, showing that the decoder uses input-dependent information. Prior draw 0 has a bright lower strip; draw 1 has a bright center; draw 2 is less distinct. Three draws can't establish distributional quality, and the reconstruction checks use training crops. Add held-out patterns and inspect full decoded grids before making a generalization claim.

Watch for posterior collapse

Low KL can look like success because each posterior matches the prior. But if reconstruction stays acceptable while changing $z$ barely changes the output, what happened? The decoder found a shortcut and stopped using its latent input. In information-theoretic terms, the mutual information $I(X;Z)$ between observations and latent codes can fall toward zero: knowing the sampled code no longer tells you much about which crop produced it. 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.^{[6]Reference 6Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937}

An expressive autoregressive decoder can model data through previously observed pixels or tokens while ignoring the global latent. Training can settle near $q_\phi(z\mid x)=p(z)$, with little input-specific information in the code.^{[7]Reference 7Generating Sentences from a Continuous Spacehttps://aclanthology.org/K16-1002/} That decoder can still generate varied, plausible outputs. An unconditional deterministic mean-image decoder optimized with MSE instead favors the dataset's mean image when it ignores $z$. Averaging is a consequence of that decoder and loss, not a definition of collapse.

There are two pressures hidden in the average KL. Let $q(z)$ be the mixture of encoder posteriors over the data and let $I_q(X;Z)$ measure information under that joint data/encoder distribution. Then:^{[8]Reference 8ELBO Surgery: Yet Another Way to Carve Up the Variational Evidence Lower Boundhttps://www.cs.columbia.edu/~blei/fogm/2022F/readings/HoffmanJohnson2016.pdf}

$$\mathbb{E}_X D_{\mathrm{KL}}(q_\phi(z\mid X)\,\|\,p(z)) = I_q(X;Z)+D_{\mathrm{KL}}(q(z)\,\|\,p(z))$$

Matching the aggregate codes to the prior and making every crop's posterior equal to the prior are different goals. If every posterior equals the prior, $I_q(X;Z)=0$: a sampled code cannot tell us which crop it came from. A large average KL, however, may reflect prior mismatch rather than useful information.

The training code already reports KL per latent dimension. One near-zero dimension can simply be unused capacity when the other distinguishes the two patterns. This illustrative monitor raises a warning only when KL is tiny in every dimension; it doesn't diagnose collapse by itself:

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.

Three interventions worth testing are:

1. KL annealing: Start with weak KL pressure and increase it toward the intended final weight. Bowman et al. anneal from zero to one for their likelihood-scaled objective.^{[7]Reference 7Generating Sentences from a Continuous Spacehttps://aclanthology.org/K16-1002/} Our mean-MSE example uses a different scale, so one isn't a universal endpoint. Zero KL weight also doesn't guarantee useful latents, as the nearly constant-crop experiment showed.
2. Free bits: The original formulation sums $\max(\tau,\mathbb{E}_X D_{\mathrm{KL},j})$ over latent groups, with each group's KL averaged over the minibatch before applying the floor.^{[9]Reference 9Improved Variational Inference with Inverse Autoregressive Flowhttps://arxiv.org/abs/1606.04934} Below $\tau$, that KL contribution has zero gradient; reconstruction can reward additional information without extra KL cost. This doesn't force a minimum mutual information. With natural logarithms, the threshold is measured in nats despite the name.
3. Decoder rebalancing: Reduce decoder capacity or weaken autoregressive conditioning. Bowman's word dropout removes some teacher-forced context so latent information becomes more useful.^{[7]Reference 7Generating Sentences from a Continuous Spacehttps://aclanthology.org/K16-1002/} It doesn't establish that the decoder must use $z$, and degrading a decoder can hurt quality.

Inspect interpolation rather than assuming it is meaningful. A decoder built from continuous layers changes continuously with $z$ in a deterministic autoencoder too; that mathematical property doesn't guarantee plausible images. The path $z(\alpha)=(1-\alpha)z_A+\alpha z_B$ can cross regions poorly represented during training in either model. Gaussian posterior support alone doesn't tell us where most probability mass lies or where decoding works well.

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.^{[6]Reference 6Neural 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 assignments are piecewise constant, so their ordinary derivatives don't supply a useful encoder training signal. VQ-VAE uses a straight-through gradient estimator for the encoder path, plus a codebook loss and a commitment loss so encoder outputs stay close to useful entries:^{[6]Reference 6Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937}

$$\lVert \operatorname{sg}[z_e(x)] - e \rVert_2^2 + \beta \lVert z_e(x) - \operatorname{sg}[e] \rVert_2^2$$

Here $\operatorname{sg}[\cdot]$ stops gradients, $e$ is the selected codebook vector, and this $\beta$ is the VQ-VAE commitment weight, not the $\beta$-VAE KL weight. In the original straight-through setup, reconstruction trains the decoder and encoder, but supplies no codebook gradient. The codebook loss updates selected entries; commitment updates the encoder. A moving-average codebook update is another option.^{[6]Reference 6Neural Discrete Representation Learning.https://arxiv.org/abs/1711.00937} The important distinction is the interface:

The nearest code doesn't change during this backward pass. Which tensors receive gradients from reconstruction alone, and which need the auxiliary losses? This two-coordinate example makes the routing visible:

A hand-picked VQ-VAE code ID can be decoded, but it isn't automatically a plausible sample. The learned prior supplies the combinations of code IDs that the generative model has been trained to use.

Continuous VAEs and discrete VQ-VAEs can both turn a large input into a smaller latent another model can consume. A later generator can learn the distribution of those latents instead of modeling every pixel directly.

Latent models reduce the expensive space

The same bottleneck contract scales beyond this tiny crop. A later generator can work in latent space, then hand its final representation to the decoder that turns it back into pixels. Rombach et al. train diffusion models in the latent space of pretrained autoencoders rather than applying all denoising operations directly to pixels.^{[10]Reference 10High-Resolution Image Synthesis with Latent Diffusion Models.https://arxiv.org/abs/2112.10752} They write the spatial downsampling factor as $f = H/h = W/w$ and evaluate the efficiency and quality trade-off across values from 2 to 32, including $f=4$ and $f=8$. The autoencoder's job is precise: compress before the expensive generation steps, then decode a final latent back to pixels.

Their compression stage uses perceptual and adversarial losses, with weak KL or vector-quantization regularization. A separate diffusion model learns the latent distribution.^{[10]Reference 10High-Resolution Image Synthesis with Latent Diffusion Models.https://arxiv.org/abs/2112.10752} Sampling Gaussian noise and decoding it directly isn't equivalent to running that generator: diffusion must first transform noise into a latent that fits the learned distribution.

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. On a 512 by 512 grid that is 262,144 positions becoming 4,096:

Channels and decoder quality still matter. For a hypothetical RGB image and four-channel latent at $f=8$, scalar values drop from $512^2\times3$ to $64^2\times4$, a 48-fold reduction, rather than 64. Neither ratio predicts runtime or compressed-file size by itself. Rombach et al. find that aggressive compression hurts their downstream results; detail discarded by the encoder can't be uniquely recovered from its code. A generator may synthesize plausible detail without reproducing the original pixels.

Debug the latent interface

Start by logging reconstruction and KL losses separately. A low total can hide a reconstruction term that dominates, or a KL term that has collapsed toward zero.

If sampling or KL returns NaN or Inf, check the scale path first: sigma = exp(0.5 * logvar), not exp(logvar). Inspect the range of logvar before adding an explicit clamp, since a clamp can hide unstable optimization.

Keep the two distributions distinct. The encoder supplies $q_\phi(z \mid x)$ for an observed crop; generation samples from the prior $p(z)$. If decoded outputs barely change when $z$ changes, the decoder may be ignoring its latent input even when reconstruction looks fine.

For a VQ-VAE, log code usage as well. If only a few code IDs are selected, the codebook's capacity is going unused and the discrete bottleneck needs a closer look.

---

Test what the code preserves