Variational Autoencoders (VAE) in TensorFlow - Practice Problems & Coding Challenges

Challenge - 5 Problems

🧠 Conceptual

intermediate

What is the main purpose of the KL divergence term in a VAE loss?

In Variational Autoencoders, the loss function includes a reconstruction loss and a KL divergence term. What does the KL divergence term do?

A It measures the distance between the encoder and decoder outputs to improve training speed.

B It calculates the difference between the input and output images to improve reconstruction quality.

C It measures how much the learned latent distribution differs from the prior distribution, encouraging the latent space to follow the prior.

D It regularizes the weights of the neural network to prevent overfitting.

💡 Hint

Think about how the latent space is shaped in a VAE.

The KL divergence term forces the learned latent distribution to be close to a prior (usually a standard normal), which helps the model generate meaningful samples and ensures smooth latent space.

❓ Predict Output

intermediate

What is the shape of the latent vector z in this VAE encoder output?

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

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers

class Encoder(tf.keras.Model):
    def __init__(self, latent_dim):
        super().__init__()
        self.dense1 = layers.Dense(64, activation='relu')
        self.dense_mean = layers.Dense(latent_dim)
        self.dense_log_var = layers.Dense(latent_dim)

    def call(self, x):
        x = self.dense1(x)
        mean = self.dense_mean(x)
        log_var = self.dense_log_var(x)
        epsilon = tf.random.normal(shape=tf.shape(mean))
        z = mean + tf.exp(0.5 * log_var) * epsilon
        return z

encoder = Encoder(latent_dim=16)
input_tensor = tf.random.normal(shape=(32, 100))
z = encoder(input_tensor)
print(z.shape)

A (100, 16)

B (16, 32)

C (32, 100)

D (32, 16)

💡 Hint

The batch size is 32 and latent_dim is 16.

The input batch has 32 samples, and the latent dimension is set to 16, so the output z shape is (32, 16).

❓ Model Choice

advanced

Which model architecture is best suited for image data in a VAE?

You want to build a VAE to generate 28x28 grayscale images. Which encoder architecture is most appropriate?

A A convolutional neural network with Conv2D layers followed by flattening and dense layers.

B A fully connected network with input flattened to 784 units followed by dense layers.

C A recurrent neural network processing the image as a sequence of pixels.

D A simple linear regression model with one dense layer.

💡 Hint

Think about spatial structure in images.

Convolutional layers capture spatial features in images better than fully connected or recurrent layers, making them ideal for image VAEs.

❓ Hyperparameter

advanced

What effect does increasing the latent dimension size have on a VAE?

If you increase the latent dimension size in a VAE, what is the most likely effect?

A The KL divergence term becomes zero, making the model deterministic.

B The model can capture more complex features but risks overfitting and slower training.

C The model will always train faster and produce better reconstructions.

D The reconstruction loss will increase because the latent space is too large.

💡 Hint

Think about model capacity and overfitting.

A larger latent space allows encoding more information but can cause overfitting and requires more training time.

❓ Metrics

expert

How to correctly compute total VAE loss in TensorFlow?

Given reconstruction loss computed as mean squared error and KL divergence computed per batch, which code snippet correctly computes the total loss for training a VAE?

TensorFlow

import tensorflow as tf

reconstruction_loss = tf.reduce_mean(tf.square(inputs - outputs))
kl_loss = -0.5 * tf.reduce_sum(1 + log_var - tf.square(mean) - tf.exp(log_var), axis=1)

# Which of the following computes total loss correctly?

A total_loss = reconstruction_loss + tf.reduce_mean(kl_loss)

B total_loss = tf.reduce_mean(reconstruction_loss + kl_loss)

C total_loss = tf.reduce_mean(reconstruction_loss) + tf.reduce_sum(kl_loss)

D total_loss = tf.reduce_sum(reconstruction_loss) + tf.reduce_sum(kl_loss)

💡 Hint

Remember to average reconstruction loss over batch and sum KL divergence per sample before averaging.

Reconstruction loss is averaged over batch, KL divergence is summed per sample then averaged over batch. Option A correctly combines these.

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