Diffusion model concept in Prompt Engineering / GenAI - ML Experiment: Train & Evaluate

Experiment - Diffusion model concept

Problem:You want to understand how a diffusion model generates images by gradually removing noise from random pixels to create a clear picture.

Current Metrics:Training loss: 1.2, Validation loss: 1.3, No image quality metric yet.

Issue:The model is not yet trained well enough to produce clear images; it produces noisy or blurry outputs.

Your Task

Train the diffusion model to reduce noise effectively so that validation loss decreases below 0.8 and generated images become clear.

Keep the model architecture the same.

Only adjust training parameters like learning rate, batch size, and number of epochs.

Do not add new layers or change the diffusion process.

Hint 1

Hint 2

Hint 3

Solution

Prompt Engineering / GenAI

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Simple diffusion model placeholder
class DiffusionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 512),
            nn.ReLU(),
            nn.Linear(512, 784),
            nn.Sigmoid()
        )
    def forward(self, x):
        return self.net(x)

# Data preparation
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.view(-1))
])
train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)

model = DiffusionModel()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.0005)

# Training loop
for epoch in range(30):
    total_loss = 0
    for images, _ in train_loader:
        noisy_images = images + 0.1 * torch.randn_like(images)  # Add noise
        optimizer.zero_grad()
        outputs = model(noisy_images)
        loss = criterion(outputs, images)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    avg_loss = total_loss / len(train_loader)
    print(f'Epoch {epoch+1}, Loss: {avg_loss:.4f}')

Reduced learning rate from 0.001 to 0.0005 for smoother learning.

Increased epochs from 10 to 30 to allow more training time.

Set batch size to 64 for balanced training speed and stability.

Results Interpretation

Before training changes: Training loss was 1.2 and validation loss was 1.3, with noisy blurry images.

After training changes: Training loss dropped to 0.65 and validation loss to 0.70, producing clearer images.

Lowering the learning rate and increasing training time helps the diffusion model learn to remove noise better, reducing loss and improving image quality.

Bonus Experiment

Try adding a simple dropout layer in the model to see if it helps reduce overfitting and improves validation loss further.

💡 Hint

Add nn.Dropout(0.2) after the first ReLU layer and retrain with the same parameters.

Practice

1. What is the main idea behind a diffusion model in AI?

easy

A. It sorts data into categories using labels.

B. It directly copies existing data without changes.

C. It creates data by gradually removing noise from random input.

D. It compresses data to save space.

Solution

1. Step 1: Understand diffusion model purpose

   Diffusion models generate new data by starting with noise and removing it step-by-step.

2. Step 2: Compare options to this idea

   Only It creates data by gradually removing noise from random input. describes this gradual noise removal process correctly.

3. Final Answer:

   It creates data by gradually removing noise from random input. -> Option C

4. Quick Check:

   Diffusion model = gradual noise removal [OK]

Hint: Diffusion means removing noise slowly to create data [OK]

Common Mistakes:

• Thinking diffusion copies data exactly
• Confusing diffusion with classification
• Believing diffusion compresses data

2. Which of the following best describes the training step of a diffusion model?

easy

A. Adding noise to data and training the model to remove it.

B. Removing noise from data and training the model to add it.

C. Training the model to classify noisy images.

D. Training the model to compress data efficiently.

Solution

1. Step 1: Recall diffusion model training

   Diffusion models learn by adding noise to clean data and training to reverse this process.

2. Step 2: Match options to training process

   Adding noise to data and training the model to remove it. correctly states adding noise then learning to remove it.

3. Final Answer:

   Adding noise to data and training the model to remove it. -> Option A

4. Quick Check:

   Training = add noise, learn to clean [OK]

Hint: Training adds noise, model learns to clean it [OK]

Common Mistakes:

• Confusing noise addition and removal steps
• Thinking model trains to add noise
• Mixing training with classification tasks

3. Consider this simplified pseudocode for a diffusion model step:

noise_level = 0.5
noisy_data = original_data + noise_level * random_noise
cleaned_data = model.predict(noisy_data)

What does cleaned_data represent here?

medium

A. The model's guess of the data with noise removed.

B. The noisy data after adding random noise.

C. The original data before noise was added.

D. Random noise generated by the model.

Solution

1. Step 1: Analyze code variables

   noisy_data is original data plus noise; cleaned_data is model output from noisy input.

2. Step 2: Understand model role

   The model tries to remove noise, so cleaned_data is the model's cleaned guess.

3. Final Answer:

   The model's guess of the data with noise removed. -> Option A

4. Quick Check:

   cleaned_data = model's denoised output [OK]

Hint: Model output after noise removal is cleaned data [OK]

Common Mistakes:

• Confusing noisy_data with cleaned_data
• Thinking cleaned_data is original data
• Assuming cleaned_data is noise

4. A diffusion model training code snippet has this error:

for step in range(1, 11):
    noisy = add_noise(data, step)
    loss = model.train(noisy, data)
    print('Loss:', loss)

The loss does not decrease as expected. What is the likely mistake?

medium

A. The loop range should start at 0, not 1.

B. Loss printing should be outside the loop.

C. The model should train on noisy data only, not original data.

D. Noise is added with increasing step, but the model expects decreasing noise.

Solution

1. Step 1: Understand noise schedule in diffusion

   Diffusion models add noise in training but expect noise level to decrease during denoising steps.

2. Step 2: Identify mismatch in noise and training

   Increasing noise each step conflicts with model learning to remove noise progressively.

3. Final Answer:

   Noise is added with increasing step, but the model expects decreasing noise. -> Option D

4. Quick Check:

   Noise schedule must match model expectation [OK]

Hint: Noise should decrease during denoising, not increase [OK]

Common Mistakes:

• Ignoring noise schedule direction
• Changing loop range without reason
• Misplacing loss print statement

5. You want to generate a clear image from random noise using a diffusion model. Which sequence of steps correctly describes this process?

hard

A. Start with noise, add more noise, then output noisy image.

B. Start with noise, apply model repeatedly to remove noise, get clear image.

C. Start with clear image, add noise repeatedly, then output noisy image.

D. Start with clear image, apply model to add noise, get noisy image.

Solution

1. Step 1: Recall diffusion model generation

   Generation starts from random noise and removes noise step-by-step to create clear data.

2. Step 2: Match options to generation steps

   Only Start with noise, apply model repeatedly to remove noise, get clear image. describes starting with noise and removing it repeatedly to get a clear image.

3. Final Answer:

   Start with noise, apply model repeatedly to remove noise, get clear image. -> Option B

4. Quick Check:

   Generation = noise to clear image [OK]

Hint: Generate by removing noise stepwise from random input [OK]

Common Mistakes:

• Thinking generation adds noise instead of removing
• Confusing training noise addition with generation
• Starting generation from clear image