70 - Autoencoders: Compress, Reconstruct and Understand What a Latent Space Is

An autoencoder is a network trained on a task that seems useless: reconstructing its own input. The trick lies in the fact that in the middle of the network there is a bottleneck, a layer with very few numbers, through which all the information must pass. To reconstruct well, the network is forced to decide what is worth keeping of an image: it learns to compress and, as a side effect, it learns a representation of the data. It is the building block from which the VAE arises, the subject of the next lessons, so it is worth understanding it well, with its limits. In this lesson you build it in PyTorch on handwritten digits, you compare it with a classical technique (PCA) which you discover is, in one precise case, the same thing, you measure how much nonlinearity helps, you use it to remove noise and finally you find out why it is not enough to generate new images.

The words of this lesson

The encoder transforms the input x into a short vector z, the code; the decoder transforms z into a reconstruction x̂ that must resemble x. The latent space is the space of possible codes: "latent" means hidden, because it is never observed directly in the data but is a representation that the network has chosen, and with k = 2 numbers it is a plane. The bottleneck is the narrowest layer: if it were as wide as the input the network could copy without learning anything. The reconstruction error is the difference between x and x̂, here the mean squared error (MSE) over the pixels. PCA finds the k directions along which the data vary the most and projects the data onto them, and it is linear. A denoising autoencoder receives as input an image corrupted by noise and must reconstruct the clean one. All this is unsupervised (or self-supervised) learning: no labels are needed, the input itself acts as the target.

The analogy: describing a face over the phone

You have to describe a person's face to a friend using only ten words, and the friend must then draw it. You cannot give the colour of every single pixel: you must choose the features that matter ("round", "glasses", "short hair", "smiling"). The ten words are the code z; you are the encoder, the friend is the decoder. The fewer the words, the more approximate the drawing; the more you and your friend train together, the more efficient your vocabulary becomes. The "vocabulary" that emerges is the latent space. Note the limit, though: if the friend, at random, receives ten words that nobody has ever said (for example "round" and "square" together), the drawing will be meaningless. We will come back to this point.

In the drawing the images have 64 pixels and the hidden layer has 128 units, as in the code that follows; the code z is the only very narrow layer. Above all the drawing recalls one thing: no label appears anywhere, the loss compares the reconstruction with the input.

The mathematics

Let E be the encoder and D the decoder, with parameters to be learned. For an image x of d pixels (here d = 64):

z  = E(x)                     (code, k numbers, k ≪ d)
x̂ = D(z) = D(E(x))            (reconstruction, d numbers)
L  = (1/d) · Σ_pixel (x_i − x̂_i)²    (MSE, averaged over the images)

There is no label: training is an ordinary gradient descent on L. The ability to reconstruct depends on the value of k: if k is small, the code can keep only the essentials.

A special case: the linear autoencoder is PCA

If encoder and decoder are linear (z = W₁ x, x̂ = W₂ z), the problem becomes finding two matrices whose product W₂W₁ has rank at most k and minimises the squared reconstruction error. It is a classical problem, solved by the Eckart-Young theorem: the best rank-k approximation is obtained with the singular value decomposition (SVD), keeping only the k largest singular values, and that is what PCA does (the scikit-learn documentation describes it as the SVD of the centred data). So a well-trained linear autoencoder converges to a subspace equivalent to that of PCA (a result by Baldi and Hornik, 1989; both works are bibliographic citations, not consulted online for this lesson). The minimum error has a closed formula: the sum of the squares of the discarded singular values, divided by the number of numbers in the dataset:

MSE_k = ( Σ_{i>k} s_i² ) / (N · d)

We can check it with the code, without training anything. The first block uses scikit-learn's load_digits in place of MNIST: 1797 8×8 images with values from 0 to 16, split 70/30 with stratification (1257 for training, 540 for testing) and rescaled to [0, 1]. The MNIST download was not run for this version, so the absolute numbers are not comparable with those of the previous version of the lesson.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.set_num_threads(2)
X, y = load_digits(return_X_y=True)                                # 1797 digits, 8x8 pixels, values 0..16
Xtr, Xte, ytr, yte = train_test_split(X / 16.0, y, test_size=0.3, random_state=0, stratify=y)
Xtr, Xte = torch.tensor(Xtr, dtype=torch.float32), torch.tensor(Xte, dtype=torch.float32)
D = Xtr.shape[1]                                                   # 64 pixels

