How to Learn Diffusion Models from Scratch, One Stage at a Time

How do you learn diffusion models?

You learn diffusion models by writing each part yourself in the order the model uses it, starting with how data is noised and ending with how a sampler turns noise back into data. The parts are the noise schedule, the closed-form forward sample, the noise prediction loss, the reverse sampling step, faster samplers, the score and flow views, guidance, the denoising network and latent diffusion. TensorTonic has a coding problem with tests for each of these, and the Build Diffusion Model from Scratch study plan collects 26 of them into seven sections.

This guide works as a diffusion model tutorial in nine stages. Each stage has a short explanation and a table that pairs every topic with a page to read and a problem to solve, so you can build a diffusion model from scratch one tested function at a time.

What is a diffusion model?

A diffusion model is a generator that learns to create data by learning to remove noise from it. During training, clean images are corrupted with Gaussian noise at a randomly chosen strength, and a network learns to predict the noise that was added. To generate, the model starts from pure random noise and removes a little of it at a time over many steps until an image is left.

The idea works because adding noise is easy and fully known, while removing it is a learning problem with a simple regression target. Each training example is a noisy input and the exact noise inside it, so the network is trained with an ordinary mean squared error loss. Image generators such as Stable Diffusion are built on this method, with the extra pieces covered in the later stages of this guide.

What do you need to know before learning diffusion models?

You need Python, NumPy, the basics of PyTorch and some probability before you start learning diffusion models. The probability is narrower than it looks, since almost everything rests on the normal distribution, its mean and variance, and the KL divergence between two Gaussians. You do not need to know GANs, and you do not need to follow the full variational derivation before you can write working code.

• For PyTorch, you should be able to write a small network and a training loop, and the guide to learning PyTorch and the PyTorch study plan cover both.
• For probability, you should understand sampling from a Gaussian with a chosen mean and variance, which the probability study plan covers.
• For networks, you should know convolutions, residual blocks and attention at the level of the deep learning interview questions guide.

The first section of the study plan covers the Gaussian foundations directly, so you can check yourself on them before the diffusion stages begin.

In what order should you learn diffusion models?

You should learn diffusion models in the order the data moves through them, which is forward noising first, then training, then sampling, and then the improvements on each of those. The forward process defines what the network has to undo, the loss depends on the noise schedule, and the sampler reuses both. The network architecture comes later than most people expect, because a small multilayer perceptron on two-dimensional points is enough to test everything before it.

Stage 1: How does the forward process add noise?

The forward process adds a small amount of Gaussian noise to the data at each of many steps, often 1,000, until nothing of the original is left. At step t it scales the previous sample by the square root of one minus beta t and adds noise with variance beta t, and the list of beta values is called the noise schedule. The original DDPM paper used betas that grow linearly, and the cosine schedule was introduced later because the linear one destroys small images too quickly.

The signal-to-noise ratio at each step is the most useful way to compare schedules. Many common schedules leave a little signal at the final step, so the model never sees pure noise in training even though sampling starts from it, and rescaling the schedule to zero terminal SNR removes that mismatch.

Stage 2: How do you jump to any noise level in one step?

You jump to any noise level by using the fact that a chain of Gaussian steps is itself a single Gaussian. If alpha bar t is the product of one minus beta over the first t steps, then the noisy sample is the square root of alpha bar t times the clean data plus the square root of one minus alpha bar t times standard normal noise. This closed form lets training pick a random timestep and produce the noisy input directly, without running the steps before it.

The formula is the reparameterisation trick from VAEs applied to the forward process. It also assumes the data has roughly the same scale as unit Gaussian noise, which is why images are usually scaled to the range from minus one to one before training.

Stage 3: What does a diffusion model learn to predict?

A diffusion model usually learns to predict the noise that was added to a clean input. Each training step picks a clean example, a random timestep and fresh Gaussian noise, builds the noisy input with the closed form, and minimises the mean squared error between the true noise and the predicted noise. This simple loss comes from the evidence lower bound, where most terms are KL divergences between two Gaussians, and Ho and colleagues found that dropping the weights on those terms gave better samples.

Predicting the noise is one of three equivalent targets. The network can also predict the clean data or a velocity that mixes the two, and any of them can be converted into the others with the schedule values at that step. The choice changes how the loss is weighted across timesteps, which is what methods such as Min-SNR weighting adjust directly.

Stage 4: How does a diffusion model generate a sample?

A diffusion model generates a sample by starting from pure Gaussian noise and applying a learned reverse step once for every timestep, from the noisiest to the cleanest. Each DDPM step uses the predicted noise to estimate the mean of the slightly cleaner sample, then adds fresh noise scaled by the posterior standard deviation, except at the final step. The mean and variance come from the posterior of the forward process, which is why the posterior problem from the previous stage matters here.

Every reverse step is a full forward pass through the network, so a thousand-step sampler is slow, and samplers that need fewer steps matter a great deal in practice.

Stage 5: How do DDIM and other samplers make generation faster?

DDIM makes generation faster by replacing the random reverse step with a deterministic one that can skip most of the timesteps. It turns the noise prediction into an estimate of the clean data and then rebuilds the sample at an earlier noise level from that estimate and the same noise prediction. It works with a network trained by the ordinary DDPM loss, so you get faster sampling with no retraining.

The work of Karras and colleagues treats sampling as solving an ordinary differential equation over the noise level. Their schedule of noise levels spends more steps near the clean end, and an Euler or Heun solver over those levels gives good samples in a few dozen steps. The Frechet distance problem lets you measure whether a faster sampler has cost you quality, since it is the calculation behind FID.

Stage 6: How do score matching and flow matching relate to diffusion?

