Level N6 · Foundations · Classic networks · runs in your browser

Autoencoders and VAEs

Can a network learn to squeeze a picture into a few numbers and back?

Side trip · best after level 10 · Probability and sampling

A handwritten digit is 784 numbers, but most of those numbers are predictable: dark background, a stroke that bends a certain way. How few numbers are enough to describe it? This level trains a network to squeeze a digit into 2 numbers and rebuild it from them. Those few numbers are called the code (or the latent). The same idea lets image generators work on small codes instead of millions of pixels (level D3).

1. Squeeze and rebuild

An autoencoder is two small networks, one after the other:

  1. the encoder maps 784 pixels → 128 hidden numbers → a code of 2 numbers;
  2. the decoder maps the 2-number code → 128 → 784 pixels, each between 0 and 1.

The narrow middle is called the bottleneck. Everything the decoder knows about the picture has to pass through it.

Number A 28×28 picture has 784 numbers. The code has 2. How many times fewer numbers is the code?
🔒 Answer the question above to unlock

Batches work as in level 5: a batch of pictures is a table with one row per picture, and each layer is X @ W + b.

Shape A batch of 32 pictures has shape (32, 784). The encoder maps each picture to a 2-number code. What shape is the batch of codes?
I got stuck here If the target is the input itself, why can’t the network just copy it?

It would, if it could: copying gives zero loss. The bottleneck forbids it. 784 numbers can’t pass through 2 unchanged, so the encoder has to keep what matters most (which digit, how slanted, how thick) and drop the rest. Nobody tells it what matters; minimizing the rebuild loss decides. No labels are used anywhere: the picture is its own target.

2. Training: rebuild the input

🔒 Answer the question above to unlock

Here is the whole autoencoder. Each block has one cell per number.

Squeeze 784 numbers through 2

Press Next to follow one picture through the network.

Drag to turn · click, then scroll to zoom
Press Next to start. · 202,258 parameters in all
Layers (5)

Look at the sizes: each block holds as many cells as it has numbers, so the code in the middle is almost nothing.

forward (this step) Select a layer to see its numbers.

The loss compares the rebuilt picture x^\hat{x} with the original xx, pixel by pixel. The easiest one to read is the mean squared error: subtract, square, average.

Number A 4-pixel picture x = [1, 0, 1, 1] is rebuilt as x̂ = [0.5, 0, 1, 0.5]. What is the mean squared error?
🔒 Answer the question above to unlock

The real autoencoder in demo.py trains for 6 epochs (an epoch is one pass over all 60,000 training digits) and takes about 12 seconds on a laptop CPU, for this autoencoder and a second model together. Afterwards its mean squared error on test digits it never saw is 0.044 per pixel.

To train, demo.py uses a different rebuild loss. Every pixel is between 0 and 1, so it treats each pixel like a yes/no answer and uses the binary cross-entropy from level 3, −(xln⁡x^+(1−x)ln⁡(1−x^))-(x \ln \hat{x} + (1 - x)\ln(1 - \hat{x})), summed over the 784 pixels. It reports the mean squared error only because that number is easy to read. Both losses are 0 for a perfect rebuild. To run it yourself, download demo.py into a folder of its own and run python demo.py there (setup).

np.mean averages every number in an array, and ** squares each number.

CodeWrite the mean squared error between a picture x and its rebuild xₕₐₜ (both NumPy arrays of the same shape).

Enter keeps the indent · Tab indents · Esc then Tab leaves the editor · ⌘/Ctrl + Enter runs

3. Move through the code space

🔒 Answer the question above to unlock

Because the code is 2 numbers, every code is a point on a flat map, and every point on the map is a picture: give it to the decoder and it draws something. Codes of real digits fall somewhere on that map. To make a new digit, you would pick a random point and decode it, using the same mean and standard deviation as the real codes.

Predict firstGuess before you look: you draw 1,000 random codes with the same mean and std as the real digits’ codes. How many land far from every real digit’s code?
🔒 Answer the question above to unlock

The code space: every point is a picture

Drag on the map (or focus it and use the arrow keys) to pick a 2-number code. The trained decoder draws it, here in your browser.

Loading the trained decoders…

the decoder’s picture

code (0.00, 0.00)

Loading…
codes of 400 real test digits 3each digit’s average code 1 and 2 std of this model’s codes, where you would draw random ones your code Highlighted dots: the digit nearest your code.

Real codes form clusters, one per digit, with gaps between them. Press Draw 200: about a third of the random codes land far from every real code (the crosses). The lab keeps only 400 real codes, so its gaps look bigger than they are; demo.py compares with 4,000 real codes and finds 11%. Nothing in the loss asks the codes to fill the map or to stay near any center: this autoencoder spread them with a std of 5.2 and 3.9, wherever rebuilding was easiest. A code in a gap decodes to something the decoder never had to draw, which is why made-up digits from a plain autoencoder often look strange.

Try it

Drag the cross into a gap between two clusters, for example between the 1s and the 0s. What does the decoder draw on the way across? Then press Random code ten times and count how often it lands in an empty region.

4. VAE: codes that fill the space

🔒 Answer the question above to unlock

A variational autoencoder (VAE) changes two things.

First, the encoder gives a range, not a point. For each code number it outputs a mean μ and a standard deviation σ. During training the code is drawn at random from that range, using a standard normal number ε (mean 0, std 1, from level 10):

z=μ+σ⋅εz = \mu + \sigma \cdot \varepsilon
Number A VAE’s encoder gives μ = 1 and σ = 0.5 for one code number. The random draw is ε = −2. What is the code z = μ + σ·ε?
🔒 Answer the question above to unlock

So the decoder learns to draw a sensible picture for a whole neighborhood of codes, not just one exact point.