# --- the best possible LINEAR compression: PCA (computed with an SVD, no training) ---
mean = Xtr.mean(dim=0)
U, S, Vt = torch.linalg.svd(Xtr - mean, full_matrices=False)       # rows of Vt = principal directions

def pca_mse(k):
    V = Vt[:k]                                                     # (k, 64)
    codes = (Xte - mean) @ V.T                                     # compress to k numbers
    recon = codes @ V + mean                                       # decompress
    return F.mse_loss(recon, Xte).item()

print(Xtr.shape, Xte.shape)
print("baseline: predict the mean image  MSE", round(F.mse_loss(mean.expand_as(Xte), Xte).item(), 4))
for k in (2, 8, 16):
    print(f"PCA with k={k:2d}:  test MSE {pca_mse(k):.4f}")

torch.Size([1257, 64]) torch.Size([540, 64])
baseline: predict the mean image  MSE 0.0739
PCA with k= 2:  test MSE 0.0528
PCA with k= 8:  test MSE 0.0245
PCA with k=16:  test MSE 0.0115

torch.linalg.svd computes the decomposition of X − mean: the rows of Vt are the principal directions. pca_mse(k) compresses with the first k directions (codes = (Xte − mean) @ V.T), reconstructs (codes @ V + mean) and measures the MSE on the test set. The baseline is always predicting the mean image (MSE 0.0739): every compression must do better. With k = 2 numbers PCA drops to 0.0528 (little), with 8 to 0.0245, with 16 to 0.0115: a compression of 64/16 = 4 times with an error equal to about a sixth of the baseline.

# PCA error from the singular values alone (Eckart-Young): MSE_k = sum of the discarded squared singular values / (N * 64)
N = Xtr.shape[0]
for k in (2, 8, 16):
    predicted = (S[k:] ** 2).sum().item() / (N * D)                # no reconstruction needed
    V = Vt[:k]
    recon = ((Xtr - mean) @ V.T) @ V + mean
    measured = F.mse_loss(recon, Xtr).item()                        # the TRAINING reconstruction error, computed directly
    print(f"k={k:2d}  from singular values {predicted:.5f}   measured {measured:.5f}")

k= 2  from singular values 0.05233   measured 0.05233
k= 8  from singular values 0.02376   measured 0.02376
k=16  from singular values 0.01093   measured 0.01093

The code computes the error with the formula ((S[k:]**2).sum() / (N·d)) and then, as a check, actually reconstructs the training images and measures the MSE: they coincide to the fifth decimal place. Careful: these are errors on the training set and those of the previous block are on the test set, so they have close but not identical values (0.02376 against 0.0245 at k = 8). Now the Eckart-Young formula is no longer a theorem to take on trust: it is a number that you saw come out equal.

The autoencoder in PyTorch, linear and nonlinear

We build two autoencoders: a linear one (a single layer going down and one going up, no nonlinearity) and a nonlinear one (64 → 128 → k → 128 → 64 with ReLU and a sigmoid at the output). We train them for 300 epochs for three values of k, with Adam at 0.001 and mini-batches of 64, and compare the test errors with those of PCA. In the constructor, if linear=True, encoder and decoder are a single nn.Linear each. In the forward the output goes through a sigmoid, to keep the pixels in (0, 1), only in the nonlinear case: a sigmoid at the output makes the network nonlinear and would break the comparison with PCA. The noise parameter of train_ae is for the next experiment.

