Variational autoencoders (VAE): new images from a continuous latent space

A generative model with a latent variable, a VAE written from scratch with the reparameterization trick, an ordered latent space, generating digits from noise compared with an ordinary autoencoder, the effect of β and collapse, interpolation, and where VAEs are used.

12 minutes Matěj Žalmánek 2026-10-16

Series AI basics Part 44 of 52

Full series syllabus 52 of 83 topics published

What you will take away

Why an ordinary autoencoder is not good at generating.
How a VAE works: a distribution instead of a point, the KL penalty and the reparameterization trick.
What β costs, why VAE images are blurry, and where VAEs are used today.

Useful to know first: Information theory: entropy, cross-entropy and KL divergence — where loss functions come from

The autoencoder from the part on representations can squeeze an image into a few numbers and put it back together from them. An obvious question: if the decoder can turn a code into a digit, couldn't we give it a code we made up and get a new digit that nobody ever drew? With an ordinary autoencoder this hardly works. The variational autoencoder is a modification that makes exactly this possible — and it is the first of the generative models we will look at in this section of the series.

It builds on the parts Autoencoders and Information theory (KL divergence).

1. From compression to generation

Definition — generative model

Let 𝒳 be the set of possible data (here 8 × 8 images with pixel brightnesses in [0, 1], i.e. 𝒳 = [0, 1]⁶⁴), n ≥ 1, and let x₁, …, x_n ∈ 𝒳 be training data, which we assume to be independent samples from an unknown distribution p_data on 𝒳. Let Θ be a set of admissible parameters and, for every θ ∈ Θ, let p_θ be a distribution on 𝒳. A generative model is a family of distributions (p_θ)_{θ ∈ Θ} together with a procedure for sampling from the distribution p_θ (producing new x ~ p_θ). Learning chooses θ so that p_θ is as close as possible to p_data, typically as the maximum likelihood estimate

θ̂ = argmax_{θ ∈ Θ} ∑_i=1ⁿ log p_θ(x_i).

A latent variable model additionally introduces a latent space 𝒵 = ℝ^k (k ≥ 1), a fixed density p(z) on it (the prior), and for every code z ∈ 𝒵 a distribution p_θ(x | z) on 𝒳 (the decoder). The distribution of the data is then

p_θ(x) = ∫_𝒵 p_θ(x | z) p(z) dz,

and sampling happens in two steps:

draw z ~ p(z);
draw x ~ p_θ(x | z) — in practice (including here) the output of the decoder, i.e. the pixel probabilities, is often taken directly.

The new x thus look like the data but are not copies of it. From here on we drop the index θ and write p(x | z), p(x).

In plain words

A classifier answers “what is this”. A generative model can answer “show me something that could be in the data”. The recipe: roll a few numbers with dice, hand them to a painter (the decoder), and the painter paints a picture from them.

p(x) = ∫ p(x | z) · p(z) dz

p(z): distribution of the code (the prior) — we choose it ourselves and keep it simple, here N(0, I), because we will need to draw from it when generating
p(x | z): the decoder — how probable the image x is given the code z; here a network that returns, for every pixel, the probability that it is lit
∫ … dz: a sum over all possible codes — the image could have come from any code; but we cannot compute an integral over the whole latent space, which is why the VAE replaces it with an estimate (below)

Generating is easy (draw z, decode), learning needs p(x) — and that is the hard integral.

The catch with an ordinary autoencoder: we do not know where to draw the code from. The encoder places the digits in the latent space wherever it pleases. In our network with a two-dimensional code, the codes of the test digits lie around the point (-3.9, -0.9) with standard deviations 4.3 and 6.0, and between the groups there are holes where the decoder has never seen anything.

2. The variational autoencoder

Definition — variational autoencoder (VAE)

Let 𝒳 = [0, 1]⁶⁴ be the set of images, k ≥ 1 the dimension of the code, I the k × k identity matrix and β > 0. All logarithms in this part are natural (log = ln). A variational autoencoder (VAE; Diederik Kingma and Max Welling, 2013) consists of three parts:

