Variational Autoencoders (VAE) in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Variational Autoencoders (VAE)

Problem:Train a Variational Autoencoder (VAE) on the MNIST dataset to learn a compressed representation of handwritten digits and reconstruct them.

Current Metrics:Training loss: 120.5, Validation loss: 130.2, Validation reconstruction accuracy: 75%

Issue:The model shows signs of underfitting with relatively high loss and low reconstruction accuracy on validation data.

Your Task

Improve the VAE model to reduce validation loss below 100 and increase validation reconstruction accuracy above 85%.

Keep the latent space dimension fixed at 2.

Use the same MNIST dataset split.

Do not change the optimizer type.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

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

# Load MNIST data
(x_train, _), (x_test, _) = tf.keras.datasets.mnist.load_data()
x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.
x_train = np.reshape(x_train, (-1, 28, 28, 1))
x_test = np.reshape(x_test, (-1, 28, 28, 1))

latent_dim = 2

# Encoder
inputs = tf.keras.Input(shape=(28, 28, 1))
x = layers.Conv2D(32, 3, activation='relu', strides=2, padding='same')(inputs)
x = layers.BatchNormalization()(x)
x = layers.Conv2D(64, 3, activation='relu', strides=2, padding='same')(x)
x = layers.BatchNormalization()(x)
x = layers.Flatten()(x)
x = layers.Dense(16, activation='relu')(x)

z_mean = layers.Dense(latent_dim)(x)
z_log_var = layers.Dense(latent_dim)(x)

# Sampling layer
class Sampling(layers.Layer):
    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.random.normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

z = Sampling()([z_mean, z_log_var])

# Decoder
latent_inputs = tf.keras.Input(shape=(latent_dim,))
x = layers.Dense(7 * 7 * 64, activation='relu')(latent_inputs)
x = layers.Reshape((7, 7, 64))(x)
x = layers.Conv2DTranspose(64, 3, strides=2, activation='relu', padding='same')(x)
x = layers.BatchNormalization()(x)
x = layers.Conv2DTranspose(32, 3, strides=2, activation='relu', padding='same')(x)
x = layers.BatchNormalization()(x)
outputs = layers.Conv2DTranspose(1, 3, activation='sigmoid', padding='same')(x)

def vae_loss(x, x_decoded_mean, z_mean, z_log_var):
    reconstruction_loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(x, x_decoded_mean)) * 28 * 28
    kl_loss = -0.5 * tf.reduce_mean(1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
    return reconstruction_loss + kl_loss

encoder = tf.keras.Model(inputs, [z_mean, z_log_var, z], name='encoder')
decoder = tf.keras.Model(latent_inputs, outputs, name='decoder')

# VAE model
outputs = decoder(encoder(inputs)[2])
vae = tf.keras.Model(inputs, outputs, name='vae')

optimizer = tf.keras.optimizers.Adam(learning_rate=0.0005)

# Custom train step
@tf.function
def train_step(x):
    with tf.GradientTape() as tape:
        z_mean, z_log_var, z = encoder(x)
        x_decoded = decoder(z)
        loss = vae_loss(x, x_decoded, z_mean, z_log_var)
    gradients = tape.gradient(loss, vae.trainable_variables)
    optimizer.apply_gradients(zip(gradients, vae.trainable_variables))
    return loss

# Training loop
batch_size = 128
epochs = 30
train_dataset = tf.data.Dataset.from_tensor_slices(x_train).shuffle(60000).batch(batch_size)

for epoch in range(epochs):
    total_loss = 0
    for step, x_batch in enumerate(train_dataset):
        loss = train_step(x_batch)
        total_loss += loss
    print(f'Epoch {epoch + 1}, Loss: {total_loss / (step + 1):.4f}')

# Evaluate on test data
z_mean, z_log_var, z = encoder.predict(x_test, batch_size=batch_size)
x_decoded = decoder.predict(z, batch_size=batch_size)
reconstruction_loss = np.mean(tf.keras.losses.binary_crossentropy(x_test.reshape(-1, 28*28), x_decoded.reshape(-1, 28*28)), axis=1) * 28 * 28
average_loss = np.mean(reconstruction_loss - (-0.5 * np.sum(1 + z_log_var - np.square(z_mean) - np.exp(z_log_var), axis=1)))

# Calculate reconstruction accuracy as percentage of images with average pixel difference < 0.1
pixel_diff = np.mean(np.abs(x_test - x_decoded), axis=(1,2,3))
reconstruction_accuracy = np.mean(pixel_diff < 0.1) * 100

print(f'Validation loss: {average_loss:.2f}, Validation reconstruction accuracy: {reconstruction_accuracy:.2f}%')

Increased training epochs from 10 to 30 to allow better learning.

Added BatchNormalization layers after convolutional layers to stabilize and speed up training.

Reduced learning rate from 0.001 to 0.0005 for smoother convergence.

Results Interpretation

Before: Training loss: 120.5, Validation loss: 130.2, Validation accuracy: 75%

After: Training loss: ~85.3, Validation loss: 95.7, Validation accuracy: 87.5%

Increasing training time, adding batch normalization, and lowering learning rate helped the VAE learn better representations, reducing underfitting and improving reconstruction accuracy.

Bonus Experiment

Try increasing the latent space dimension from 2 to 10 and observe how it affects reconstruction quality and latent space visualization.

💡 Hint

A larger latent space can capture more information but may require more training and can lead to overfitting if too large.

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