class AutoEncoder(nn.Module):
    def __init__(self, k, hidden=128, linear=False):
        super().__init__()
        self.linear = linear
        act = nn.ReLU()
        if linear:                                                 # one linear map down, one linear map up
            self.enc = nn.Linear(D, k)
            self.dec = nn.Linear(k, D)
        else:
            self.enc = nn.Sequential(nn.Linear(D, hidden), act, nn.Linear(hidden, k))
            self.dec = nn.Sequential(nn.Linear(k, hidden), act, nn.Linear(hidden, D))
    def forward(self, x):
        z = self.enc(x)                                            # the code: k numbers
        out = self.dec(z)
        return (out if self.linear else torch.sigmoid(out)), z     # nonlinear AE: pixels squashed into (0, 1)

def train_ae(k, linear=False, epochs=300, seed=0, noise=0.0):
    torch.manual_seed(seed)
    model = AutoEncoder(k, linear=linear)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    for _ in range(epochs):
        perm = torch.randperm(len(Xtr))
        for i in range(0, len(Xtr), 64):
            x = Xtr[perm[i:i + 64]]
            x_in = (x + noise * torch.randn_like(x)).clamp(0, 1) if noise > 0 else x   # denoising: corrupt the INPUT only
            recon, _ = model(x_in)
            loss = F.mse_loss(recon, x)                            # ... but compare with the CLEAN image
            opt.zero_grad(); loss.backward(); opt.step()
    return model

@torch.no_grad()
def test_mse(model):
    return F.mse_loss(model(Xte)[0], Xte).item()

for k in (2, 8, 16):
    lin = train_ae(k, linear=True)
    non = train_ae(k)
    print(f"k={k:2d}   linear AE {test_mse(lin):.4f}   nonlinear AE {test_mse(non):.4f}   (PCA {pca_mse(k):.4f})")

k= 2   linear AE 0.0532   nonlinear AE 0.0355   (PCA 0.0528)
k= 8   linear AE 0.0251   nonlinear AE 0.0115   (PCA 0.0245)
k=16   linear AE 0.0116   nonlinear AE 0.0058   (PCA 0.0115)

The linear autoencoder comes close to PCA: 0.0251 against 0.0245 (k = 8) and 0.0116 against 0.0115 (k = 16); with k = 2 it is 0.0532 against 0.0528. It always stays a hair above PCA, as one expects from an optimiser that after a finite number of epochs has not yet reached the exact optimum. The nonlinear autoencoder does clearly better: for the same k, the error is lower by 33% (0.0355 against 0.0528 with 2 numbers), by 53% (0.0115 against 0.0245 with 8) and by 50% (0.0058 against 0.0115 with 16). It is the gain of nonlinearity: the data lie on a "curved surface" inside the pixel space, not on a flat subspace, and a nonlinear network can follow it. One seed and 540 test examples: the order PCA, linear, nonlinear is clear, the second decimal place is not.

Removing noise: the denoising autoencoder

A surprisingly useful variant: you corrupt the input with noise but ask to reconstruct the clean image. Since random noise cannot be compressed into k numbers, the bottleneck discards it. The noise is Gaussian with standard deviation 0.5, added to the pixels and then clipped to [0, 1]; the line loss = F.mse_loss(recon, x) in train_ae compares with the clean image even when the input is noisy.

# Denoising: the network sees a NOISY image and must reconstruct the CLEAN one.
torch.manual_seed(1)
noisy_test = (Xte + 0.5 * torch.randn_like(Xte)).clamp(0, 1)
plain = train_ae(16)                                  # trained on clean images
denoising = train_ae(16, noise=0.5)                   # trained on noisy inputs -> clean targets
with torch.no_grad():
    print("noisy input vs clean        MSE:", round(F.mse_loss(noisy_test, Xte).item(), 4))
    print("plain AE on noisy input     MSE:", round(F.mse_loss(plain(noisy_test)[0], Xte).item(), 4))
    print("denoising AE on noisy input MSE:", round(F.mse_loss(denoising(noisy_test)[0], Xte).item(), 4))

noisy input vs clean        MSE: 0.1175
plain AE on noisy input     MSE: 0.0877
denoising AE on noisy input MSE: 0.0317

