Variational Autoencoders (VAE) in TensorFlow - Model Metrics & Evaluation

Metrics & Evaluation - Variational Autoencoders (VAE)

Which metric matters for Variational Autoencoders (VAE) and WHY

For VAEs, the key metrics are the reconstruction loss and the KL divergence. Reconstruction loss measures how well the model can recreate the input data from its compressed form. KL divergence measures how close the learned data distribution is to the desired normal distribution. Together, they ensure the model learns useful features and generates realistic new data.

Confusion matrix or equivalent visualization

VAEs are unsupervised and do not use confusion matrices. Instead, we track losses during training:

Epoch 1/10
Reconstruction Loss: 0.1200
KL Divergence: 0.0800
Total Loss: 0.2000

Epoch 2/10
Reconstruction Loss: 0.1100
KL Divergence: 0.0750
Total Loss: 0.1850

... (loss values decrease as training progresses)
    

This shows the model improving at reconstructing inputs and learning a smooth latent space.

Precision vs Recall tradeoff (or equivalent) with concrete examples

VAEs balance reconstruction loss and KL divergence. If reconstruction loss is too low but KL divergence is high, the model memorizes data but latent space is not smooth, hurting generation quality. If KL divergence is too low but reconstruction loss is high, the model generates smooth latent space but poor reconstructions.

Example: Imagine compressing photos. If you compress too much (high KL), photos become blurry (high reconstruction loss). If you compress too little (low KL), photos are sharp but the compression is not useful for generating new images.

What "good" vs "bad" metric values look like for VAEs

Good: Reconstruction loss and KL divergence both steadily decrease and stabilize at low values. This means the model reconstructs inputs well and learns a smooth latent space.

Bad: Reconstruction loss stays high or KL divergence collapses to zero. This means the model either fails to reconstruct inputs or ignores the latent space, leading to poor generation.

Common pitfalls in VAE metrics

• KL vanishing: KL divergence becomes zero early, causing the model to ignore latent variables.
• Overfitting: Very low reconstruction loss on training but high on validation means poor generalization.
• Ignoring latent space: If KL divergence is too low, the latent space is not meaningful for generation.
• Misinterpreting losses: Total loss is a sum; focusing on only one loss can mislead about model quality.

Self-check question

Your VAE model shows a reconstruction loss of 0.05 but KL divergence is near zero. Is this good? Why or why not?

Answer: This is not good. The model reconstructs inputs well but ignores the latent space (KL near zero). It likely memorizes data and cannot generate diverse new samples.

Key Result

For VAEs, balancing low reconstruction loss and meaningful KL divergence ensures good data reconstruction and smooth latent space for generation.

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