Variational Lossy Autoencoder
A VAE's latent code only holds what its decoder can't work out alone. Choose what the decoder can see, and you choose what the latent must remember.
Start here (no maths needed)
When a program squashes a picture into a short code, it only stores what it can't easily guess back, so you can choose what goes into the code.
Think of it like this
Imagine describing a drawing to a friend who will redraw it. You only have a few words. If your friend is great at drawing small details, you don't waste words on them. You say the big things: "a big number 7, slanting right". Your friend fills in the rest. Their drawing will have the same big shape, but the little details may differ. That's why it's called "lossy": some details get lost, on purpose.
How it works
- One part of the program, the encoder, squashes a picture into a short code.
- Another part, the decoder, rebuilds the picture from that code.
- Space in the code is costly, so the program only stores what the decoder can't figure out alone.
- The researchers let the decoder see only a few nearby pixels. It can draw small details, but not the big shape.
- So the big shape has to go into the code. The tiny details are left for the decoder to fill in.
Why robots care
A robot's camera sees a lot: carpet patterns, shadows, the cup it needs to grab. If the robot squashes each picture into a short code, it wants that code to hold where the cup is, not the carpet pattern. This idea helps choose what the robot remembers.
Try it
- Pick a simple drawing and describe it to a friend in 10 words or fewer, without letting them see it. Compare their drawing with the original. What stayed the same? What changed?
- This page has no interactive demo. Try the drawing game with two different friends and see whether the big shape stays the same while the small details change.
Words to know
- Autoencoder
- A program that squashes data into a short code and then rebuilds it.
- Code
- The short summary the picture is squashed into. Also called the latent.
- Decoder
- The part that rebuilds the picture from the code.
- Lossy
- Keeps the important parts and lets some details go.
Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.
1. Summary
TL;DR
People kept trying to pair VAEs with powerful autoregressive decoders (RNNs, PixelCNNs), and kept finding the latent code ignored. The usual explanation was "optimisation is hard". This paper argues it's not just optimisation: even at the optimum, a VAE prefers to model anything the decoder can handle on its own, and only puts the rest into the latent. The authors turn that into a tool. Give the decoder only a small local window of pixels, and it can model texture but not global shape, so shape must go into the latent. The result is a deliberately lossy code, plus state-of-the-art likelihoods on several image benchmarks.
- Diagnosis: via bits-back coding, using the latent always costs extra bits when the approximate posterior is imperfect. If the decoder can model the data without
z, ignoringzis optimal. - Lossy code: a PixelCNN decoder with a small receptive field (6 masked conv layers with 3×3 filters) handles local detail. The latent keeps global structure.
- Measured: on binarised MNIST the latent used about 19.2 bits per image, vs 37.3 bits for the same VAE with a factorised decoder.
- Better prior: an autoregressive-flow prior on
zmatches an inverse-autoregressive-flow posterior in training cost, but gives a deeper generative path. - Likelihoods: 79.03 nats on statically binarised MNIST, and 2.95 bits/dim on CIFAR-10, the best then among variationally trained latent-variable models.
2. The core idea
A VAE maximises the evidence lower bound (ELBO):
L(x) = E_{q(z|x)}[ log p(x|z) ] − KL( q(z|x) ‖ p(z) )The first term rewards reconstruction. The second charges for every bit the encoder q(z|x) puts into z beyond the prior p(z). Read as a code length (bits-back coding, the idea from Day 4), the expected cost is:
E[−L(x)] = E[ −log p(x) ] + E[ KL( q(z|x) ‖ p(z|x) ) ]
≥ H(data) + E[ KL( q(z|x) ‖ p(z|x) ) ]H(data): the best possible code length, the entropy of the data.KL(q(z|x) ‖ p(z|x)): the penalty for an approximate posterior that doesn't exactly match the true one. In practice it's never zero.- So routing information through
zalways pays this tax. Ifp(x|z)can model the data with no help fromz, the cheapest solution isq(z|x) = p(z): the latent carries nothing. That's posterior collapse, explained as a preference, not a bug.
The VLAE uses this on purpose. Its decoder only sees a local window:
p_local(x|z) = Π_i p( x_i | z, x_{window(i)} )Whatever is predictable from nearby pixels (texture, stroke edges) gets modelled locally. Long-range structure (the digit's identity and shape) can't be, so it must come through z.
3. Deep dive
"Lossy" is the feature
Decode the same z twice and you get two images with the same global shape but different local detail: stroke width and exact pixel masks change. That's the point. The latent is a summary of what matters globally, and the decoder fills in plausible detail. On CIFAR-10, widening the decoder window (4×2, then 5×3, then 7×4) let the decoder take over more structure, so the latent kept only rough shapes. Colour often went missing from the latent, since it's easy to predict locally. Conditioning the decoder on a grayscale version of the image forced colour back into z. Change what the decoder sees, change what the code stores.
The paper is honest that this isn't magic. On OMNIGLOT, where small patches carry meaning, the lossy code sometimes lost the character's identity. You still have to decide which information your task needs.
Autoregressive flow prior
A plain Gaussian prior is a poor fit for a rich latent. The authors make the prior an autoregressive flow: start from simple noise ε and transform it, z = f(ε). Rewriting the bound shows this is the same, on the encoder side, as using an inverse autoregressive flow posterior, with the same training cost. But on the generation side the model gets an extra deep transformation for free. On MNIST, swapping an IAF posterior for an equivalent AF prior improved test likelihood (79.88 to 79.30 nats without the autoregressive decoder).
Where it sits in the course
- Autoregressive decoders are the "predict the next symbol" models of Day 1, applied to pixels.
- The code-length view links straight to MDL (Day 4, Day 5) and to Day 26: a good representation is a short description of the part of the data you care about.
- Receptive-field control is the same lever as in Day 11, used for a different reason.
4. Why it matters for robotics
- Latent world models. Model-based RL agents often compress camera frames into a latent with a VAE-style encoder, then learn dynamics in that latent. A latent spent on carpet texture is capacity not spent on where the object is. This paper gives a principled way to steer what the latent keeps.
- Posterior collapse is a real failure mode. If you give a latent-variable model a powerful autoregressive decoder (over frames or action sequences), expect it to ignore the latent unless you limit the decoder or use tricks like free bits, which this paper also used.
- Decide the information split by design. For example, a decoder that already sees the previous frame can only gain from
zby learning what changed. That's an application of the paper's principle, not a result from it, but it's a useful way to design robot representations. - Lossy is fine for control. A policy rarely needs pixel-perfect reconstruction. A compact code of task-relevant structure is often easier to learn from and cheaper to run on board.
5. How to study it
- Read sections 2 and 3 of the paper slowly (about 1.5 h). Re-derive the step from the bits-back code length to the KL penalty yourself.
- Look at the paper's figures of original vs "decompressed" images on MNIST and CIFAR-10. Note what changes and what stays.
- In the 30u30 Day 23 folder, run
setup.py, readimplementation.py(masked convolutions, PixelCNN decoder, flow prior), then runtrain_minimal.pywith the flow prior on. - Do the exercises in order: masked convolution, a PixelCNN prior, reproducing posterior collapse by growing the receptive field, the flow prior, then the full VLAE.
- Stretch: log the KL term (bits in the latent) as you widen the decoder's window. It should fall as the decoder takes over.
6. Self-check
Why does a VAE with a powerful decoder ignore its latent, even if optimisation were perfect?
Using the latent costs an extra KL penalty whenever the approximate posterior isn't exact. If the decoder can model the data alone, setting the posterior equal to the prior is cheaper.
How does the VLAE force global structure into z?
Its decoder only sees a small local window of previous pixels. Local detail can be modelled there, but long-range shape can't, so it has to come from z.
What happens to the latent as the decoder's receptive field grows?
The decoder models more structure itself, so the latent keeps less, ending with only rough shape.
Why is an autoregressive-flow prior "free"?
On the encoder path it's mathematically equivalent to an IAF posterior with the same training cost, but it adds a deeper transformation on the generation path.
How did the authors get colour back into the latent on CIFAR-10?
They made the decoder condition only on a grayscale version of the image, so colour could no longer be predicted locally and had to be stored in z.
7. Further reading
- The original paper
- 30u30 Day 23: implementation, notebook, cheatsheet, exercises
- Next: Day 24, GPipe