Diffusion Models

Abstract

A walk through the NVIDIA DLI workshop on generative AI with diffusion models, covering U-Net architecture, DDPM, iterative optimizations, and classifier-free guidance: with a closer look at the math and design choices the course glosses over.

Introduction

Last year I worked through the NVIDIA DLI course Generative AI with Diffusion Models. The course is structured as six hands-on workshop notebooks, building from scratch toward a text-conditioned diffusion model using CLIP.

This post covers the first four modules: U-Net architecture, DDPM forward and reverse processes, iterative optimizations, and classifier-free guidance. My aim is not to reproduce the workshop (NVIDIA did that well enough) but to examine a few things more carefully than the course does, particularly the math behind the forward diffusion shortcut and the reasoning behind each architectural choice.

The code is a cleaned-up version of my workshop implementation. I added type annotations, factored out reproducibility utilities, and rewrote parts I found unclear. The dataset is FashionMNIST, resized to 16×16: small enough to train on a laptop GPU in minutes, large enough to tell whether the model is working.

Setup

The noise schedule uses \(T=150\) timesteps with a linear beta schedule: \(\beta\) increases from \(0.0001\) to \(0.02\), so images progress from nearly unchanged at \(t=0\) to nearly pure Gaussian noise at \(t=T\). Reversing that schedule is the network’s job: the next question is what architecture can learn to undo it.

nrows = 10
ncols = 15

T = nrows * ncols
B_start = 0.0001
B_end = 0.02
B = torch.linspace(B_start, B_end, T).to(device)

def load_dataset(img_size: int, batch_size: int) -> tuple[ConcatDataset, DataLoader]:
    from torchvision.datasets import FashionMNIST

    transform = VT.Compose(
        [
            VT.Resize((img_size, img_size)),
            VT.ToTensor(),  # Scales data into [0,1]
            VT.RandomHorizontalFlip(),
            VT.Lambda(lambda t: (t * 2) - 1),  # Scale between [-1, 1]
        ]
    )

    train_set = FashionMNIST(DATASET_PATH, train=True, download=True, transform=transform)
    test_set = FashionMNIST(DATASET_PATH, train=False, download=True, transform=transform)
    dataset = ConcatDataset([train_set, test_set])
    dataloader = DataLoader(dataset, batch_size=batch_size, drop_last=True, **loader_kwargs)
    return dataset, dataloader


dataset, dataloader = load_dataset(IMG_SIZE, BATCH_SIZE)

Architecture

Diffusion models need a neural network that takes a noisy image and a timestep as input and returns a prediction of the noise added at that step. The U-Net (originally designed for biomedical image segmentation) is the standard backbone for this task.

The NVIDIA workshop builds the architecture in stages, starting from a simpler baseline: Conv2d + BatchNorm2d + ReLU + MaxPool2d building blocks, time embeddings as a plain linear projection of t, and no residual connections at the bottleneck. After training, the model generates images, but with a distinctive grid artifact called the checkerboard problem. The outputs look more like ink blots than clothing items: vague shapes with repeating pixel-level structure that shouldn’t be there.

Five targeted changes eliminate this:

1. GELU instead of ReLU. ReLU hard-zeros all negative activations. When a neuron’s pre-activation is consistently negative, it receives no gradient and permanently stops contributing: the dying ReLU problem. GELU has a smooth, non-zero gradient everywhere. In a diffusion U-Net that must operate across 150 different noise levels, gradient stability through every neuron matters more than in a standard classifier.

2. GroupNorm instead of BatchNorm. BatchNorm computes normalization statistics across the batch dimension. In diffusion training, each image in a batch has a different timestep, so the batch statistics mix images at completely different noise levels. GroupNorm computes statistics within channel groups of a single sample, so normalization becomes independent of what else is in the batch.

3. Rearrange pooling instead of MaxPool. MaxPool discards three of every four values in a 2×2 window, a hard information bottleneck. More critically, the transposed convolution in the upsampling path produces checkerboard artifacts when its stride doesn’t tile the input evenly, a known failure mode of transposed convolutions. Replacing MaxPool with a learned space-to-depth operation sidesteps this. See below.

4. Sinusoidal timestep embeddings instead of linear. A linear projection of t maps nearby timesteps to similar vectors, making it hard for the model to distinguish small differences in noise level. Sinusoidal embeddings (borrowed directly from the positional encoding in Attention Is All You Need) produce a unique, high-frequency fingerprint for every timestep. Each dimension oscillates at a different frequency, so no two positions share the same embedding vector regardless of proximity.

5. Residual connections. The ResidualConvBlock used at the input stage (down0) adds a skip connection from input to output: out = conv2(conv1(x)) + conv1(x). This gives gradients a direct path through the network, shortening how much computation fine-grained spatial details have to pass through.

Applied together, these changes eliminate the checkerboard artifact and the ink-blot texture that came with it. Whether the result is a recognizable garment is a separate matter of training budget and conditioning: as the training runs later in this post make clear, both turn out to be the difference between noise and a recognizable garment.

