No-Confusion Diffusion

A simple and minimal exercise repo to get familiar with basic Diffusion Models.
See the paper here.

_{Portrait: Von Gottlieb Biermann / Nach Christian Albrecht Jensen — Gauß-Gesellschaft Göttingen e.V. (Foto: A. Wittmann). Gemeinfrei, Wikimedia Commons.}

Getting Started

This repository is build using pytorch for development and Omegaconf for configuration management. All config files can be found in configs. To keep compute and complexity of the task manageable, this repo only features two simple tasks/datasets: fashion-mnist (see fashion_mnist.yaml) and cifar10 (see cifar10.yaml).

However, others can easily be added.

Prerequisites

uv is used as the project and package manager, and I highly recommend using it to create the virtual environment and run the code:

Install uv using:

curl -LsSf https://astral.sh/uv/install.sh | sh # macOS and Linux

Windows? Please don't use Windows!

Running the code

The main entry point of this project is main.py.

You can run the whole code as is, with the solution already provided, to check if everything runs correctly:

uv run python main.py --config configs/fashion_mnist.yaml --solution

Note: Having a GPU is encouraged for faster training.

You may override any default argument:

uv run python main.py --config configs/fashion_mnist.yaml training.epochs=100

For a full run that should yield good results on a GPU (e.g. T4 on Colab):

uv run python main.py --config configs/fashion_mnist.yaml --solution training.epochs=50 data.batch_size=256

If you have a strong GPU (e.g. RTX 5090):

uv run python main.py --config configs/fashion_mnist.yaml --solution training.epochs=50 data.batch_size=512

Generated samples will be saved to outputs/samples.png after training.

Your Task

There are generally two components to (basic) Diffusion:

1. The actual model performing the noise prediction: f(img, timestep) = noise.
2. The Diffusion Process, meaning noise schedule shape, number of timesteps, sampling strategy (DDPM vs DDIM vs others), loss weighting across timesteps, and so on.

For the --solution flag, bot hare already implemented. When omitting the --solution flag, you have to implement 1. (the model), while 2. is still taken care of automatically.

Component 1

Implement the DiffusionModel class in src/model.py.

Your model receives:

• x: noisy images, shape (B, C, H, W)
• t: integer timesteps, shape (B,), values in [0, T)

And returns the predicted noise, same shape as x.

Everything else (training, diffusion process, sampling) is already provided. The solution model and sampler live in src/solution/ -- feel free to look at the sampler to understand the reverse process, but the model is a HuggingFace wrapper and won't help you much architecturally.

Once your model is implemented, run:

uv run python main.py --config configs/fashion_mnist.yaml

Component 2

As an extra step/goal, you can implement the Diffusion Process (includes the reverse process) around the model, which is also quite interesting but more complex.

Tips

A SinusoidalTimestepEmbedding module is available in src/utils.py for mapping integer timesteps to dense vectors. Import and use it freely.

Some architecture starter tips:

• Start with a simple ConvNet + time embedding projection, see if it trains at all
• Add skip connections (UNet-style) for better gradient flow
• GroupNorm + SiLU is a solid combo for normalization and activation
• Attention layers can further improve sample quality

Challenges

Level	Component	Goal
1	1	Implement a working model that produces recognizable Fashion-MNIST samples
2	1	Beat the solution model's validation loss on Fashion-MNIST
3	1	Get recognizable results on CIFAR-10 (--config configs/cifar10.yaml)
4 (Stretch)	2	Implement the DDPM sampling loop yourself in src/sampler.py, then run with --custom-sampler