Variational Autoencoder in Computer Vision - ML Experiment: Train & Evaluate

Experiment - Variational Autoencoder

Problem:We want to learn a compact representation of images using a Variational Autoencoder (VAE). The current model reconstructs training images well but performs poorly on new images, showing signs of overfitting.

Current Metrics:Training loss: 0.15, Validation loss: 0.45, Training accuracy (reconstruction quality): 92%, Validation accuracy: 70%

Issue:The model overfits the training data, indicated by much better training loss and accuracy compared to validation. This means it does not generalize well to new images.

Your Task

Reduce overfitting so that validation loss decreases and validation accuracy improves to at least 85%, while keeping training accuracy below 90% to avoid overfitting.

You can modify the model architecture (e.g., add dropout or batch normalization).

You can adjust training hyperparameters like learning rate and batch size.

You cannot change the dataset or add more data.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

Computer Vision

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

# Load MNIST dataset
(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

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

# Encoder
encoder_inputs = layers.Input(shape=(28, 28, 1))
x = layers.Conv2D(32, 3, activation='relu', strides=2, padding='same')(encoder_inputs)
x = layers.BatchNormalization()(x)
x = layers.Dropout(0.3)(x)
x = layers.Conv2D(64, 3, activation='relu', strides=2, padding='same')(x)
x = layers.BatchNormalization()(x)
x = layers.Dropout(0.3)(x)
x = layers.Flatten()(x)
x = layers.Dense(16, activation='relu')(x)
z_mean = layers.Dense(latent_dim, name='z_mean')(x)
z_log_var = layers.Dense(latent_dim, name='z_log_var')(x)
z = Sampling()([z_mean, z_log_var])
encoder = Model(encoder_inputs, [z_mean, z_log_var, z], name='encoder')

# Decoder
latent_inputs = layers.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, padding='same', activation='relu')(x)
x = layers.BatchNormalization()(x)
x = layers.Dropout(0.3)(x)
x = layers.Conv2DTranspose(32, 3, strides=2, padding='same', activation='relu')(x)
x = layers.BatchNormalization()(x)
x = layers.Dropout(0.3)(x)
decoder_outputs = layers.Conv2DTranspose(1, 3, activation='sigmoid', padding='same')(x)
decoder = Model(latent_inputs, decoder_outputs, name='decoder')

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

    def call(self, inputs):
        z_mean, z_log_var, z = self.encoder(inputs)
        reconstructed = self.decoder(z)
        kl_loss = -0.5 * tf.reduce_mean(
            1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var), axis=1)
        self.add_loss(tf.reduce_mean(kl_loss))
        return reconstructed

vae = VAE(encoder, decoder)

optimizer = tf.keras.optimizers.Adam(learning_rate=0.0005)
vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())

# Early stopping callback
early_stop = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)

history = vae.fit(x_train, x_train,
                  epochs=50,
                  batch_size=64,
                  validation_data=(x_test, x_test),
                  callbacks=[early_stop])

Added dropout layers with rate 0.3 in encoder and decoder to reduce overfitting.

Added batch normalization layers to stabilize and speed up training.

Reduced learning rate from default 0.001 to 0.0005 for smoother convergence.

Added early stopping to stop training when validation loss stops improving.

Fixed KL loss calculation to reduce over batch dimension correctly.

Results Interpretation

Before: Training loss 0.15, Validation loss 0.45, Training accuracy 92%, Validation accuracy 70%
After: Training loss 0.22, Validation loss 0.28, Training accuracy 88%, Validation accuracy 86%

Adding dropout and batch normalization helped reduce overfitting by making the model less dependent on specific training examples. Lowering the learning rate and using early stopping helped the model converge better and generalize to new data.

Bonus Experiment

Try using a larger latent dimension (e.g., 10) and see how it affects reconstruction quality and overfitting.

💡 Hint

Increasing latent dimension can improve reconstruction but may increase overfitting. Use the same regularization techniques to control it.

Practice

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

easy

A. To compress data into a smaller, meaningful representation and generate new data samples

B. To classify images into different categories

C. To detect edges in images

D. To increase the resolution of images

Solution