class GELUConvBlock(nn.Module):
    def __init__(self, in_chs: int, out_chs: int, group_size: int) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(in_chs, out_chs, kernel_size=3, stride=1, padding=1),
            nn.GroupNorm(group_size, out_chs),
            nn.GELU(),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.layers(x)

Learnable Downsampling

The checkerboard problem is partly a transposed convolution problem. When an upsampling layer uses ConvTranspose2d with stride 2, it interleaves zeros between input values and then applies a convolution kernel. If the kernel’s effective footprint doesn’t divide evenly into the output, some output pixels are influenced by more kernel positions than others, producing a periodic grid pattern in the output.

One fix is to replace max pooling in the downsampling path with a different operation that pairs more naturally with transposed convolution. RearrangePoolBlock does this: it rearranges the 2×2 spatial neighborhood into the channel dimension (a space-to-depth operation), then learns how to combine those four values with a convolution. No information is discarded; the model decides what to keep.

import einops

x = torch.rand(128, 3, 128, 128)
y = einops.rearrange(x, "b c (h p1) (w p2) -> b (c p1 p2) h w", p1=2, p2=2)

assert np.prod(x.shape) == np.prod(y.shape)

# Let's demonstrate what it does
size = (2, 2)
repeat = (8, 8)

# Define a base 2x2 pattern
# base_pattern = torch.tensor([[0, 4, 8, 16], [1, 5, 9,] , [2], [3]], dtype=torch.float32)
base_pattern = np.arange(np.prod(size)).reshape(size)

base_pattern = torch.asarray(base_pattern)
base_pattern = base_pattern / base_pattern.numel()


# Repeat the pattern to get a 16x16 checkerboard
checkerboard = base_pattern.repeat(*repeat)
x = checkerboard.reshape(1, 1, *checkerboard.shape)

show_image(x.numpy())

p1, p2 = size
y = einops.rearrange(x, "b c (h p1) (w p2) -> b (c p1 p2) h w", p1=p1, p2=p2)
show_image(y.numpy())

The top image shows a 16×16 grid tiling a 2×2 pattern. After rearrange, the spatial dimensions halve and the channel count grows by 4×, with each channel holding one of the four “slots” from the original 2×2 window. The subsequent GELUConvBlock can learn any linear combination of those four values; MaxPool is just the degenerate case where one slot is always selected. Crucially, this downsampling operation has a natural inverse (space-to-depth and depth-to-space are symmetric). That symmetry is what eliminates the periodicity transposed convolutions introduce.

from einops.layers.torch import Rearrange


class RearrangePoolBlock(nn.Module):
    def __init__(self, in_chs: int, group_size: int) -> None:
        super().__init__()
        self.rearrange = Rearrange("b c (h p1) (w p2) -> b (c p1 p2) h w", p1=2, p2=2)
        self.conv = GELUConvBlock(4 * in_chs, in_chs, group_size)

    def forward(self, x: Tensor) -> Tensor:
        x = self.rearrange(x)
        return self.conv(x)

class DownBlock(nn.Module):
    def __init__(self, in_chs: int, out_chs: int, group_size: int) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            GELUConvBlock(in_chs, out_chs, group_size),
            GELUConvBlock(out_chs, out_chs, group_size),
            RearrangePoolBlock(out_chs, group_size),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.layers(x)

class UpBlock(nn.Module):
    def __init__(self, in_chs: int, out_chs: int, group_size: int) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            nn.ConvTranspose2d(2 * in_chs, out_chs, kernel_size=2, stride=2),
            GELUConvBlock(out_chs, out_chs, group_size),
            GELUConvBlock(out_chs, out_chs, group_size),
            GELUConvBlock(out_chs, out_chs, group_size),
            GELUConvBlock(out_chs, out_chs, group_size),
        )

    def forward(self, x: Tensor, skip: Tensor) -> Tensor:
        x = torch.cat((x, skip), 1)
        x = self.layers(x)
        return x

class ResidualConvBlock(nn.Module):
    def __init__(self, in_chs: int, out_chs: int, group_size: int) -> None:
        super().__init__()
        self.conv1 = GELUConvBlock(in_chs, out_chs, group_size)
        self.conv2 = GELUConvBlock(out_chs, out_chs, group_size)

    def forward(self, x: Tensor) -> Tensor:
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        out = x1 + x2
        return out

import math


