Autoencoders: Compression as Representation

"My whole job is to forget you correctly. Anyone can throw away pixels; the art is throwing away the ones you will not miss."

A Bottleneck Layer With Excellent Taste

In the previous chapter you saw the landscape of generative models and the central role of the latent space, the low-dimensional surface that natural images cluster near inside their vast pixel space. This section gives you the first concrete way to find that surface, and it does so with an idea older and simpler than any generative theory: compression. You compressed images with classical transforms in Chapter 4, throwing away high-frequency coefficients that the eye barely notices. An autoencoder does the same thing, except it learns the transform from the data instead of fixing it in advance, and the transform it learns is tuned to the specific images it sees rather than to images in general. The illustration below captures the bottleneck in a single image.

1. The Encoder, the Bottleneck, and the Decoder Beginner

An autoencoder has three parts. An encoder $f$ maps an input $x$ (an image, flattened or kept as a tensor) to a code $z = f(x)$ that lives in a much smaller space. A decoder $g$ maps that code back to the input space, producing a reconstruction $\hat{x} = g(z)$. The code $z$ sits in the bottleneck, the narrowest layer of the network, and its dimension is the single most important design choice. Training minimizes a reconstruction loss that measures how close $\hat{x}$ is to $x$, most commonly the mean squared error

averaged over the training set. Nothing else is involved: no labels, no target other than the input itself. This is why the autoencoder is the purest example of self-supervised learning, where the supervision signal is manufactured from the unlabeled data. Figure 31.1.1 shows the hourglass shape that gives the architecture its name.

The code below implements precisely this for MNIST, whose images are $28 \times 28 = 784$ pixels. The encoder is a small multilayer perceptron that shrinks 784 inputs down to 32; the decoder mirrors it back up to 784 with a sigmoid so the outputs live in $[0, 1]$ like the normalized pixels. Read the two nn.Sequential stacks against the hourglass of Figure 31.1.1.

# Minimal undercomplete autoencoder for MNIST: an encoder MLP that
# funnels 784 pixels into a 32-dim bottleneck, and a decoder that mirrors
# it back. The forward pass returns both the reconstruction and the code.
import torch
import torch.nn as nn

class Autoencoder(nn.Module):
    def __init__(self, in_dim=784, code_dim=32):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(in_dim, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, code_dim),          # bottleneck: 32-dim code
        )
        self.decoder = nn.Sequential(
            nn.Linear(code_dim, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, in_dim), nn.Sigmoid(),  # outputs in [0, 1]
        )

    def forward(self, x):
        z = self.encoder(x)        # compress
        x_hat = self.decoder(z)    # reconstruct
        return x_hat, z

model = Autoencoder()
x = torch.rand(8, 784)            # a batch of 8 flattened images
x_hat, z = model(x)
print("code shape:", z.shape)         # code shape: torch.Size([8, 32])
print("reconstruction shape:", x_hat.shape)  # torch.Size([8, 784])

Training is an ordinary supervised loop where the target happens to be the input. The next block runs it on real MNIST for a few epochs and reports the reconstruction error falling, which is the signal that the code is capturing more and more of each digit.

# Train the autoencoder on real MNIST as a self-supervised loop:
# the target is the input image itself, the labels are discarded, and
# the falling MSE signals that 32 numbers increasingly suffice per digit.
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
                         transforms.Lambda(lambda t: t.view(-1))])  # flatten to 784
train = datasets.MNIST("./data", train=True, download=True, transform=tf)
loader = DataLoader(train, batch_size=256, shuffle=True)

model = Autoencoder()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()

for epoch in range(5):
    total = 0.0
    for x, _ in loader:                 # labels are ignored: self-supervised
        x_hat, _ = model(x)
        loss = loss_fn(x_hat, x)         # target IS the input
        opt.zero_grad(); loss.backward(); opt.step()
        total += loss.item() * x.size(0)
    print(f"epoch {epoch}: mse = {total / len(train):.4f}")
# epoch 0: mse = 0.0461
# epoch 1: mse = 0.0292
# ...
# epoch 4: mse = 0.0204

2. Why the Bottleneck Must Be Undercomplete Beginner

Why force the code to be small at all? Imagine the opposite: a code dimension equal to or larger than the input, an overcomplete bottleneck. Then the encoder and decoder can collude to learn the identity function, copying each pixel through untouched. Reconstruction is perfect, the loss is zero, and the code has learned nothing about the structure of digits; it is just the image again in a different layout. A code that perfectly reconstructs by memorizing is useless as a representation.

An undercomplete bottleneck (code dimension smaller than the input) makes the identity shortcut impossible. There are simply not enough numbers to store every pixel, so the network must spend its limited code budget on whatever is most reconstructable: the regularities shared across digits. It learns that strokes are continuous, that a digit has roughly one connected blob of ink, that certain pen movements are common and others never occur. The code becomes a description in terms of these regularities rather than in terms of raw pixels. The bottleneck width is therefore a knob on an explicit trade: too wide and the code memorizes, too narrow and the reconstruction degrades because genuine variation is being discarded. Table 31.1.1 shows the trade concretely on MNIST.

The two-dimensional code in the first row is special because you can plot it. The next block trains a code-2 autoencoder and scatters the codes of test digits, colored by their true label, to see whether the bottleneck has organized the digits without ever being told what a label is.

# Train a code-2 autoencoder and scatter the test-set codes colored by
# their true label. The labels are used only for coloring, never for
# training, so any clustering is evidence the bottleneck found structure.
import torch, matplotlib.pyplot as plt

ae2 = Autoencoder(code_dim=2)            # 2D bottleneck so we can plot it
# ... (train ae2 with the loop from subsection 1) ...

