Skip to content
Python Central Hub

Variational Autoencoders (VAE)

What you’ll learn

  • why a plain autoencoder’s latent space is unsafe to sample from, and what makes a variational autoencoder probabilistic and generative instead
  • how the encoder outputs a mean and a log-variance instead of a single coding, and why the reparameterization trick lets you still backpropagate through a random sampling step
  • how the loss combines a reconstruction loss with a latent (KL-divergence) loss that pulls every coding toward a standard Gaussian cloud
  • how to generate brand-new images by sampling random codings and decoding them, and how to blend two images together with semantic interpolation

From a single coding to a cloud of possibilities

A regular undercomplete autoencoder learns to map each input to one fixed point in latent space. Nothing forces those points to sit close together, or to fill the space without gaps — so if you picked a random point and decoded it, you’d probably get garbage. That’s fine for reconstruction, but useless for generating new data.

A variational autoencoder (VAE) fixes this by making the encoder output a small Gaussian distribution for each input instead of one exact point: a mean μ and a log-variance log(σ²). The actual coding z is then sampled randomly from that distribution. Because every input gets smeared into a little cloud, and the loss actively pushes all those clouds toward a shared standard Gaussian, the latent space ends up smooth and densely packed — any point you sample from a standard Gaussian is now likely to decode into something plausible.

diagram Variational Autoencoder: sample before you decode mermaid
The encoder outputs a mean and log-variance instead of one coding; z is sampled from that distribution using the reparameterization trick, then decoded.

The reparameterization trick

You can’t backpropagate through a random sampling operation directly — gradients don’t flow through random_normal()random_normal(). The fix is the reparameterization trick: instead of sampling zz directly from N(μ, σ)N(μ, σ), sample noise εε from a fixed N(0, 1)N(0, 1) and compute z = μ + σ * εz = μ + σ * ε. Now the randomness lives entirely in εε (which needs no gradient), while μμ and σσ stay fully differentiable.

Géron’s implementation has the encoder output log(σ²)log(σ²) rather than σσ directly — it’s more numerically stable and trains faster, since σ = exp(log(σ²) / 2)σ = exp(log(σ²) / 2):

A Sampling layer using the reparameterization trick
import numpy as np
import tensorflow as tf
 
codings_size = 10
 
class Sampling(tf.keras.layers.Layer):
    def call(self, inputs):
        mean, log_var = inputs
        return tf.random.normal(tf.shape(log_var)) * tf.exp(log_var / 2) + mean
 
inputs = tf.keras.layers.Input(shape=[28, 28])
z = tf.keras.layers.Flatten()(inputs)
z = tf.keras.layers.Dense(150, activation="selu")(z)
z = tf.keras.layers.Dense(100, activation="selu")(z)
codings_mean = tf.keras.layers.Dense(codings_size)(z)          # mu
codings_log_var = tf.keras.layers.Dense(codings_size)(z)       # log(sigma^2)
codings = Sampling()([codings_mean, codings_log_var])
 
variational_encoder = tf.keras.Model(
    inputs=[inputs], outputs=[codings_mean, codings_log_var, codings]
)
 
X_batch = np.random.rand(4, 28, 28).astype("float32")
mean, log_var, z_sample = variational_encoder.predict(X_batch, verbose=0)
print("codings_mean shape:", mean.shape)
print("codings_log_var shape:", log_var.shape)
print("sampled codings shape:", z_sample.shape)
A Sampling layer using the reparameterization trick
import numpy as np
import tensorflow as tf
 
codings_size = 10
 
class Sampling(tf.keras.layers.Layer):
    def call(self, inputs):
        mean, log_var = inputs
        return tf.random.normal(tf.shape(log_var)) * tf.exp(log_var / 2) + mean
 
inputs = tf.keras.layers.Input(shape=[28, 28])
z = tf.keras.layers.Flatten()(inputs)
z = tf.keras.layers.Dense(150, activation="selu")(z)
z = tf.keras.layers.Dense(100, activation="selu")(z)
codings_mean = tf.keras.layers.Dense(codings_size)(z)          # mu
codings_log_var = tf.keras.layers.Dense(codings_size)(z)       # log(sigma^2)
codings = Sampling()([codings_mean, codings_log_var])
 
variational_encoder = tf.keras.Model(
    inputs=[inputs], outputs=[codings_mean, codings_log_var, codings]
)
 