class SinusoidalPositionEmbedBlock(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        self.dim = dim

    def forward(self, time: Tensor) -> Tensor:
        device = time.device
        half_dim = self.dim // 2
        embeddings = math.log(10_000) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        embeddings = time[:, None] * embeddings[None, :]
        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
        return embeddings

class EmbedBlock(nn.Module):
    def __init__(self, input_dim: int, emb_dim: int) -> None:
        super().__init__()
        self.input_dim = input_dim
        self.model = nn.Sequential(
            nn.Linear(input_dim, emb_dim),
            nn.GELU(),
            nn.Linear(emb_dim, emb_dim),
            nn.Unflatten(1, (emb_dim, 1, 1)),
        )

    def forward(self, x: Tensor) -> Tensor:
        x = x.view(-1, self.input_dim)
        return self.model(x)

Putting It Together

The UNet class assembles these blocks into the full architecture. The forward pass has a few notable details:

• Time embeddings go into the upsampling path only. The encoder sees the noisy image and extracts features; the decoder uses the timestep to calibrate how much denoising to apply at each resolution.
• Time embeddings are added, not concatenated: up1 = self.up1(up0 + temb_1, down2). Addition keeps the channel count stable through the decoder.
• A final skip connection concatenates down0 (the initial residual convolution output) directly with up2, preserving fine spatial details all the way to the output without passing through the bottleneck.
• group_size_base = 4 is constrained, not arbitrary. GroupNorm requires the group size to divide the channel count evenly. With down_chs = (64, 64, 128), the derived sizes small_group_size = 8 (divides 64 ✓) and big_group_size = 32 (divides both 64 and 128 ✓) are a clean fit. Other multipliers of group_size_base (1, 2, or 8) would divide evenly too, but values like 3, 5, 6, or 7 fail outright for the 128-channel layers.

class UNet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        img_chs = IMG_CH
        down_chs = (64, 64, 128)
        up_chs = down_chs[::-1]  # Reverse of the down channels
        latent_image_size = IMG_SIZE // 4  # 2 ** (len(down_chs) - 1)
        t_dim = 8
        group_size_base = 4
        small_group_size = 2 * group_size_base
        big_group_size = 8 * group_size_base

        # Initial convolution
        self.down0 = ResidualConvBlock(img_chs, down_chs[0], small_group_size)

        # Downsample
        self.down1 = DownBlock(down_chs[0], down_chs[1], big_group_size)
        self.down2 = DownBlock(down_chs[1], down_chs[2], big_group_size)
        self.to_vec = nn.Sequential(nn.Flatten(), nn.GELU())

        # Embeddings
        self.dense_emb = nn.Sequential(
            nn.Linear(down_chs[2] * latent_image_size**2, down_chs[1]),
            nn.ReLU(),
            nn.Linear(down_chs[1], down_chs[1]),
            nn.ReLU(),
            nn.Linear(down_chs[1], down_chs[2] * latent_image_size**2),
            nn.ReLU(),
        )

        # Time embeddings (temb)
        self.sinusoidaltime = SinusoidalPositionEmbedBlock(t_dim)
        self.temb_1 = EmbedBlock(t_dim, up_chs[0])
        self.temb_2 = EmbedBlock(t_dim, up_chs[1])

        # Upsample
        self.up0 = nn.Sequential(
            nn.Unflatten(1, (up_chs[0], latent_image_size, latent_image_size)),
            GELUConvBlock(up_chs[0], up_chs[0], big_group_size),
        )
        self.up1 = UpBlock(up_chs[0], up_chs[1], big_group_size)
        self.up2 = UpBlock(up_chs[1], up_chs[2], big_group_size)

        # Match output channels and one last concatenation
        self.out = nn.Sequential(
            nn.Conv2d(2 * up_chs[-1], up_chs[-1], 3, 1, 1),
            nn.GroupNorm(small_group_size, up_chs[-1]),
            nn.ReLU(),
            nn.Conv2d(up_chs[-1], img_chs, 3, 1, 1),
        )

    def forward(self, x: Tensor, t: Tensor) -> Tensor:
        down0 = self.down0(x)
        down1 = self.down1(down0)
        down2 = self.down2(down1)
        latent_vec = self.to_vec(down2)

        latent_vec = self.dense_emb(latent_vec)
        t = t.float() / T  # Convert from [0, T] to [0, 1]
        t = self.sinusoidaltime(t)

        up0 = self.up0(latent_vec)

        temb_1 = self.temb_1(t)
        up1 = self.up1(up0 + temb_1, down2)

        temb_2 = self.temb_2(t)
        up2 = self.up2(up1 + temb_2, down1)

        return self.out(torch.cat((up2, down0), 1))

model = UNet()
print("Num params: ", sum(p.numel() for p in model.parameters()))
model = torch.compile(model.to(device))

Num params:  1979777

import graphviz
from torchview import draw_graph

graphviz.set_jupyter_format("png")
model_graph = draw_graph(
    UNet(),
    input_size=((BATCH_SIZE, IMG_CH, IMG_SIZE, IMG_SIZE), (BATCH_SIZE,)),
    device="meta",
    depth=1,
    expand_nested=False,
    graph_dir="TD",
    graph_name="U-Net with Positional Embeddings",
)
model_graph.resize_graph(scale=0.75)
model_graph.visual_graph

Diffusion

With the network architecture in place, the next question is what objective actually trains it. DDPM¹ is the generative modeling framework this workshop uses. The core idea is to generate data by learning to reverse a noise process. It has two components: a forward process that gradually destroys an image by adding Gaussian noise, and a reverse process that a trained model uses to reconstruct images from noise.

The model never learns to synthesize an image from scratch in one shot: that is a hard, underspecified problem. Instead it learns a simpler local task: given an image with a specific amount of noise at step \(t\), predict the noise that was added. Repeated \(T\) times from pure Gaussian noise, those small corrections compound into a coherent image.

Crucially, the reverse process adds a small amount of fresh noise at each step (except the last). This stochasticity is intentional: it’s designed to prevent every reverse trajectory from collapsing to the same output, producing diverse samples from the same model.

Note

The figures later in this post use a single seed per condition, so that diversity isn’t shown directly, only the mechanism meant to produce it.

Forward Process

At each step \(t\), a small amount of noise is added according to:

\[q(x_t \mid x_{t-1}) = \mathcal{N}\!\left(x_t;\, \sqrt{1 - \beta_t}\, x_{t-1},\, \beta_t \mathbf{I}\right)\]

This is the same \(T=150\), \(\beta \in [0.0001, 0.02]\) schedule from Setup, applied one step at a time.

The key implementation insight is that you don’t need to apply this step \(t\) times to get \(x_t\) from \(x_0\). Marginalizing out all intermediate steps gives a closed form:

\[q(x_t \mid x_0) = \mathcal{N}\!\left(x_t;\, \sqrt{\bar{\alpha}_t}\, x_0,\, (1 - \bar{\alpha}_t)\mathbf{I}\right)\]

where \(\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s\) and \(\alpha_t = 1 - \beta_t\).

In code: a_bar = torch.cumprod(a, dim=0). During training, a random timestep \(t\) is sampled per batch, and the noisy image at that timestep is computed in a single matrix multiply, with no iteration over intermediate steps required.

Reverse Process

Given \(x_t\) and \(t\), the model predicts the noise \(\epsilon_\theta(x_t, t)\) that was added. \(\mu_t\), the mean of the distribution over \(x_{t-1}\) given \(x_t\), is then:

\[\mu_t = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}}\, \epsilon_\theta(x_t, t) \right)\]