the encoder — a neural network that assigns to an image x ∈ 𝒳 vectors μ(x) ∈ ℝ^k and log σ²(x) ∈ ℝ^k (logarithm taken componentwise); these determine the distribution of the code q(z | x) = N(μ(x), diag σ²(x)), i.e. the coordinates z₁, …, z_k are independent and z_j ~ N(μ_j(x), σ_j²(x)) (the normal distribution);
the decoder — a neural network that assigns to a code z ∈ ℝ^k a distribution p(x | z) on 𝒳 (here via a vector of pixel probabilities y(z) ∈ (0, 1)⁶⁴);
the prior p(z) = N(0, I), the standard normal distribution on ℝ^k.

The parameters of both networks are learned by minimizing the average over the training images of the loss

L(x) = −E_{z ~ q(z | x)}[log p(x | z)] + β · D_KL(q(z | x) ‖ N(0, I)),

where for densities q, p on ℝ^k, D_KL(q ‖ p) = ∫ q(z) log q(z)/p(z) dz is the continuous analogue of the KL divergence from the part on information (with the natural logarithm, in nats). The first term of the loss is the reconstruction error (here the binary cross-entropy of the pixels, see below), the second is a penalty for how much the code distribution differs from the prior. For q = N(μ, diag σ²) it comes out as

D_KL(q ‖ N(0, I)) = ½ ∑_j=1^k (μ_j² + σ_j² − log σ_j² − 1).

A VAE thus encodes an image into a distribution, not a point, and penalizes it for deviating from the distribution that will later be used for generating. With β = 1, −L(x) is a lower bound on log p(x) (the ELBO, theorem below).

In plain words

The encoder may not say “this seven lives exactly here”, only “somewhere in this little cloud”. The decoder then has to be able to draw a seven from any point of the cloud. And the KL penalty pushes all the clouds towards the origin and does not let them get too small — so they overlap, fill the space, and no holes are left between them.

L(x) = −E_{z ~ q(z | x)}[ log p(x | z) ] + β · D_KL( q(z | x) ‖ N(0, I) )

−E_{z ~ q}[log p(x | z)]: reconstruction error — how badly the decoder draws x from a code drawn from the cloud. The expectation over the whole cloud (not just its centre) forces the decoder to draw correctly in the whole neighbourhood of the code; the minus is there because we want the log-probability as large as possible and the loss as small as possible
β: weight of the penalty — sets the trade-off between accuracy and order in the space; β = 1 is exactly the ELBO, other values are tried in the section on β
D_KL(q(z | x) ‖ N(0, I)): distance of the cloud from the distribution we will draw from — without it the encoder would shrink the clouds to points and scatter them far apart (reconstruction would be better), and we would be back to an ordinary autoencoder

The loss for one image; training uses the average over a batch. Both terms are measured in nats (natural logarithm).

D_KL(q(z | x) ‖ N(0, I)) = ½ ∑_j=1^k ( μ_j² + σ_j² − log σ_j² − 1 )

D_KL: the penalty from the previous formula, computed exactly for a normal distribution; the code coordinates are independent, hence the sum over j
μ_j²: shift of the cloud’s centre from the origin — it grows quadratically, so distant codes cost a lot and the encoder pulls them towards zero, where we will be drawing
σ_j² − log σ_j²: width of the cloud — the function s − log s has its minimum 1 exactly at s = 1; it penalizes a cloud that is too wide, but above all one that is too narrow (−log σ² → ∞ for σ → 0), so that the cloud does not turn into a point

The “− 1” only shifts the minimum to zero: D_KL = 0 exactly when μ = 0 and σ = 1, i.e. when q is itself N(0, I). The derivatives the program adds to the gradient: ∂D_KL/∂μ_j = μ_j and ∂D_KL/∂(log σ_j²) = 1/2(σ_j² − 1).