X_batch = np.random.rand(4, 28, 28).astype("float32")
mean, log_var, z_sample = variational_encoder.predict(X_batch, verbose=0)
print("codings_mean shape:", mean.shape)
print("codings_log_var shape:", log_var.shape)
print("sampled codings shape:", z_sample.shape)
text
codings_mean shape: (4, 10)
codings_log_var shape: (4, 10)
sampled codings shape: (4, 10)
text
codings_mean shape: (4, 10)
codings_log_var shape: (4, 10)
sampled codings shape: (4, 10)

The DenseDense layers producing codings_meancodings_mean and codings_log_varcodings_log_var share the same input, but the encoder has three outputs — you only ever feed codingscodings (the third one) into the decoder; the other two exist so you can compute the latent loss.

The loss: reconstruction + latent (KL) loss

A VAE’s loss has two parts:

  1. Reconstruction loss — the usual term (binary cross-entropy here) that pushes the decoder’s output to look like the input.
  2. Latent loss — the KL divergence between the encoder’s distribution and a standard Gaussian N(0, 1)N(0, 1). It penalizes codings that drift away from that shared Gaussian, which is exactly what keeps the latent space smooth and samplable.

Using the log(σ²)log(σ²) trick, the latent loss for one instance simplifies to:

Building and training the full variational autoencoder
import tensorflow as tf
 
decoder_inputs = tf.keras.layers.Input(shape=[codings_size])
x = tf.keras.layers.Dense(100, activation="selu")(decoder_inputs)
x = tf.keras.layers.Dense(150, activation="selu")(x)
x = tf.keras.layers.Dense(28 * 28, activation="sigmoid")(x)
outputs = tf.keras.layers.Reshape([28, 28])(x)
variational_decoder = tf.keras.Model(inputs=[decoder_inputs], outputs=[outputs])
 
_, _, codings = variational_encoder(inputs)
reconstructions = variational_decoder(codings)
variational_ae = tf.keras.Model(inputs=[inputs], outputs=[reconstructions])
 
# Equation 17-4 (Geron): the latent loss, using the log(sigma^2) trick
latent_loss = -0.5 * tf.reduce_sum(
    1 + codings_log_var - tf.exp(codings_log_var) - tf.square(codings_mean), axis=-1
)
variational_ae.add_loss(tf.reduce_mean(latent_loss) / 784.0)
variational_ae.compile(loss="binary_crossentropy", optimizer="rmsprop")
 
print("total trainable weights:", len(variational_ae.trainable_weights))
Building and training the full variational autoencoder
import tensorflow as tf
 
decoder_inputs = tf.keras.layers.Input(shape=[codings_size])
x = tf.keras.layers.Dense(100, activation="selu")(decoder_inputs)
x = tf.keras.layers.Dense(150, activation="selu")(x)
x = tf.keras.layers.Dense(28 * 28, activation="sigmoid")(x)
outputs = tf.keras.layers.Reshape([28, 28])(x)
variational_decoder = tf.keras.Model(inputs=[decoder_inputs], outputs=[outputs])
 
_, _, codings = variational_encoder(inputs)
reconstructions = variational_decoder(codings)
variational_ae = tf.keras.Model(inputs=[inputs], outputs=[reconstructions])
 
# Equation 17-4 (Geron): the latent loss, using the log(sigma^2) trick
latent_loss = -0.5 * tf.reduce_sum(
    1 + codings_log_var - tf.exp(codings_log_var) - tf.square(codings_mean), axis=-1
)
variational_ae.add_loss(tf.reduce_mean(latent_loss) / 784.0)
variational_ae.compile(loss="binary_crossentropy", optimizer="rmsprop")
 
print("total trainable weights:", len(variational_ae.trainable_weights))

The latent loss is divided by 784 because Keras’ "binary_crossentropy""binary_crossentropy" computes the mean over all 784 pixels, not the sum — dividing keeps the two loss terms on a comparable scale (you’ll just want a larger learning rate to compensate).

You’d train this exactly like any other autoencoder: variational_ae.fit(X_train, X_train, epochs=50, batch_size=128, ...)variational_ae.fit(X_train, X_train, epochs=50, batch_size=128, ...), using the input as its own target.

Another implementation style: a subclassed Model with train_steptrain_step (Chollet)

Géron’s add_loss()add_loss() approach above works well with the functional API. Chollet’s Deep Learning with Python builds the exact same VAE a different way: subclass keras.Modelkeras.Model and override train_step()train_step() directly. This is the pattern you reach for whenever training departs from plain supervised learning — and a VAE, which trains on a combination of two losses computed from three different tensors (z_meanz_mean, z_log_varz_log_var, and the reconstruction), is a textbook case.