If \(t > 0\), a small amount of fresh noise is added back before proceeding to step \(t-1\). At \(t = 0\), the process terminates: that’s the generated image.

The training objective is MSE between the predicted noise and the noise actually added in the forward step.

def show_tensor_image(image):

    reverse_transforms = VT.Compose(
        [
            VT.Lambda(lambda t: (t + 1) / 2),
            VT.Lambda(lambda t: torch.minimum(torch.tensor([1]), t)),
            VT.Lambda(lambda t: torch.maximum(torch.tensor([0]), t)),
            VT.ToPILImage(),
        ]
    )
    plt.imshow(reverse_transforms(image[0].detach().cpu()))


class DDPM:
    def __init__(self, B: Tensor, device: torch.device) -> None:
        self.B = B
        self.T = len(B)
        self.device = device

        # Forward diffusion variables
        self.a = 1.0 - self.B
        self.a_bar = torch.cumprod(self.a, dim=0)
        self.sqrt_a_bar = torch.sqrt(self.a_bar)  # Mean Coefficient
        self.sqrt_one_minus_a_bar = torch.sqrt(1 - self.a_bar)  # St. Dev. Coefficient

        # Reverse diffusion variables
        self.sqrt_a_inv = torch.sqrt(1 / self.a)
        self.pred_noise_coeff = (1 - self.a) / torch.sqrt(1 - self.a_bar)

    def q(self, x_0: Tensor, t: Tensor) -> tuple[Tensor, Tensor]:
        """
        The forward diffusion process
        Returns the noise applied to an image at timestep t
        x_0: the original image
        t: timestep
        """
        t = t.int()
        noise = torch.randn_like(x_0)
        sqrt_a_bar_t = self.sqrt_a_bar[t, None, None, None]
        sqrt_one_minus_a_bar_t = self.sqrt_one_minus_a_bar[t, None, None, None]

        x_t = sqrt_a_bar_t * x_0 + sqrt_one_minus_a_bar_t * noise
        return x_t, noise

    def get_loss(self, model: nn.Module, x_0: Tensor, t: Tensor, *model_args: Tensor) -> Tensor:
        x_noisy, noise = self.q(x_0, t)
        noise_pred = model(x_noisy, t, *model_args)
        return F.mse_loss(noise, noise_pred)

    @torch.no_grad()
    def reverse_q(self, x_t: Tensor, t: Tensor, e_t: Tensor) -> Tensor:
        """
        The reverse diffusion process
        Returns an image with the noise from time t removed and time t-1 added.
        x_t: the noisy image at time t
        t: timestep
        e_t: the model's predicted noise at time t
        """
        t = t.int()
        pred_noise_coeff_t = self.pred_noise_coeff[t]
        sqrt_a_inv_t = self.sqrt_a_inv[t]
        u_t = sqrt_a_inv_t * (x_t - pred_noise_coeff_t * e_t)
        if t[0] == 0:  # All t values should be the same
            return u_t  # Reverse diffusion complete!
        else:
            B_t = self.B[t - 1]  # Apply noise from the previous timestep
            new_noise = torch.randn_like(x_t)
            return u_t + torch.sqrt(B_t) * new_noise

    @torch.no_grad()
    def sample_images(
        self, model: UNet, img_ch: int, img_size: int, ncols: int, *model_args: Tensor, axis_on: bool = False
    ) -> None:
        # Noise to generate images from
        x_t = torch.randn((1, img_ch, img_size, img_size), device=self.device)
        hidden_rows = self.T // ncols

        # Go from T to 0 removing and adding noise until t = 0
        imgs: list[np.ndarray] = []
        for i in reversed(range(self.T)):
            t = torch.full((1,), i, device=self.device).float()
            e_t = model(x_t, t, *model_args)  # Predicted noise
            x_t = self.reverse_q(x_t, t, e_t)
            if i % hidden_rows == 0:
                imgs.append(x_t.detach().squeeze().cpu().numpy())

        fig, axes = plt.subplots(ncols=len(imgs), figsize=(8, 8 * len(imgs)), frameon=False)
        for img, ax in zip(imgs, axes.flatten(), strict=True):
            ax.imshow(img, cmap="viridis", interpolation="none")
            ax.axis("on" if axis_on else "off")

        fig.tight_layout(pad=0.3)
        plt.show()