The noisy image has MSE 0.1175 with respect to the clean one. An ordinary autoencoder (trained on clean images), given the noisy image, reduces it to 0.0877: it helps a little because the bottleneck filters, but it was not trained to do so. The denoising autoencoder, trained with noisy input and clean target, reaches 0.0317: a cut of almost four times with respect to the starting noise. The lesson is subtle: by changing only what is fed as input, the same network learns a different task.

The limit: an autoencoder is not a generator

If the decoder is good at turning codes into images, why not pick a code at random and make a new image from it? Because we do not know where to pick. The network has no constraint on the position and scale of the codes: it chooses the ones that suit it. Let us measure it with an autoencoder with a two-dimensional code (k = 2), so that the latent space is a plane, and look at where the codes of real images fall and how "full" the area around them is. The function computes, for each sampled point, the distance from the nearest real code (torch.cdist, then min), and takes its median over a reference set of 900 real codes.

# Can we GENERATE with an autoencoder? Where do the codes of real digits live, and what lies between?
ae2 = train_ae(2)                                        # 2-dimensional code, so that we can reason about it
with torch.no_grad():
    codes = ae2.enc(Xtr)                                 # (1257, 2): where the REAL digits live in code space
print("code mean:", [f"{v:.2f}" for v in codes.mean(0)], "  std:", [f"{v:.2f}" for v in codes.std(0)])

reference = codes[:900]                                  # the "known" real codes
def median_distance_to_nearest_real_code(z):
    return torch.cdist(z, reference).min(dim=1).values.median().item()

torch.manual_seed(0)
with torch.no_grad():
    z_real = codes[900:]                                 # other real codes (never in the reference set)
    z_gauss = torch.randn(5000, 2)                       # a standard Gaussian: the prior a VAE will use
    z_box = (torch.rand(5000, 2) - 0.5) * 2 * codes.abs().max(0).values   # uniform in the bounding box of the codes
print("median distance to the nearest real code:")
print("  other real codes     :", round(median_distance_to_nearest_real_code(z_real), 3))
print("  N(0, 1) samples      :", round(median_distance_to_nearest_real_code(z_gauss), 3))
print("  uniform in the box   :", round(median_distance_to_nearest_real_code(z_box), 3))

code mean: ['2.89', '-1.14']   std: ['3.82', '4.40']
median distance to the nearest real code:
  other real codes     : 0.198
  N(0, 1) samples      : 0.244
  uniform in the box   : 1.697

The codes of the 1257 training examples are centred around (2.89, −1.14) with standard deviation 3.8-4.4: far from the 0-1 scale usually used for sampling. A uniform sample in the "box" containing all the codes has a median at distance 1.70 from the nearest real code, about eight times the distance between real codes (0.198): a large part of the box is empty, and the decoder has never seen what to do there. The points with a standard normal distribution fall near the cloud (0.244), but only because the cloud is so large that it includes the origin: there is no reason, in the construction of the autoencoder, for this to happen. The point is methodological: the autoencoder does not tell you which distribution to sample from. If you do not know the distribution of the codes, you cannot generate. The VAE, in the next two lessons, solves exactly this: it forces the latent space to have a known shape (a standard normal), so that sampling from that shape gives sensible codes.

Trade-offs

The size of the code is the first compromise: small means strong compression and loss of detail, large means faithful reconstruction but the risk of copying without learning a useful representation. Between linear and nonlinear, PCA is deterministic, instantaneous (one SVD) and interpretable; the nonlinear autoencoder compresses better (here between 33% and 53% less error) but must be trained, has hyperparameters and can give different results per seed. Finally the MSE favours average, slightly blurred reconstructions, as in the lesson on frame prediction.

Real pitfalls