A VAE as a subclassed Model with a custom train_step (Chollet)
import tensorflow as tf
from tensorflow import keras
 
class Sampler(keras.layers.Layer):
    # Chollet's Sampler takes z_mean and z_log_var as two SEPARATE arguments,
    # rather than as one (mean, log_var) tuple like Géron's Sampling layer above.
    def call(self, z_mean, z_log_var):
        batch_size = tf.shape(z_mean)[0]
        z_size = tf.shape(z_mean)[1]
        epsilon = tf.random.normal(shape=(batch_size, z_size))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon
 
class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.sampler = Sampler()
        # Track running averages of each loss so they show up in fit()'s logs
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = keras.metrics.Mean(name="reconstruction_loss")
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")
 
    @property
    def metrics(self):
        return [self.total_loss_tracker, self.reconstruction_loss_tracker, self.kl_loss_tracker]
 
    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var = self.encoder(data)
            z = self.sampler(z_mean, z_log_var)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(keras.losses.binary_crossentropy(data, reconstruction), axis=(1, 2))
            )
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            total_loss = reconstruction_loss + tf.reduce_mean(kl_loss)
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {
            "total_loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
        }
 
# vae = VAE(encoder, decoder)
# vae.compile(optimizer=keras.optimizers.Adam(), run_eagerly=True)
# vae.fit(mnist_digits, epochs=30, batch_size=128)   # no target passed — train_step ignores it
A VAE as a subclassed Model with a custom train_step (Chollet)
import tensorflow as tf
from tensorflow import keras
 
class Sampler(keras.layers.Layer):
    # Chollet's Sampler takes z_mean and z_log_var as two SEPARATE arguments,
    # rather than as one (mean, log_var) tuple like Géron's Sampling layer above.
    def call(self, z_mean, z_log_var):
        batch_size = tf.shape(z_mean)[0]
        z_size = tf.shape(z_mean)[1]
        epsilon = tf.random.normal(shape=(batch_size, z_size))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon
 
class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.sampler = Sampler()
        # Track running averages of each loss so they show up in fit()'s logs
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = keras.metrics.Mean(name="reconstruction_loss")
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")
 
    @property
    def metrics(self):
        return [self.total_loss_tracker, self.reconstruction_loss_tracker, self.kl_loss_tracker]
 
    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var = self.encoder(data)
            z = self.sampler(z_mean, z_log_var)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(keras.losses.binary_crossentropy(data, reconstruction), axis=(1, 2))
            )
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            total_loss = reconstruction_loss + tf.reduce_mean(kl_loss)
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {
            "total_loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
        }
 
# vae = VAE(encoder, decoder)
# vae.compile(optimizer=keras.optimizers.Adam(), run_eagerly=True)
# vae.fit(mnist_digits, epochs=30, batch_size=128)   # no target passed — train_step ignores it

Two things are worth noticing here. First, compile(loss=None)compile(loss=None) and fit()fit() with no target: since train_step()train_step() computes and applies the loss itself, Keras never needs an external loss=loss= argument or a yy to compare against — only mnist_digitsmnist_digits goes in, no target at all. Second, the KL term is written with tf.reduce_sum(..., axis=(1, 2))tf.reduce_sum(..., axis=(1, 2)) reducing over the reconstruction’s spatial dimensions before averaging over the batch — mathematically the same Kullback–Leibler penalty as Géron’s Equation 17-4 above, just wired through train_step()train_step() instead of add_loss()add_loss().

Generating brand-new data

Once trained, generating a new image is almost embarrassingly simple: sample a random coding from N(0, 1)N(0, 1) and decode it — no input image required at all.

Generating new samples by decoding random codings
import tensorflow as tf
 
tf.random.set_seed(42)
 
new_codings = tf.random.normal(shape=[5, codings_size])
generated = variational_decoder(new_codings)
print("sampled codings shape:", new_codings.shape)
print("generated images shape:", generated.shape)
Generating new samples by decoding random codings
import tensorflow as tf
 
tf.random.set_seed(42)
 
new_codings = tf.random.normal(shape=[5, codings_size])
generated = variational_decoder(new_codings)
print("sampled codings shape:", new_codings.shape)
print("generated images shape:", generated.shape)
text
sampled codings shape: (5, 10)
generated images shape: (5, 28, 28)
text
sampled codings shape: (5, 10)
generated images shape: (5, 28, 28)