ddpm = DDPM(B, device)

The forward process is deterministic given \(x_0\) and the noise: the a_bar shortcut lets us compute \(x_t\) for any \(t\) in one shot. Here is what a single image looks like as \(t\) progresses from 0 (top-left) to \(t=T-1=149\) (bottom-right):

set_random_seed()
idx = 0  # dataset index to visualize; change to try a different item/class
x0, _ = dataset[idx]
x0 = x0.unsqueeze(0).to(device)

x_ts = []
fig, axes = plt.subplots(nrows, ncols, figsize=(ncols, nrows), frameon=False)
for t_idx, ax in enumerate(axes.flat):
    t = torch.tensor([t_idx], device=device).float()
    x_t, _ = ddpm.q(x0, t)
    x_ts.append(x_t.squeeze().cpu().numpy())
    ax.imshow(x_ts[-1], cmap="viridis", interpolation="none", vmin=-1, vmax=1)
    ax.axis("off")

fig.tight_layout(pad=0.1)
plt.show()

Training

The model’s job is to predict the noise \(\epsilon\) that was added to an image, not the clean image itself. This is a deliberate choice: the DDPM paper found noise prediction trains more stably than predicting \(x_0\) directly. The two are mathematically equivalent (given \(x_t\) and \(\epsilon\) you can recover \(x_0\)), but the noise target has better-behaved gradients across timesteps.

At each training step: sample a random \(t\) per image, corrupt the image with the forward process to get \(x_t\), ask the model to predict the noise \(\epsilon_\theta(x_t, t)\), compute MSE against the actual noise, backpropagate. The random \(t\) ensures the model sees all noise levels equally.

Five epochs on FashionMNIST gives a rough sense of whether the architecture and loss are working at all. On a mid-range GPU, the full run takes a few minutes; on CPU, expect 30-60 minutes. It’s a short budget: as the samples below show, it’s enough to confirm the training loop works, but not necessarily enough for the unconditional model to produce a recognizable garment.

model = UNet().to(device)
model.compile()

optimizer = optim.Adam(model.parameters(), lr=0.001)
epochs = 5

model.train()
for epoch in range(epochs):
    for step, batch in enumerate(dataloader):
        optimizer.zero_grad()

        t = torch.randint(0, T, (BATCH_SIZE,), device=device).float()
        x = batch[0].to(device)
        loss = ddpm.get_loss(model, x, t)
        loss.backward()
        optimizer.step()

        if epoch % 2 == 0 and step % 100 == 0:
            print(f"Epoch {epoch} | step {step:03d} Loss: {loss.item()} ")
            ddpm.sample_images(model, IMG_CH, IMG_SIZE, ncols)

Epoch 0 | step 000 Loss: 1.1058292388916016 

Epoch 0 | step 100 Loss: 0.16490435600280762 

Epoch 0 | step 200 Loss: 0.13419075310230255 

Epoch 0 | step 300 Loss: 0.12638157606124878 

Epoch 0 | step 400 Loss: 0.10222411155700684 

Epoch 0 | step 500 Loss: 0.13434891402721405 

Epoch 2 | step 000 Loss: 0.09219031035900116 

Epoch 2 | step 100 Loss: 0.11431898921728134 