ae2.eval()
codes, labels = [], []
with torch.no_grad():
    for x, y in DataLoader(test_set, batch_size=512):
        _, z = ae2(x)
        codes.append(z); labels.append(y)
codes = torch.cat(codes).numpy()
labels = torch.cat(labels).numpy()

plt.scatter(codes[:, 0], codes[:, 1], c=labels, cmap="tab10", s=4)
plt.colorbar(label="true digit"); plt.title("2D autoencoder latent space")
plt.show()
# The 10 digit classes form loosely separated clusters,
# even though the network never saw a single label.

3. The Linear Case Is PCA Intermediate

There is a beautiful special case that anchors the autoencoder to classical statistics. Strip out the nonlinearities: let the encoder be a single linear map $z = W_e x$ and the decoder a single linear map $\hat{x} = W_d z$, and train with mean squared error. The optimal solution is then a familiar one. The subspace spanned by the decoder's columns is exactly the subspace of the top principal components of the data, the same principal component analysis you would get from the eigenvectors of the covariance matrix. A linear undercomplete autoencoder, in other words, learns to project onto the directions of greatest variance, which is the optimal linear compression in the squared-error sense. The two views agree for a simple reason. The variance a projection keeps and the squared error it discards always add up to the fixed total variance of the data, so keeping the most variance and losing the least reconstruction error are the same objective seen from two sides.

The formal statement is the Eckart-Young-Mirsky theorem. With centered data and covariance $\boldsymbol{\Sigma} = \tfrac{1}{N}\sum_i \mathbf{x}_i \mathbf{x}_i^\top = \mathbf{U}\boldsymbol{\Lambda}\mathbf{U}^\top$ and eigenvalues sorted $\lambda_1 \ge \lambda_2 \ge \cdots \ge \lambda_d$, the rank-$k$ linear map that minimizes the reconstruction error has decoder columns spanning $\operatorname{span}\{\mathbf{u}_1,\dots,\mathbf{u}_k\}$, the top-$k$ eigenvectors of $\boldsymbol{\Sigma}$, and the error it cannot avoid is exactly the discarded tail of the spectrum:

Any optimum is related to the principal axes by an invertible map of the code, so it spans the same top principal subspace rather than reproducing the eigenvectors exactly. Nonlinearity is precisely what lets a deep autoencoder beat this bound, following a curved manifold instead of a flat subspace.

This is worth holding onto for two reasons. First, it tells you what an autoencoder buys over PCA: the nonlinear encoder and decoder let it learn a curved low-dimensional surface rather than a flat subspace, which captures the genuinely nonlinear structure of images (a rotated digit traces a curve through pixel space, not a line). Second, it gives you a sanity baseline. If your fancy nonlinear autoencoder does not beat PCA's reconstruction at the same code dimension, something is wrong with your training. The next block makes the comparison directly.

# Sanity baseline: a 32-component PCA is the best LINEAR compressor at
# this code size. A trained nonlinear autoencoder should beat its
# reconstruction MSE; if it does not, the training is broken.
import torch
from sklearn.decomposition import PCA

# X_train: (N, 784) tensor of flattened MNIST, X_test similarly
pca = PCA(n_components=32).fit(X_train.numpy())
recon = pca.inverse_transform(pca.transform(X_test.numpy()))
pca_mse = ((X_test.numpy() - recon) ** 2).mean()
print(f"PCA-32 reconstruction MSE: {pca_mse:.4f}")    # PCA-32 reconstruction MSE: 0.0241

# A trained nonlinear autoencoder with the same 32-dim code should do better,
# because it can bend the latent surface to follow the data manifold.
print(f"Nonlinear AE-32 MSE:       {ae_mse:.4f}")      # Nonlinear AE-32 MSE: 0.0204

4. The Fatal Flaw: A Latent Space Full of Holes Intermediate

Now the catch that makes this whole chapter necessary. We have a generator-shaped object: an encoder, a decoder, and a low-dimensional code. It is tempting to think we can generate new digits by picking a random code $z$ and decoding it. We cannot, and the reason is structural. Nothing in the training objective ever asked the codes to fill the latent space in any particular way. The encoder is free to scatter the training codes wherever minimizing reconstruction is easiest, leaving large empty regions between clusters. Decode a point from one of those empty regions and you get an incoherent smear, because the decoder was never trained on anything that maps there. The illustration below makes the point with a postal-worker metaphor.

Two distinct problems live here. First, the codes have no fixed scale or shape, so you do not even know what range to sample from; the codes of subsection 2 might span $[-40, 40]$ in one axis and $[-2, 2]$ in another. Second, even sampling within the occupied range, you will frequently land between clusters where the decoder has never been trained, in a hole. Figure 31.1.2 contrasts the swiss-cheese latent of a plain autoencoder with the dense, fillable latent we want.

The experiment below makes the failure visceral. We sample a code from a standard normal (a perfectly reasonable guess for where codes might live) and decode it; the result is noise, because that code almost certainly fell in a hole.

# Show why a plain autoencoder cannot generate: sample a random Gaussian
# code and decode it. Because training never constrained where valid codes
# live, this code almost surely lands in a hole and decodes to a smear.
import torch

model.eval()
with torch.no_grad():
    z_random = torch.randn(1, 32)        # a guess at "a typical code"
    fake = model.decoder(z_random)       # decode it
print("output range:", fake.min().item(), "to", fake.max().item())
# output range: 0.41 to 0.58   <- a flat mid-gray band, no black ink or white paper
# Displaying `fake` shows an incoherent gray smear, not a digit.
# The plain autoencoder has no notion of which codes are "valid",
# so generation by random sampling fails. This is the gap the VAE closes.