Variational Autoencoders (VAE) in TensorFlow

Introduction

A Variational Autoencoder (VAE) helps us learn to create new data similar to what we have seen. It learns a simple way to represent complex data, like pictures, so we can make new ones or understand the data better.

When you want to generate new images similar to a set of photos you have.

When you want to compress data into a smaller form that still keeps important details.

When you want to learn the main features of data for tasks like anomaly detection.

When you want to explore creative AI, like making new music or art based on examples.

When you want to understand how data is structured in a simple way.

Syntax

TensorFlow

class VAE(tf.keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def sample(self, eps=None):
        if eps is None:
            eps = tf.random.normal(shape=(100, latent_dim))
        return self.decode(eps, apply_sigmoid=True)

    def encode(self, x):
        mean, logvar = self.encoder(x)
        return mean, logvar

    def reparameterize(self, mean, logvar):
        eps = tf.random.normal(shape=tf.shape(mean))
        return eps * tf.exp(logvar * 0.5) + mean

    def decode(self, z, apply_sigmoid=False):
        logits = self.decoder(z)
        if apply_sigmoid:
            probs = tf.sigmoid(logits)
            return probs
        return logits

The VAE class combines an encoder and decoder model.

Reparameterization trick helps us train the model by allowing gradients to flow.

Examples

Encode input to get mean and log variance, then sample z, then decode to reconstruct.

TensorFlow

mean, logvar = encoder(input_data)
z = reparameterize(mean, logvar)
reconstructed = decoder(z)

Using the VAE class methods to encode, sample, and decode data.

TensorFlow

vae = VAE(encoder, decoder)
mean, logvar = vae.encode(x)
z = vae.reparameterize(mean, logvar)
x_recon = vae.decode(z, apply_sigmoid=True)

Sample Model

This program builds a simple VAE for MNIST digits. It trains on one batch and prints the loss.

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers
import numpy as np

latent_dim = 2

# Encoder model
encoder_inputs = tf.keras.Input(shape=(28, 28, 1))
x = layers.Flatten()(encoder_inputs)
x = layers.Dense(64, activation='relu')(x)
mean = layers.Dense(latent_dim)(x)
logvar = layers.Dense(latent_dim)(x)
encoder = tf.keras.Model(encoder_inputs, [mean, logvar], name='encoder')

# Decoder model
decoder_inputs = tf.keras.Input(shape=(latent_dim,))
x = layers.Dense(64, activation='relu')(decoder_inputs)
x = layers.Dense(28 * 28, activation=None)(x)
x = layers.Reshape((28, 28, 1))(x)
decoder = tf.keras.Model(decoder_inputs, x, name='decoder')

class VAE(tf.keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def sample(self, eps=None):
        if eps is None:
            eps = tf.random.normal(shape=(100, latent_dim))
        return self.decode(eps, apply_sigmoid=True)

    def encode(self, x):
        mean, logvar = self.encoder(x)
        return mean, logvar

    def reparameterize(self, mean, logvar):
        eps = tf.random.normal(shape=tf.shape(mean))
        return eps * tf.exp(logvar * 0.5) + mean

    def decode(self, z, apply_sigmoid=False):
        logits = self.decoder(z)
        if apply_sigmoid:
            probs = tf.sigmoid(logits)
            return probs
        return logits

# Loss function
@tf.function
def compute_loss(model, x):
    mean, logvar = model.encode(x)
    z = model.reparameterize(mean, logvar)
    x_logit = model.decode(z)

    cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)
    logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])
    kl_div = -0.5 * tf.reduce_sum(1 + logvar - tf.square(mean) - tf.exp(logvar), axis=1)
    return -tf.reduce_mean(logpx_z - kl_div)

# Training step
@tf.function
def train_step(model, x, optimizer):
    with tf.GradientTape() as tape:
        loss = compute_loss(model, x)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# Load data
(x_train, _), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train.astype('float32') / 255.
x_train = np.expand_dims(x_train, -1)
train_dataset = tf.data.Dataset.from_tensor_slices(x_train).shuffle(60000).batch(128)

# Create VAE
vae = VAE(encoder, decoder)
optimizer = tf.keras.optimizers.Adam(1e-3)

# Train for 1 epoch
for batch in train_dataset.take(1):
    loss = train_step(vae, batch, optimizer)
print(f'Loss after 1 batch: {loss.numpy():.4f}')

Important Notes

The loss combines how well the model reconstructs the input and how close the latent space is to a normal distribution.

Training VAEs usually takes many epochs; here we show only one batch for simplicity.

Summary

VAEs learn to compress and generate data by learning a smooth latent space.

They use an encoder to get mean and variance, then sample using reparameterization.

The decoder reconstructs data from the sampled points.

Practice

1. What is the main purpose of the encoder in a Variational Autoencoder (VAE)?

easy

A. To learn the mean and variance of the latent space distribution

B. To reconstruct the input data from latent variables

C. To generate new data samples directly

D. To calculate the loss function during training

Solution

1. Step 1: Understand the encoder role in VAE

   The encoder processes input data and outputs parameters of a probability distribution, typically mean and variance, representing the latent space.