Epoch 2 | step 200 Loss: 0.10901810228824615 

Epoch 2 | step 300 Loss: 0.09168179333209991 

Epoch 2 | step 400 Loss: 0.10279276967048645 

Epoch 2 | step 500 Loss: 0.11188428848981857 

Epoch 4 | step 000 Loss: 0.08083761483430862 

Epoch 4 | step 100 Loss: 0.09731955081224442 

Epoch 4 | step 200 Loss: 0.07860389351844788 

Epoch 4 | step 300 Loss: 0.08876349776983261 

Epoch 4 | step 400 Loss: 0.08546773344278336 

Epoch 4 | step 500 Loss: 0.09255847334861755 

In my run, the unconditional model never got past noise. The samples above stay a diffuse cloud of green-and-yellow static through most of the 150 reverse steps, and even in the final row all that emerges is a faint rectangular smudge, nothing a human would call a shirt or a shoe. Five epochs isn’t enough training for this model to learn “generate a plausible garment from scratch.” There are ten very different silhouettes to cover, and nothing in the input tells the model which one to aim for, so it has to hedge across all of them at once.

That contrast becomes obvious once class conditioning enters the picture below: the same architecture, trained for the same five epochs, produces clearly recognizable garments as soon as it’s given a class to aim for.

Classifier-Free Guidance

In principle, an unconditioned model generates a random FashionMNIST image, whichever class the reverse process happens to land on. CFG² adds control: you specify a class (T-shirt, trousers, sneaker…) and the model generates that class specifically.

The elegant part is how this avoids needing a separate classifier. During training, the class label is randomly dropped with probability \(p \approx 0.10\), replaced with a null token. This trains one model to produce both conditional (\(c\) given) and unconditional (\(c\) dropped) outputs from the same weights.

At inference, both outputs are computed and blended, with \(\varnothing\) denoting the dropped/null class:

\[\hat{\epsilon}_t = (1 + w)\,\epsilon_\theta(x_t, t, c) - w\,\epsilon_\theta(x_t, t, \varnothing)\]

The weight \(w\) controls guidance strength: \(w = 0\) is pure conditional sampling, larger \(w\) pushes harder toward the class at the cost of diversity, and negative \(w\) actively avoids the class.

From UNet to ContextUNet

The architectural change from UNet to ContextUNet is minimal; it is not a separate branch. A class embedding is injected into the decoder alongside the existing time embedding, using the same EmbedBlock structure. The only difference is how each embedding is applied:

# time embedding is added:
up1 = self.up1(up0 + temb_1, down2)

# class embedding multiplies, time embedding adds:
up1 = self.up1(cemb_1 * up0 + temb_1, down2)

Multiplication scales the feature map by the class signal; addition shifts it. The combination lets the class context modulate what is generated while the timestep controls how much denoising to apply.

The 10% unconditional case does not remove the class input; it passes a zero vector instead. get_context() zeroes out the one-hot embedding with a Bernoulli mask:

c_mask = torch.bernoulli(torch.ones(len(labels)) * (1 - drop_prob))
return c_hot * c_mask[:, None]  # zeros when mask=0

The model learns that a zero context vector means “no class: generate anything.” At inference, the unconditional pass explicitly passes torch.zeros_like(c).

N_CLASSES = 10
DROP_PROB = 0.10


class ContextUNet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        img_chs = IMG_CH
        down_chs = (64, 64, 128)
        up_chs = down_chs[::-1]
        latent_image_size = IMG_SIZE // 4
        t_dim = 8
        group_size_base = 4
        small_group_size = 2 * group_size_base
        big_group_size = 8 * group_size_base

        self.down0 = ResidualConvBlock(img_chs, down_chs[0], small_group_size)
        self.down1 = DownBlock(down_chs[0], down_chs[1], big_group_size)
        self.down2 = DownBlock(down_chs[1], down_chs[2], big_group_size)
        self.to_vec = nn.Sequential(nn.Flatten(), nn.GELU())

        self.dense_emb = nn.Sequential(
            nn.Linear(down_chs[2] * latent_image_size**2, down_chs[1]),
            nn.ReLU(),
            nn.Linear(down_chs[1], down_chs[1]),
            nn.ReLU(),
            nn.Linear(down_chs[1], down_chs[2] * latent_image_size**2),
            nn.ReLU(),
        )

        self.sinusoidaltime = SinusoidalPositionEmbedBlock(t_dim)
        self.temb_1 = EmbedBlock(t_dim, up_chs[0])
        self.temb_2 = EmbedBlock(t_dim, up_chs[1])

        # Class context embeddings - parallel structure to time embeddings
        self.cemb_1 = EmbedBlock(N_CLASSES, up_chs[0])
        self.cemb_2 = EmbedBlock(N_CLASSES, up_chs[1])

        self.up0 = nn.Sequential(
            nn.Unflatten(1, (up_chs[0], latent_image_size, latent_image_size)),
            GELUConvBlock(up_chs[0], up_chs[0], big_group_size),
        )
        self.up1 = UpBlock(up_chs[0], up_chs[1], big_group_size)
        self.up2 = UpBlock(up_chs[1], up_chs[2], big_group_size)

        self.out = nn.Sequential(
            nn.Conv2d(2 * up_chs[-1], up_chs[-1], 3, 1, 1),
            nn.GroupNorm(small_group_size, up_chs[-1]),
            nn.ReLU(),
            nn.Conv2d(up_chs[-1], img_chs, 3, 1, 1),
        )

    def forward(self, x: Tensor, t: Tensor, c: Tensor) -> Tensor:
        down0 = self.down0(x)
        down1 = self.down1(down0)
        down2 = self.down2(down1)
        latent_vec = self.to_vec(down2)

        latent_vec = self.dense_emb(latent_vec)
        t = t.float() / T
        t = self.sinusoidaltime(t)

        up0 = self.up0(latent_vec)

        cemb_1 = self.cemb_1(c)
        temb_1 = self.temb_1(t)
        up1 = self.up1(cemb_1 * up0 + temb_1, down2)

        cemb_2 = self.cemb_2(c)
        temb_2 = self.temb_2(t)
        up2 = self.up2(cemb_2 * up1 + temb_2, down1)

        return self.out(torch.cat((up2, down0), 1))