Second, a penalty pulls every picture’s range toward the standard normal, N(0, 1). It is called the KL term. For one code number,

KL=12(μ2+σ2−1−ln⁡σ2)\text{KL} = \tfrac{1}{2}\left(\mu^2 + \sigma^2 - 1 - \ln \sigma^2\right)

It is 0 exactly when μ = 0 and σ = 1, and grows as μ moves away from 0 or σ moves away from 1. (ln is the natural log; ln 1 = 0.)

Number Use KL = ½(μ² + σ² − 1 − ln σ²). For one code number, μ = 2 and σ = 1 (so ln σ² = ln 1 = 0). What is the KL term?

The total loss for one picture is the rebuild loss summed over its 784 pixels, plus the KL term summed over its code numbers. The rebuild part wants codes far apart and precise, so two different pictures never get the same code. The KL part wants every code close to 0 with a std of 1. The compromise packs the codes tightly around 0, overlapping, with no big gaps.

The balance between the two parts depends on that sum. Summed over 784 pixels, a typical rebuild loss is in the hundreds, while the KL term for 2 code numbers is a few units.

ChooseSomeone writes the VAE loss as the rebuild error AVERAGED over the 784 pixels (instead of summed), plus the KL term. What happens in training?
ChooseIn demo.py, the plain autoencoder’s random codes landed far from every real code 11% of the time (about a third in the lab, which keeps fewer codes). A VAE draws random codes from N(0, 1). How often do they land far from every real code?
🔒 Answer the question above to unlock

The code space: every point is a picture

Drag on the map (or focus it and use the arrow keys) to pick a 2-number code. The trained decoder draws it, here in your browser.

Loading the trained decoders…

the decoder’s picture

code (0.00, 0.00)

Loading…
codes of 400 real test digits 3each digit’s average code 1 and 2 std of this model’s codes, where you would draw random ones your code Highlighted dots: the digit nearest your code.

The cost is that pictures are a little less sharp: the VAE’s test loss is 0.046 per pixel against the plain autoencoder’s 0.044. In return, a random code from N(0, 1) almost always decodes to something digit-like. Press Draw 200 again: about 1 in 7 random codes land far from every real code, against about a third for the plain autoencoder. (demo.py, with 4,000 real codes, finds 1% against 11%.)

Go deeper Why write the sample as μ + σ·ε?

Training needs the gradient of the loss with respect to μ and σ (level 6). “Draw z at random from a normal with mean μ and std σ” has no slope, because of the random draw. Writing z=μ+σεz = \mu + \sigma \varepsilon moves the randomness into ε, which doesn’t depend on anything. Then ∂z/∂μ=1\partial z/\partial \mu = 1 and ∂z/∂σ=ε\partial z/\partial \sigma = \varepsilon, and backpropagation works as usual. This rewrite is called the reparameterization trick. In code the encoder usually outputs ln⁡σ2\ln \sigma^2 instead of σ, so that any number it outputs gives a valid σ (always positive).

Now write the sampling step. * and + work number by number on arrays of the same shape.

CodeWrite the VAE’s sampling step: given arrays mu, sigma and a standard normal draw eps (all the same shape), return the codes.

Enter keeps the indent · Tab indents · Esc then Tab leaves the editor · ⌘/Ctrl + Enter runs

🔒 Answer the question above to unlock

Last, the whole VAE loss for one picture. np.log is the natural log, ln. .sum() adds up every number of an array.

CodeWrite the VAE loss for one picture: the binary cross-entropy of every pixel, summed, plus the KL term of every code number, summed.

Enter keeps the indent · Tab indents · Esc then Tab leaves the editor · ⌘/Ctrl + Enter runs

5. Where this goes next

Image generators don’t run diffusion (level D1) on millions of pixels. They first train an autoencoder, then add and remove noise on the small codes, and decode at the end. In level D3 a 784-pixel digit becomes a 16-number code: 49 times fewer numbers to denoise.

Two details connect this level to D3. First, real image models use a VAE-style encoder with a small KL term, so their codes are smooth and about size 1. Diffusion (D1) assumes the signal has a variance of about 1, so the codes are also scaled to std 1 before noise is added. Second, their codes are not one flat list: a big picture becomes a small grid of codes (for example 64 × 64 positions with 4 numbers each), so D3’s idea of cutting into patches still applies to the codes.

I got stuck here Is this the same encoder and decoder as in the Transformer?

Same words, different jobs. Here the encoder compresses a picture into a code and the decoder rebuilds the picture from it. In the 2017 Transformer (side trip N7), the encoder reads a source sentence into one vector per word, and the decoder writes a new sentence while looking at them. Neither of those squeezes anything into a bottleneck.

Recap

a summary for when you finish the level

The key formulas and common mistakes appear here once you clear the level.

You can now

  • Follow the shapes through an autoencoder: (B, 784) → (B, 2) codes → (B, 784) pixels.
  • Compute the rebuild loss and a VAE’s KL term by hand.
  • Write a VAE’s sampling step and its full loss in NumPy.

Keep in mind

  • Mean squared error: np.mean((x - x_hat) ** 2)
  • Sampling: , with ε from N(0, 1)
  • , 0 exactly when μ = 0 and σ = 1
  • VAE loss: rebuild loss summed over the pixels + KL summed over the code numbers

Common mistakes

  • Averaging the rebuild loss over the pixels: the KL term then wins, and every code decodes to the same unclear digit.
  • Expecting random codes of a plain autoencoder to decode well: they often land in gaps between the clusters.
‹ N5. Seq2seq + attentionN7. 2017 Transformer ›

Press ? for keyboard shortcuts

Reading mode · every part open, no stars