1. Step 1: Understand VAE's role in data compression

   VAEs encode input data into a smaller latent space that captures important features.

2. Step 2: Recognize VAE's generative ability

   VAEs can generate new data by sampling from the latent space, unlike simple compression methods.

3. Final Answer:

   To compress data into a smaller, meaningful representation and generate new data samples -> Option A

4. Quick Check:

   VAE purpose = compress + generate [OK]

Hint: Remember: VAE compresses and creates new data [OK]

Common Mistakes:

• Confusing VAE with classifiers
• Thinking VAE only compresses without generation
• Mixing VAE with image enhancement methods

2. Which part of the Variational Autoencoder architecture introduces randomness during training?

easy

A. The encoder output layer

B. The reparameterization trick

C. The decoder input layer

D. The loss function calculation

Solution

1. Step 1: Identify where randomness is needed

   Randomness is introduced to sample from the latent space during training.

2. Step 2: Recognize the reparameterization trick's role

   This trick allows gradients to flow through random sampling by expressing sampling as a deterministic function plus noise.

3. Final Answer:

   The reparameterization trick -> Option B

4. Quick Check:

   Randomness in VAE = reparameterization trick [OK]

Hint: Randomness comes from reparameterization trick [OK]

Common Mistakes:

• Thinking encoder output is random by itself
• Confusing decoder input with randomness source
• Assuming loss function adds randomness

3. Given the following PyTorch snippet for sampling latent variables in a VAE, what is the shape of z if mu and log_var have shape (batch_size, latent_dim)?

std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
z = mu + eps * std

medium

A. (latent_dim, batch_size)

B. (1, latent_dim)

C. (batch_size, 1)

D. (batch_size, latent_dim)

Solution

1. Step 1: Analyze shapes of mu and log_var

   Both have shape (batch_size, latent_dim), so std and eps will match this shape.

2. Step 2: Calculate z shape from element-wise operations

   Since z = mu + eps * std is element-wise, z keeps the same shape as mu and std.

3. Final Answer:

   (batch_size, latent_dim) -> Option D

4. Quick Check:

   z shape = mu shape = (batch_size, latent_dim) [OK]

Hint: Element-wise ops keep shape same as inputs [OK]

Common Mistakes:

• Swapping batch_size and latent_dim dimensions
• Assuming z is a vector, not a batch
• Confusing shape of eps with mu

4. You implemented a VAE but notice the generated images are very blurry. Which debugging step is most likely to improve image sharpness?

medium

A. Use a smaller batch size during training

B. Remove the KL divergence term from the loss function

C. Increase the latent space dimension to capture more details

D. Replace the decoder with a simple linear layer

Solution

1. Step 1: Understand cause of blurry images

   Blurriness often means the latent space is too small to capture enough detail.

2. Step 2: Choose the fix that increases detail capacity

   Increasing latent dimension allows encoding more features, improving sharpness.

3. Final Answer:

   Increase the latent space dimension to capture more details -> Option C

4. Quick Check:

   More latent dims = better detail [OK]

Hint: Bigger latent space means clearer images [OK]

Common Mistakes:

• Removing KL divergence breaks training stability
• Changing batch size rarely fixes blurriness
• Simplifying decoder reduces model power

5. You want to generate new images similar to your training set using a trained VAE. Which step correctly describes how to generate these images?

hard

A. Sample random points from a standard normal distribution, then pass them through the decoder

B. Feed random noise directly into the encoder and use its output as images

C. Use the encoder to compress training images, then output those compressed vectors as images

D. Train the VAE again with random noise as input to generate images

Solution

1. Step 1: Understand VAE latent space distribution

   The latent space is modeled as a standard normal distribution after training.

2. Step 2: Generate images by sampling and decoding

   Sampling random points from this distribution and decoding them produces new images similar to training data.

3. Final Answer:

   Sample random points from a standard normal distribution, then pass them through the decoder -> Option A

4. Quick Check:

   Generate = sample latent + decode [OK]

Hint: Sample latent space, decode to generate images [OK]

Common Mistakes:

• Feeding noise into encoder instead of decoder
• Using encoder outputs as images directly
• Retraining VAE unnecessarily for generation