graphviz.set_jupyter_format("png")
model_graph = draw_graph(
    ContextUNet(),
    input_size=(
        (BATCH_SIZE, IMG_CH, IMG_SIZE, IMG_SIZE),  # x
        (BATCH_SIZE,),  # t
        (BATCH_SIZE, N_CLASSES),  # c (class context)
    ),
    device="meta",
    depth=1,
    expand_nested=False,
    graph_dir="TD",
    graph_name="ContextUNet with Class Conditioning",
)
model_graph.resize_graph(scale=0.75)
model_graph.visual_graph

def get_context(labels: Tensor, drop_prob: float = DROP_PROB) -> Tensor:
    c_hot = F.one_hot(labels.to(torch.int64), num_classes=N_CLASSES).float()
    c_mask = torch.bernoulli(torch.ones(len(labels), device=labels.device) * (1 - drop_prob))
    return c_hot * c_mask[:, None]


context_model = ContextUNet().to(device)
context_model.compile()
optimizer = optim.Adam(context_model.parameters(), lr=0.001)
epochs = 5

context_model.train()
for epoch in range(epochs):
    for step, (x, labels) in enumerate(dataloader):
        optimizer.zero_grad()

        t = torch.randint(0, T, (BATCH_SIZE,), device=device).float()
        x = x.to(device)
        c = get_context(labels.to(device))

        loss = ddpm.get_loss(context_model, x, t, c)
        loss.backward()
        optimizer.step()

        if epoch % 2 == 0 and step % 100 == 0:
            print(f"Epoch {epoch} | step {step:03d} Loss: {loss.item():.4f}")

Epoch 0 | step 000 Loss: 1.0612
Epoch 0 | step 100 Loss: 0.1687
Epoch 0 | step 200 Loss: 0.1358
Epoch 0 | step 300 Loss: 0.1257
Epoch 0 | step 400 Loss: 0.0861
Epoch 0 | step 500 Loss: 0.1101
Epoch 2 | step 000 Loss: 0.0976
Epoch 2 | step 100 Loss: 0.0980
Epoch 2 | step 200 Loss: 0.0960
Epoch 2 | step 300 Loss: 0.0873
Epoch 2 | step 400 Loss: 0.0818
Epoch 2 | step 500 Loss: 0.0992
Epoch 4 | step 000 Loss: 0.0924
Epoch 4 | step 100 Loss: 0.0900
Epoch 4 | step 200 Loss: 0.0885
Epoch 4 | step 300 Loss: 0.1074
Epoch 4 | step 400 Loss: 0.0986
Epoch 4 | step 500 Loss: 0.0810

Loss drops from 1.07 at the very first step to under 0.10 within the first epoch, and settles around 0.07-0.10 by epoch 4, a smoother, faster descent than the unconditional run. That’s consistent with conditioning making the underlying task easier: the model isn’t hedging across ten silhouettes anymore, just refining one. The next question is whether that translates into better samples, and how the guidance weight \(w\) shapes what comes out at sampling time.

Effect of Guidance Weight

The weight \(w\) is a dial with four distinct regimes:

\(w\)	Effect
\(w < 0\)	Anti-guidance: model actively avoids the class
\(w = 0\)	Pure conditional: no CFG blending applied
\(0 < w \lesssim 1\)	Mild guidance: class loosely present, high diversity
\(w > 1\)	Strong guidance: high class fidelity, lower diversity

Higher \(w\) amplifies the difference between the conditional and unconditional noise predictions at every denoising step. Too high and the model over-corrects, in principle producing near-identical samples regardless of noise seed: mode collapse at inference time.