2. Step 2: Differentiate encoder from decoder

   The decoder uses sampled latent variables to reconstruct data, not the encoder.

3. Final Answer:

   To learn the mean and variance of the latent space distribution -> Option A

4. Quick Check:

   Encoder outputs mean and variance [OK]

Hint: Encoder outputs mean and variance, decoder reconstructs [OK]

Common Mistakes:

• Confusing encoder with decoder
• Thinking encoder generates final output
• Assuming encoder calculates loss

2. Which of the following is the correct way to sample latent variables in a VAE using the reparameterization trick in TensorFlow?

easy

A. z = tf.random.normal(shape) + mean * variance

B. z = mean + tf.random.normal(shape) * variance

C. z = mean * tf.random.normal(shape) + variance

D. z = mean + tf.random.normal(shape) * tf.sqrt(variance)

Solution

1. Step 1: Recall reparameterization formula

   The latent variable z is sampled as z = mean + epsilon * stddev, where epsilon ~ N(0,1) and stddev = sqrt(variance).

2. Step 2: Match formula to options

   z = mean + tf.random.normal(shape) * tf.sqrt(variance) correctly uses tf.sqrt(variance) to get stddev and multiplies by random normal noise.

3. Final Answer:

   z = mean + tf.random.normal(shape) * tf.sqrt(variance) -> Option D

4. Quick Check:

   Sample = mean + noise * stddev [OK]

Hint: Multiply noise by sqrt(variance), not variance directly [OK]

Common Mistakes:

• Using variance instead of its square root
• Multiplying mean by noise instead of adding
• Mixing order of operations

3. Given the following TensorFlow code snippet for a VAE encoder output, what is the shape of the sampled latent vector z?

mean = tf.constant([[0.0, 0.0], [1.0, 1.0]])
log_var = tf.constant([[0.0, 0.0], [0.0, 0.0]])
stddev = tf.exp(0.5 * log_var)
epsilon = tf.random.normal(shape=tf.shape(mean))
z = mean + epsilon * stddev

medium

A. (2, 2)

B. (2,)

C. (1, 2)

D. (2, 1)

Solution

1. Step 1: Check shapes of mean and log_var

   Both mean and log_var have shape (2, 2), meaning 2 samples with 2 latent dimensions each.

2. Step 2: Determine shape of epsilon and z

   epsilon is sampled with shape matching mean, so (2, 2). z is computed element-wise, so it also has shape (2, 2).

3. Final Answer:

   (2, 2) -> Option A

4. Quick Check:

   Output shape matches input batch and latent dims [OK]

Hint: Sample shape matches mean shape exactly [OK]

Common Mistakes:

• Assuming latent vector is 1D
• Ignoring batch dimension
• Confusing log variance shape

4. Identify the error in this VAE decoder code snippet and select the correct fix:

def decode(z):
    x_recon = tf.keras.layers.Dense(784, activation='sigmoid')(z)
    return x_recon

# Usage
z = tf.random.normal(shape=(32, 2))
x_reconstructed = decode(z)

medium

A. The Dense layer should be created inside the function, not called directly on z

B. The Dense layer should be defined once outside the function and called inside

C. Activation should be 'relu' instead of 'sigmoid' for reconstruction

D. The input z shape is incorrect; it should be (784, 2)

Solution

1. Step 1: Understand Keras layer usage

   Keras layers should be defined once (e.g., in __init__) and then called multiple times. Defining inside function creates new layers each call, which is inefficient and breaks training.

2. Step 2: Identify correct fix

   The Dense layer should be defined once outside the function and called inside suggests defining the Dense layer once outside and calling it inside decode, which is the proper pattern.

3. Final Answer:

   The Dense layer should be defined once outside the function and called inside -> Option B

4. Quick Check:

   Define layers once, call repeatedly [OK]

Hint: Define layers once, call inside functions [OK]

Common Mistakes:

• Defining layers inside functions repeatedly
• Confusing activation functions for reconstruction
• Misunderstanding input shapes

5. You want to train a VAE on image data with 28x28 pixels. Which loss function combination correctly balances reconstruction and latent space regularization?

hard

A. Mean squared error reconstruction loss only

B. KL divergence only without reconstruction loss

C. Binary cross-entropy reconstruction loss + KL divergence between latent and standard normal

D. Categorical cross-entropy reconstruction loss + L2 regularization on weights

Solution

1. Step 1: Understand VAE loss components

   VAE loss combines reconstruction loss (measuring how well output matches input) and KL divergence (regularizing latent space to be close to standard normal).

2. Step 2: Match loss to image data and VAE theory

   For 28x28 images with pixel values in [0,1], binary cross-entropy is suitable for reconstruction. KL divergence is essential for latent regularization.

3. Final Answer:

   Binary cross-entropy reconstruction loss + KL divergence between latent and standard normal -> Option C

4. Quick Check:

   Reconstruction + KL divergence = correct VAE loss [OK]

Hint: Use reconstruction loss + KL divergence for VAE [OK]

Common Mistakes:

• Ignoring KL divergence term
• Using only reconstruction loss
• Using wrong loss for image data