1. Bottleneck too wide. With a code as large as the input the network learns the identity: low MSE and a useless representation. A perfect reconstruction is not a good sign by itself.
2. Thinking the code is interpretable. No component of z means "thickness" or "slant" by construction. If you want dimensions with a meaning, additional constraints are needed.
3. Sampling codes at random from the ordinary autoencoder. The latent space has holes and an arbitrary scale; the images decoded at those points have no guarantee of meaning.
4. Evaluating only with the MSE. A good MSE can coexist with blurred images; look at the images and use additional metrics if needed.
5. Comparing a non-converged linear autoencoder with PCA. Here it always stays a hair above PCA (0.0532 against 0.0528 at k = 2): PCA is an exact reference, the autoencoder must converge to it.
6. Using a denoising autoencoder with noise different from the real one. If in production the noise has a different shape from the one seen in training, the model may not clean it; I did not measure it.

The six pitfalls have in common trust in a single number: a low MSE, a code that looks tidy, a comparison with PCA. Each is unmasked by looking at a second indicator, be it the reconstructed image, the distribution of the codes or the exact reference of PCA.

A worked example by hand: the PCA of three points and perfect compression

Three points in the plane: (1, 2), (2, 4), (3, 6). They all lie on the line y = 2x, so only one number per point is really needed, not two. Centring: the mean is (2, 4); the centred points are (−1, −2), (0, 0), (1, 2). Principal direction: the centred points are multiples of (1, 2); the unit direction is u = (1, 2)/√5 = (0.447, 0.894), and it is the one PCA finds with k = 1.

Compressing: the code of each point is the projection onto u: (−1, −2)·u = −√5 = −2.236, 0, +2.236, three numbers instead of six. Reconstructing: the reconstructed point is mean + code·u; for the first, (2, 4) + (−2.236)·(0.447, 0.894) = (1, 2), error 0. If the points had been (1, 2.2), (2, 3.8), (3, 6), the best line would have been an approximation and the reconstruction would have had an error equal to what lies "off the line"; the Eckart-Young theorem says that no other line does better. The nonlinear autoencoder generalises this idea: instead of a line it can use a curve, and it compresses better data that lie on a curve, like handwritten digits.

A couple of calculations as an exercise: the compression factor of a 64 → 16 → 64 autoencoder is 64/16 = 4; a linear 64 → 8 → 64 autoencoder with biases has 64·8 + 8 + 8·64 + 64 = 1096 parameters; and if the sum of the squares of the discarded singular values of 10 images of 4 pixels is 120, the reconstruction MSE is 120 / (10·4) = 3.0.

Frequently asked questions

When does it make sense to use an autoencoder instead of PCA? When the data have nonlinear structure and you need compact codes, or when you want a basis for a generative model. If the data are almost linear, PCA is simpler, instantaneous and interpretable.

Can the autoencoder be used to find anomalies? It is a very common application: you train on "normal" data and use the reconstruction error as a score. I did not measure it here, so I give no numbers.

How do I choose k? Vary k and look at the error curve: there is often an "elbow" beyond which the improvements are small. Here the PCA error drops a lot between k = 2 and 8 (from 0.0528 to 0.0245) and less between 8 and 16 (to 0.0115).

Can the code be used for something else? As the input of a classifier, to search for similar images, to visualise data with k = 2: these are common uses, not results of this lesson.

Does a deeper autoencoder compress better? More capacity allows approximating more complex structures, but increases the risk of learning the identity if the bottleneck is not narrow enough. Here I used a single hidden layer of 128 and did not try other depths.

What you can do now

You can build an autoencoder in PyTorch and you know that it compresses through a bottleneck without labels. In the linear case it comes close to PCA (0.0251 against 0.0245 at k = 8) and the Eckart-Young formula gives the error from the SVD alone (0.05233 at k = 2, both from the formula and measured on the training set). Nonlinearity lowers the error by 33-53% for the same code. If you corrupt the input you get a denoiser (from 0.1175 to 0.0317). The limit is that the space of codes has no known shape: with k = 2 the codes are centred at (2.89, −1.14) with standard deviation 3.8-4.4, and uniform samples in their box fall on average 1.70 from the nearest real code. The VAE, in the next lesson, adds what is missing: a distribution over the latent space.

In the series, the previous lesson is FractalNet: Deep Networks Without Explicit Shortcuts, With Drop-Path.