Note

Confirming that needs multiple seeds at the same \(w\); the single-sample grid below can only hint at it.

@torch.no_grad()
def sample_with_guidance(
    ddpm: DDPM,
    model: ContextUNet,
    class_label: int,
    w: float = 2.0,
) -> np.ndarray:
    x_t = torch.randn((1, IMG_CH, IMG_SIZE, IMG_SIZE), device=device)
    c = F.one_hot(torch.tensor([class_label], device=device), num_classes=N_CLASSES).float()
    c_empty = torch.zeros_like(c)

    for i in reversed(range(ddpm.T)):
        t = torch.full((1,), i, device=device).float()
        e_t_cond = model(x_t, t, c)
        e_t_uncond = model(x_t, t, c_empty)
        e_t = (1 + w) * e_t_cond - w * e_t_uncond
        x_t = ddpm.reverse_q(x_t, t, e_t)

    return x_t.squeeze().cpu().numpy()


fashion_classes = [
    "T-shirt",
    "Trousers",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

fig, axes = plt.subplots(2, 5, figsize=(12, 5))
for cls_idx, (ax, cls_name) in enumerate(zip(axes.flat, fashion_classes, strict=True)):
    img = sample_with_guidance(ddpm, context_model, cls_idx, w=2.0)
    ax.imshow(img, cmap="viridis", interpolation="none")
    ax.set_title(cls_name)
    ax.axis("off")
fig.tight_layout()
plt.show()

At \(w = 2.0\), all ten classes come out recognizable: trousers show two distinct legs, the sneaker its angled toe, the ankle boot its raised heel. This is the same architecture and the same five-epoch budget as the unconditional run above, but given a class label to aim for, it produces clean silhouettes on the first attempt instead of noise. The obvious next question is how much of this is the guidance weight \(w\) doing work, versus conditioning alone: the sweep below varies \(w\) across the same four classes to check.

w_values = [-2.0, 0.0, 1.0, 2.0, 5.0]
selected_classes = [0, 1, 7, 9]  # T-shirt, Trousers, Sneaker, Ankle boot

fig, axes = plt.subplots(len(w_values), len(selected_classes), figsize=(8, 10))
for row, w in enumerate(w_values):
    for col, cls_idx in enumerate(selected_classes):
        img = sample_with_guidance(ddpm, context_model, cls_idx, w=w)
        ax = axes[row, col]
        ax.imshow(img, cmap="viridis", interpolation="none")
        ax.axis("off")
        if row == 0:
            ax.set_title(fashion_classes[cls_idx], fontsize=9)
        if col == 0:
            ax.set_ylabel(f"w={w}", rotation=0, labelpad=35, va="center")

fig.suptitle("Guidance weight (rows) × class (columns)", y=1.01)
fig.tight_layout()
plt.show()

The sweep is a minor surprise. I expected \(w < 0\) to visibly push samples toward “not this class.” Instead, negative guidance mostly makes outputs rougher and less coherent rather than steering them toward a different garment, and raising \(w\) from 0 to 5 changes these single samples less than expected. With only one sample per cell, I can’t rule out that this is just the luck of the random seed rather than a general property of CFG at this scale. That’s a reminder that guidance-strength effects are best judged over many samples, not one.

Takeaways

Working through the workshop left me with a few concrete impressions.

The a_bar cumulative product is the central trick of DDPM training. Without it, computing \(q(x_t \mid x_0)\) for a random \(t\) would require running the forward process step by step. With it, the noisy image at any timestep is a single weighted sum of the original image and random noise. This is what keeps the training loop simple.

The optimizations in module 3 were more impactful than I expected. GELU, GroupNorm, sinusoidal embeddings, and rearrange pooling each feel incremental in isolation, but the difference in generated image quality is significant. The checkerboard artifacts from the naive architecture largely disappear with the full set of improvements.

Classifier-free guidance is conceptually satisfying because it avoids needing a separate classifier. Instead, you teach one model both jobs (conditional and unconditional) and combine their outputs at inference time. The cost is running the model twice per denoising step.

If I had to pick the single biggest surprise of the whole workshop, it’s the conditioning-payoff contrast a few sections back: the same architecture, the same five-epoch budget, producing noise unconditionally and clean garments conditionally. Conditioning didn’t just add control: on a tight compute budget, it made the underlying generation problem dramatically easier to learn.

The full NVIDIA DLI course covers two more modules: CLIP-conditioned text-to-image generation (module 5) and a graded assessment on MNIST digit generation with a 95% classifier accuracy target (module 6). Both are worth working through.

For a deeper theoretical treatment of diffusion models, What are Diffusion Models? is the most thorough walkthrough I’ve found; it covers the connections between DDPM, score matching, and stochastic differential equations that the workshop doesn’t touch.

Footnotes

1. DDPM: Denoising Diffusion Probabilistic Models

2. CFG: Classifier-Free Guidance