Theorem — the ELBO is a lower bound on the log-likelihood

Let p(z) be a density on ℝ^k (the prior), for every z ∈ ℝ^k let p(x | z) be a distribution on 𝒳 (the decoder), let x ∈ 𝒳 be an image with p(x) = ∫ p(x | z) p(z) dz > 0, and let

p(z | x) = p(x | z) p(z)/p(x)

be the posterior density of the code (Bayes’ theorem). Let q(z | x) be any density on ℝ^k that is positive only where p(z | x) is positive, and for which the expectations E_q = E_{z ~ q(z | x)} below exist and are finite. Then

log p(x) = E_q[log p(x | z)] − D_KL(q(z | x) ‖ p(z)) + D_KL(q(z | x) ‖ p(z | x)),

and therefore

log p(x) ≥ E_q[log p(x | z)] − D_KL(q(z | x) ‖ p(z)),

with equality if and only if q(· | x) = p(· | x).

Proof. log p(x) does not depend on z, so it equals its own expectation over q. By Bayes’ theorem, for every z with q(z | x) > 0

p(x) = p(x | z) p(z)/p(z | x) = p(x | z) · p(z)/q(z | x) · q(z | x)/p(z | x),

where we multiplied the fraction by q(z | x) in the numerator and the denominator. We take the logarithm and the expectation over z ~ q(z | x):

log p(x) = E_q[log p(x | z)] + E_q[log p(z)/q(z | x)] + E_q[log q(z | x)/p(z | x)].

By definition, the second term is −D_KL(q(z | x) ‖ p(z)), the third is D_KL(q(z | x) ‖ p(z | x)). The KL divergence is always non-negative and equal to zero exactly for identical distributions (Gibbs’ inequality from the part on information theory, which holds for densities as well), so dropping the third term gives both the inequality and the condition for equality. ∎

The right-hand side of the inequality is the ELBO (evidence lower bound), and with β = 1 it is exactly −L(x). By minimizing L we therefore raise a lower bound on log p(x) without having to compute the integral. The bound is the tighter, the closer q(z | x) is to the true posterior distribution p(z | x).

Definition — binary cross-entropy of pixels

Let x ∈ [0, 1]⁶⁴ be an image (pixel brightnesses x_j, j = 1, …, 64) and y ∈ (0, 1)⁶⁴ the output of the decoder, where we understand y_j as the probability that pixel j is lit. The binary cross-entropy of the image x with respect to the prediction y is

BCE(x, y) = −∑_j=1⁶⁴ [ x_j log y_j + (1 − x_j) log(1 − y_j) ].

For a black-and-white image x ∈ {0, 1}⁶⁴ and a decoder under which the pixels are lit independently, each with probability y_j(z), exactly BCE(x, y(z)) = −log p(x | z). For grey brightnesses x_j ∈ (0, 1) the same expression is used as the reconstruction error.

An example for one pixel: it is lit (x = 1) and the decoder says 0.9 → −ln 0.9 ≈ 0.105; it says 0.5 → ln 2 ≈ 0.693; it is not lit (x = 0) and the decoder says 0.9 → −ln 0.1 ≈ 2.303. A confident mistake costs the most. Because the output is a sigmoid, y_j = 1/1 + e^−a_j, the derivative with respect to its input a_j is simply ∂BCE/∂a_j = y_j − x_j — just as in logistic regression.

Definition — reparameterization trick

Let μ(x) ∈ ℝ^k and σ(x) ∈ (0, ∞)^k be the outputs of the encoder for an image x. Instead of drawing z ~ q(z | x) = N(μ(x), diag σ²(x)) directly, the reparameterization trick proceeds as follows:

draw ε = (ε₁, …, ε_k) with independent ε_j ~ N(0, 1) — noise whose distribution does not depend on the network’s parameters;
for j = 1, …, k set z_j = μ_j(x) + σ_j(x) · ε_j.