Score matching and flow matching are two other ways to train the same kind of noise-to-data generator. The score of a distribution is the gradient of its log density, and a network trained to predict noise is, up to a scale factor that depends on the timestep, an estimate of the score of the noised data. Denoising score matching trains that score directly, and annealed Langevin dynamics samples with it by taking small gradient steps plus noise at a sequence of falling noise levels.

Flow matching trains a network to predict a velocity that carries a noise sample to a data sample along a chosen path. With a straight path between the two, the target velocity is simply the data minus the noise, and sampling integrates that velocity with a few Euler steps. Rectified flow uses these straight paths, which is why it degrades more gracefully than DDPM when you cut the number of sampling steps.

Stage 7: How does classifier-free guidance steer a diffusion model?

Classifier-free guidance steers a diffusion model by combining a conditional and an unconditional noise prediction at every sampling step. During training the condition, such as a class label or a text embedding, is replaced with an empty token for a small fraction of examples, so one network learns both predictions. At sampling time the guided prediction is the unconditional prediction plus a guidance scale times the difference between the conditional and unconditional ones.

A larger guidance scale makes samples follow the condition more closely and reduces their variety, so the scale is a trade-off you tune. Text-to-image models read the prompt through cross-attention, where the image features provide the queries and the text embeddings provide the keys and values.

Stage 8: Should the denoising network be a U-Net or a Transformer?

The denoising network can be either a U-Net or a Transformer, and the choice mostly depends on scale. A U-Net shrinks the image in an encoder, grows it back in a decoder, and carries fine detail across with skip connections, while a sinusoidal timestep embedding is added inside each residual block so the network knows how noisy its input is. Self-attention at the lower resolutions lets distant parts of the image affect each other.

A Diffusion Transformer, or DiT, cuts the input into patches and processes them with Transformer blocks. It conditions each block with adaptive layer normalization, where the timestep and class embedding produce a scale and a shift, and adaLN-Zero adds a gate initialised to zero so that each block starts as the identity. DiTs scale more predictably with compute, which is why many recent large image and video models use them.

Stage 9: What is latent diffusion?

Latent diffusion is diffusion run on the compressed latent space of an autoencoder instead of on raw pixels. An autoencoder is trained first to encode an image into a much smaller tensor and decode it back, and the diffusion model is then trained and sampled entirely on those latents. Stable Diffusion works this way, and its autoencoder reduces each spatial side of the image by a factor of eight, which makes every denoising step far cheaper.

Nothing in the earlier stages changes when you move to latents, since the schedule, the loss and the sampler all work on whatever tensor they are given. The new parts are the autoencoder and the final decode, and the VAE lessons explain the KL term that keeps the latent space smooth enough to generate from.

Can you learn diffusion models without a GPU?

Yes, you can learn diffusion models without a GPU, because every component can be written and tested on small arrays or two-dimensional toy data. A GPU only matters when you train an image model at a useful size, and even then a small model on small images trains quickly. The skills you are building are the formulas and the shapes, and those are identical on a laptop and on a cluster.

The coding problems linked in this guide run in the browser and are checked against tests, so there is nothing to install. Our guide on practising ML without a GPU covers where to find compute when you want to train something larger.

How do you put the pieces together into a model that generates?

You put the pieces together by training on a toy dataset first, checking each part against a known answer, and only then moving to images. The Generate Images with Diffusion and Rectified Flow project follows this path in PyTorch, training DDPM and rectified flow on a 2D spiral and then on 14 × 14 clothing images, and it measures how many sampling steps each model needs.

1. Scale the data so it has roughly the same range as unit Gaussian noise.
2. Build the schedule and confirm that alpha bar falls steadily towards zero at the last step.
3. Noise a batch to the last step and check that its mean is close to zero and its standard deviation close to one.
4. Check that a network which always predicts zero noise gets a loss of about one, which is the variance of the noise.
5. Train on the 2D toy data until the loss falls clearly below that baseline.
6. Sample with the full DDPM loop, plot the samples against the data, and then try DDIM with far fewer steps.

What should you build after your first diffusion model?

After your first diffusion model, the useful next steps are sampling in very few steps, conditioning on text, and moving to Transformer backbones. Each of these is one of the stages above taken further, and the table pairs each direction with something to read and a project or plan to work through.

If you are preparing for interviews, the computer vision interview questions guide has a section on generative models that compares GANs, VAEs and diffusion and explains how a diffusion model is trained. For a wider list of papers to build in order, see the ML papers worth implementing from scratch.

What mistakes do people make when learning diffusion models?

• Starting with the full derivation before writing any code slows most people down, so write the forward sample and the loss first and read the derivation once they work.
• Mixing up beta, alpha and alpha bar, or indexing them off by one step, gives samples that are slightly wrong with no error message, so print the values at the first and last step.
• Feeding images in the range from zero to one instead of minus one to one leaves the data on a different scale from the noise.
• Sampling the timestep once per batch instead of once per example makes the loss jump around and slows training.
• Adding fresh noise on the final reverse step leaves visible grain in the output.
• Judging progress by the training loss alone hides poor samples, since the loss stays fairly flat late in training, so look at samples and measure a distance such as FID.
• Jumping straight to a large U-Net on real images makes every bug expensive, and a small network on 2D points shows the same bugs in seconds.

Practice this

Then work through the full Build Diffusion Model from Scratch study plan, or train both models end to end in the Generate Images with Diffusion and Rectified Flow project.

Related guides

• How to build an LLM from scratch, one component at a time
• Computer vision interview questions with answers
• How to learn PyTorch, from tensors to a training loop
• 15 ML papers worth implementing from scratch, and in what order
• How to read an ML paper and implement it
• Maths for machine learning
• How to learn computer vision, from pixels to vision transformers