You can also perform semantic interpolation: encode two real images to get two codings, blend them (e.g., 0.5 * z1 + 0.5 * z20.5 * z1 + 0.5 * z2), and decode the blend. Because the latent space is smooth, the result looks like a genuine in-between image — not a double-exposed overlay you’d get by blending pixels directly.

Concept vectors: editing images by moving through latent space

Because the latent space is smooth, certain directions in it correspond to a single, meaningful visual attribute — a concept vector. Chollet’s book gives the classic example: train a VAE on a dataset of faces, average the codings of many smiling faces, average the codings of many non-smiling faces, and subtract one mean from the other. The result — a smile vector — can be added to any face’s coding, and decoding z + smile_vectorz + smile_vector produces the same face, smiling. It’s the same idea you’ve seen with word embeddings (king - man + woman ≈ queenking - man + woman ≈ queen), just applied to a latent space of images instead of a latent space of words:

Building and applying a concept vector
import tensorflow as tf
 
tf.random.set_seed(0)
 
# Toy stand-ins for the codings of 50 smiling and 50 non-smiling training faces
smiling_codings = tf.random.normal(shape=[50, codings_size]) + 0.6
not_smiling_codings = tf.random.normal(shape=[50, codings_size]) - 0.6
 
# The concept vector is just the difference between the two group means
smile_vector = tf.reduce_mean(smiling_codings, axis=0) - tf.reduce_mean(not_smiling_codings, axis=0)
 
# Nudge any coding along that direction before decoding
z = tf.random.normal(shape=[1, codings_size])
z_smiling = z + smile_vector
 
before = variational_decoder(z)
after = variational_decoder(z_smiling)
print("smile vector shape:", smile_vector.shape)
print("decoded shapes match:", before.shape == after.shape)
Building and applying a concept vector
import tensorflow as tf
 
tf.random.set_seed(0)
 
# Toy stand-ins for the codings of 50 smiling and 50 non-smiling training faces
smiling_codings = tf.random.normal(shape=[50, codings_size]) + 0.6
not_smiling_codings = tf.random.normal(shape=[50, codings_size]) - 0.6
 
# The concept vector is just the difference between the two group means
smile_vector = tf.reduce_mean(smiling_codings, axis=0) - tf.reduce_mean(not_smiling_codings, axis=0)
 
# Nudge any coding along that direction before decoding
z = tf.random.normal(shape=[1, codings_size])
z_smiling = z + smile_vector
 
before = variational_decoder(z)
after = variational_decoder(z_smiling)
print("smile vector shape:", smile_vector.shape)
print("decoded shapes match:", before.shape == after.shape)
text
smile vector shape: (10,)
decoded shapes match: True
text
smile vector shape: (10,)
decoded shapes match: True

Notice this never needed per-pixel labels or any extra supervision beyond a rough split into “has the attribute” and “doesn’t” — the VAE itself trained with no labels at all. The same trick discovers vectors for glasses, hair color, age, or pose; Chollet credits artist Tom White with a well-known demonstration of exactly this smile vector, trained on the CelebA celebrity-faces dataset. One gotcha: if your two groups differ in more than the target attribute (say, the smiling faces also skew younger), the concept vector picks up that correlation too — a good concept vector needs the two groups to be otherwise as similar as possible.

Visualize it

Blue dots are the codings of real training images, pulled by the KL loss into a Gaussian cloud around the origin. Watch a new amber point get sampled from that same cloud every couple of seconds and decoded (right) into a brand-new pattern — nearby points in latent space decode into visibly similar outputs, which is exactly what makes the space safe to sample from:

sketch Sampling the latent space to generate new data p5.js
Blue dots: codings of real training images. Amber dot: a freshly sampled z, decoded on the right into a new pattern.

Mini-checkpoint

Why not just use the encoder’s mean μ directly as the coding, and skip sampling entirely?

  • Because nothing would then force nearby points in latent space to decode to similar outputs. The random sampling step is what makes the loss “notice” a whole neighborhood around each μμ during training — pushing the decoder to make that entire neighborhood decode sensibly, not just the exact point μμ. Skip the sampling, and you’re back to an ordinary (unsafe-to-sample) autoencoder.

🧪 Try It Yourself

Exercise 1 – Build the Sampling Layer

Exercise 2 – Compute the Latent (KL) Loss

Exercise 3 – Generate New Samples by Decoding Random Codings

Next

Continue to Generative Adversarial Networks (GANs) — instead of one network learning a samplable distribution, two networks compete: a generator that fakes data, and a discriminator that tries to catch it.

If this helped you, consider buying me a coffee ☕

Buy me a coffee

Was this page helpful?

Let us know how we did