The z obtained this way has exactly the distribution q(z | x) (a linear transformation of a normal variable is normal), and at the same time it is a differentiable function of the parameters: ∂z_j/∂μ_j = 1 and ∂z_j/∂σ_j = ε_j. The gradient therefore flows through backpropagation all the way into the encoder, which is not possible with plain sampling z ~ N(μ, σ²).

In plain words

We “factor the randomness out”: instead of “draw a point from the cloud” we say “draw generic noise and shift and stretch it as the encoder says”. Shifting and stretching are ordinary arithmetic, and that can be differentiated.

z_j = μ_j(x) + σ_j(x) · ε_j, σ_j = exp(½ · log σ_j²), ε_j ~ N(0, 1)

μ_j(x): centre of the cloud, an output of the encoder — ∂z_j/∂μ_j = 1, the decoder’s gradient passes into the encoder unchanged
σ_j(x): width of the cloud — the encoder does not return σ but log σ², which may be any real number; the exponential always turns it into a positive width, so the network does not have to respect any constraint. ∂z_j/∂σ_j = ε_j
ε_j ~ N(0, 1): all of the randomness — it does not depend on any parameter, so nothing has to be differentiated through it; when generating, we skip the encoder and draw z directly from N(0, I)

A shifted and stretched normal variable is again normal: z_j ~ N(μ_j, σ_j²), i.e. exactly a sample from q(z | x).

Example — one test digit (0) in the trained VAE with k = 2

encoder: μ = (0.285, 1.587), log σ² = (−2.695, −1.931) → σ = (0.260, 0.381)

j = 1: ½ (0.081 + 0.068 + 2.695 − 1) ≈ 0.922

j = 2: ½ (2.518 + 0.145 + 1.931 − 1) ≈ 1.797 → D_KL ≈ 2.72 nats

reparameterization with ε = (0.5, −1.0): z = (0.285 + 0.260 · 0.5, 1.587 − 0.381 · 1.0) = (0.415, 1.206)

(colours as in the formulas for D_KL and reparameterization above)

Most of the penalty here comes not from the distance to the origin but from the narrowness of the cloud: the term −log σ² contributes 2.695 and 1.931. The encoder would like the code to be as precise as possible; D_KL keeps it wider. The average D_KL over the test digits is 3.42 nats. The reconstruction error of this digit (from the centre μ) is 34.6 nats over 64 pixels.

def krok_vae(davka, beta=1.0):
    for x in davka:
        mu, logvar = koder(x)                  # k + k numbers
        sigma = exp(0.5 * logvar)
        eps = normal(0, 1, size=len(mu))       # randomness from outside
        z = mu + sigma * eps                   # reparameterization trick
        y = dekoder(z)                         # pixel probabilities
        rek = -sum(x * log(y) + (1 - x) * log(1 - y))
        kl = 0.5 * sum(mu**2 + exp(logvar) - logvar - 1)
        zpetne_sireni(rek + beta * kl)         # gradient flows z → mu, sigma → encoder
    adam_krok()                                # average gradient of the batch

Our network: 64 pixels → 32 neurons (tanh) → μ and log σ² (k numbers each) → 32 neurons → 64 pixels (sigmoid), trained with Adam for 150 epochs on 600 synthetic digits. The same network, only without the noise ε and without the KL penalty, is an ordinary autoencoder — so the comparison is fair.

3. An ordered latent space

Codes of 200 test digits (k = 2); dashed circles of radius 1 and 2 around the origin

Both networks group identical digits together. The autoencoder, however, scatters them over a large area — both circles of the normal distribution are small compared with it. The VAE codes lie around the origin with standard deviations 1.14 and 1.04, that is, exactly where we will be drawing from.

What the whole space looks like is shown by a grid: regularly spaced points of the latent space, each passed through the decoder.

Decoder on a grid of 10 × 10 points of the latent space

In the VAE the digits flow smoothly into one another and almost every cell is a readable digit. To be fair: in the plane the autoencoder’s decoder is also quite smooth, and its grid also contains mostly digits. The difference only shows in full with a code of more dimensions.

