I Built a Diffusion Model From Scratch — Here’s Everything That Actually Matters

I spent the last couple of days implementing a diffusion model from scratch on Kaggle. Building the forward process, the U-Net, the noise schedule, and the sampler from the ground up.

The Core Idea

Diffusion models work in two phases.

Forward process: Take a real image. Add Gaussian noise gradually over T timesteps until the image is completely destroyed — pure noise. This is the “corruption” phase and requires no learning.

Reverse process: Train a neural network to undo this corruption. At inference time start from pure noise and denoise step by step until a clean image emerges.

The key insight is that you don’t simulate the forward process step by step during training. The reparameterization trick lets you jump to any noisy timestep directly:

x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * noise

So during training you sample a random timestep t, add exactly the right amount of noise in one operation, and train the model to predict what noise was added. The loss is just MSE between predicted and actual noise.

L = MSE(ε_θ(x_t, t), ε)

Simple objective. Hard to get right in practice.

Cosine vs Linear Noise Schedule

This is one of those things that looks like a minor hyperparameter but actually matters a lot.

Linear schedule adds noise uniformly across timesteps. The problem is the image becomes almost completely destroyed very early and the last timesteps contribute almost nothing useful to the learning signal.

Cosine schedule (from OpenAI’s improved DDPM paper) keeps noise additions small at the beginning and end of the process with a smooth progression in the middle. Every timestep carries meaningful information.

f(t) = cos((t/T + 0.008) / 1.008 × π/2)²
beta_t = 1 - f(t) / f(t-1)

I used cosine by default. The sample quality difference is noticeable especially at lower timesteps.

U-Net Architecture — What’s Different for Diffusion

The backbone is a U-Net but with additions specific to diffusion models.

Sinusoidal Timestep Embeddings

The model needs to know which timestep it’s currently denoising at. This is encoded as a sinusoidal embedding — same idea as positional encodings in Transformers:

python

embeddings = torch.exp(torch.arange(half_dim) * -log(10000) / (half_dim - 1))
embeddings = t[:, None].float() * embeddings[None, :]
embeddings = torch.cat([embeddings.sin(), embeddings.cos()], dim=-1)

This gives a continuous representation of the timestep that the model can use to calibrate how aggressively to denoise.

AdaGN Timestep Conditioning

The timestep embedding is projected to scale and shift values that modulate every residual block:

python

time_out = self.time_mlp(time_emb)
scale, shift = time_out.chunk(2, dim=1)
h = h * (scale + 1) + shift

Instead of just adding the time embedding the model learns to scale and shift the activations. This is significantly more expressive — the model can completely change its behavior at different timesteps.

Memory-Efficient Attention

At the bottleneck I added self-attention using PyTorch’s scaled_dot_product_attention with Flash Attention:

python

with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_mem_efficient=True):
    out = F.scaled_dot_product_attention(q, k, v, scale=self.scale)

This was specifically to avoid OOM errors on T4 GPUs while still getting global context at the bottleneck. Standard attention at 128×128 resolution would blow memory immediately.

GroupNorm over BatchNorm

I used GroupNorm throughout instead of BatchNorm. With batch size 16 and image generation tasks, GroupNorm is more stable — it normalizes within each sample independently of the batch.

EMA — Non-Negotiable

EMA (Exponential Moving Average) maintains a smoothed copy of the model weights during training. Only the EMA model is used for inference.

The difference between raw training weights and EMA weights for image generation is significant. Raw weights produce noisier, less coherent samples. EMA weights are much more stable. It adds almost zero overhead and should always be used for diffusion models.

DDIM — 8× Faster Sampling

Standard DDPM sampling runs all T=400 timesteps. That’s slow.

DDIM (Denoising Diffusion Implicit Models) uses a subset of timesteps with a deterministic update rule:

pred_x0 = (x_t - sqrt(1-alpha_bar) * pred_noise) / sqrt(alpha_bar)
x_{t-1} = sqrt(alpha_bar_prev) * pred_x0 + sqrt(1-alpha_bar_prev) * pred_noise

With eta=0.0 the sampling is fully deterministic — same noise vector always produces the same image. I used 50 DDIM steps instead of 400, giving 8× speedup with comparable quality.

Image Reconstruction

For reconstruction I took a target image, added noise up to timestep t=300, then ran DDIM denoising back to t=0. The result is a reconstruction of the original image that can be evaluated with SSIM and PSNR.

This is useful for tasks like inpainting and image editing where you want to modify an image rather than generate from scratch.

Training Setup

Dataset     : CelebA-HQ 256 (128×128 resize)
T           : 400 timesteps
Schedule    : Cosine
Batch size  : 16
Epochs      : 50 (early stopping, patience=10)
Optimizer   : AdamW (lr=2e-4, weight decay=1e-4)
Scheduler   : CosineAnnealingLR
Grad clip   : 1.0
Precision   : AMP (bfloat16)
GPUs        : Dual T4 on Kaggle

Mixed precision and gradient clipping were both essential. Without grad clipping the training occasionally produced NaN losses in early epochs.

Key Takeaways

The noise schedule is not a minor detail. Cosine vs linear makes a real difference in sample quality especially on fewer timesteps.

EMA weights always for inference. Never use raw training weights for generation. The overhead is negligible.

Flash Attention makes the difference on T4. Without it standard self-attention at meaningful resolutions causes OOM. With it you get global context without blowing memory.

DDIM is the practical sampler. 400 steps is too slow for anything interactive. 50 steps with DDIM is fast and sharp.

🤗 Live Demo: HuggingFace Space