4. Generation: new digits from noise

We draw 200 codes and have them decoded: for the VAE from N(0, I), for the autoencoder (to give it a fair chance) from a normal distribution fitted to its actual codes. We measure quality by the distance (mean squared error) of a sample to the nearest training digit. Real test digits are on average 0.017 from the nearest training digit, random noise 0.159.

Definition — distance to the nearest training digit and the fitted distribution

Let x⁽¹⁾, …, x⁽⁶⁰⁰⁾ ∈ [0, 1]⁶⁴ be the training digits. The distance of an image y ∈ [0, 1]⁶⁴ to the nearest training digit is

d(y) = min_{i ∈ {1, …, 600}} 1/64 ∑_j=1⁶⁴ (y_j − x⁽ⁱ⁾_j)²,

i.e. the mean squared error (MSE) to the most similar digit. For drawn images y⁽¹⁾, …, y⁽²⁰⁰⁾, the quality of the generator is the average

d̄ = 1/200 ∑_s=1²⁰⁰ d(y^(s)).

A small d̄ means “the samples look like some real digits” — beware, a generator that merely copies the training data would also score zero.

Let c⁽¹⁾, …, c^(N) ∈ ℝ^k be the codes that the autoencoder’s encoder assigns to N = 200 test digits. The fitted distribution is N(m, diag s²), where for j = 1, …, k

m_j = 1/N ∑_i=1^N c⁽ⁱ⁾_j, s_j = √1/N ∑_i=1^N (c⁽ⁱ⁾_j − m_j)².

The autoencoder thus draws z_j = m_j + s_j · ε_j with independent ε_j ~ N(0, 1), i.e. each coordinate from a normal distribution with the same mean and standard deviation as its actual codes have (for k = 2: m = (−3.9, −0.9), s = (4.3, 6.0)). The coordinates are drawn independently — the fitted distribution does not capture any dependence between them.

Distance of drawn images to the nearest training digit (smaller = more plausible) autoencoder VAE real digits random noise

code size `k`	AE reconstruction (MSE)	VAE reconstruction (MSE)	AE samples	VAE samples
2	0.051	0.063	0.042	0.044
8	0.025	0.036	0.090	0.054
16	0.017	0.039	0.112	0.055

With a two-dimensional code both networks do about equally well (0.042 versus 0.044). With eight dimensions the autoencoder falls apart: its samples have a distance of 0.090, more than halfway to noise, while the VAE holds at 0.054. With sixteen it is 0.112 versus 0.055. In many dimensions the autoencoder’s codes lie on a thin, twisted surface, and a random point from the fitted normal distribution practically never lands on it.

The first ten drawn images, code with 8 dimensions

The price: the VAE reconstructs worse. With k = 16 its reconstruction error is 0.039 versus 0.017 for the autoencoder, and its images are blurrier. The code is noisy, so the decoder learns to draw the “average” of what could be in the neighbourhood of the code. Blurriness is a well-known weakness of VAEs.

5. Reconstruction versus order: β

The weight β determines how strongly the KL penalty pushes the codes towards the normal distribution. Five VAEs with a code of k = 2 and different β:

β	KL (nats per image)	reconstruction (MSE)	samples (distance)
0.1	8.18	0.053	0.044
0.5	4.71	0.057	0.044
1	3.42	0.063	0.044
2	2.07	0.076	0.052
5	0.05	0.131	0.075

A small β leaves the codes looser (KL 8.2) and reconstruction is more accurate. A large β squeezes them towards zero. At β = 5 the KL drops to 0.05: the encoder has stopped storing anything in the code, and the decoder keeps drawing roughly the same thing, whatever code it gets. This is called posterior collapse.

Definition — the nat, and KL as the price of the code

The nat is the unit of information corresponding to the natural logarithm: an event with probability p ∈ (0, 1] carries −ln p nats (instead of −log₂ p bits). Since ln p = ln 2 · log₂ p,

1 nat = 1/ln 2 bits ≈ 1.443 bits.

Let x be a random image with the data distribution p_data (here a uniformly chosen test digit), let the code z given x have the encoder’s distribution q(z | x), and let p(z) = N(0, I) be the prior. Denote by

q(z) = E_{x ~ p_data}[q(z | x)]

the distribution of the codes of all images together, and by I(x; z) the mutual information of the image and the code (the continuous analogue of the definition from the part on information). Then

E_{x ~ p_data}[D_KL(q(z | x) ‖ p(z))] = I(x; z) + D_KL(q(z) ‖ p(z)) ≥ I(x; z),

so the average KL divergence over the data is an upper bound on the information the code carries about the image. It is the “price” the encoder pays for whatever it stores in the code about the image.

At β = 0.1 this is at most 8.18 nats ≈ 11.8 bits per image, at β = 5 only 0.05 nats ≈ 0.07 bits — less than a tenth of a bit, so the code says almost nothing about the image and the decoder has no choice but to keep drawing the same thing.

Images drawn by the VAE (k = 2) at different β

Conversely, β greater than 1 is used deliberately by the β-VAE (Irina Higgins et al., 2017): stronger pressure towards independent code coordinates means that individual dimensions often correspond to understandable properties (rotation, stroke thickness).

6. Smooth transitions

We take the codes of two test digits (networks with an eight-dimensional code, for each digit the image the VAE reconstructs best) and walk along the straight line between them in eight steps: z(t) = (1 − t) · z_a + t · z_b for t = 0, 1/7, 2/7, …, 1, where z_a, z_b are the centres μ of the two digits. We decode every intermediate point.

Interpolation in the latent space (k = 8)

The transition is smooth for both networks: the intermediate steps are plausible shapes between the two digits. Along a straight line between two real codes we do not get far from the data, so here the ordinary autoencoder does not lag behind — its weakness is only drawing blindly. The same principle in large models makes it possible to smoothly change, say, a facial expression or the style of an image.

7. Where VAEs are used

Compression for other generators. Stable Diffusion (latent diffusion, Robin Rombach et al., 2022) does not work with pixels but with the code of an autoencoder with a small KL penalty: it compresses a 512 × 512 × 3 image to 64 × 64 × 4, and diffusion runs in that smaller space. Diffusion will be the subject of one of the next parts.
Discrete codes. VQ-VAE (Aäron van den Oord et al., 2017) picks entries from a learned codebook instead of a continuous code. An image thus becomes a sequence of “tokens” that a transformer can generate.
Scientific data. Molecules, single-cell data, sound — wherever a smooth space is useful in which one can search and interpolate.

In image sharpness VAEs lose to GANs and to diffusion models. But they have stable training, a genuine probabilistic interpretation, and an encoder for free.

8. Summary

A generative model learns the distribution of the data so that it can be sampled; with a latent variable it draws a code and decodes it.
A VAE encodes into a distribution N(μ, σ²), and loss = reconstruction + β · KL to N(0, I). The trick z = μ + σε makes it possible to differentiate through the sampling.
VAE codes lie around the origin and fill the space; autoencoder codes are scattered and full of holes.
With an 8-dimensional code: autoencoder samples 0.090, VAE 0.054 (real digits 0.017, noise 0.159).
The price is worse reconstruction and blurriness; too large a β leads to collapse (KL 0.05).

Want to get hands-on? Have an AI agent train a VAE with a two-dimensional code on the MNIST dataset and plot a grid of decoded images — you will see a whole “map” of handwritten digits.

Do you see it differently? Write to support@coden-app.com.

Variational autoencoders (VAE): new images from a continuous latent space

1. From compression to generation

2. The variational autoencoder

3. An ordered latent space

4. Generation: new digits from noise

5. Reconstruction versus order: β

6. Smooth transitions

7. Where VAEs are used

8. Summary

An Agent on your computer, driven from your phone