Guided diffusion from scratch with 2-D examples

A developer published a tutorial on guided diffusion using 2-D examples, demonstrating vanilla conditional flow matching and classifier-free guidance for generating images conditioned on labels such as 'cat' or 'dog'. The post includes code for a conditional flow model and explains how to scale guidance strength without a separate classifier.

This is a direct follow-up to the earlier 2-D diffusion post . The code below is trimmed to the essentials; the full, runnable version lives in the original notebook:github.com/litlig/notebooks/guided 2d diffusion.ipynb Problem framing problem-framing In an unguided diffusion model, we sample a noise image and follow the diffusion process to arrive at an image. Often, though, we want to prompt the model to generate a specific kind of image; such generation is called guided generation . We continue with the 2-dimensional example — a two-pixel image. We say the cluster in the upper right are cat images, and the cluster in the lower left are dog images. Vanilla guidance vanilla-guidance With a text label/prompt \ y\ , we want to learn a guided vector field that moves the probability mass toward the distribution of images conditioned on the label, \ p {data} \cdot\,|\,y \ . Intuitively, we can follow the same procedure as unguided flow matching. For each training example, we sample a pair of image and label or text caption . We train a neural network to estimate the conditional vector field, giving it the label \ y\ as an additional input. The loss function is still the mean squared error: \ \min {\theta}\ \mathbb{E} {t \sim U 0,1 ,\ z,y \sim p {data},\ x \sim p \cdot\,|\,z }\ \big\| f t^\theta x, y - v t x\,|\,z \big\|^2\ When the target \ z\ is given, the conditional vector field is computed the same way as before: \ v t x\,|\,z = \dot{a} t\, x 0 + \dot{b} t\, z\ The label enters the network through an embedding, which is concatenated with \ x t\ and \ t\ : python class ConditionalFlowModel nn.Module : def init self, embedding dim=2, num classes=2 : super . init self.embedding = nn.Embedding num classes, embedding dim input size = 2 + 1 + embedding dim xt + t + label embedding self.net = nn.Sequential nn.Linear input size, 128 , nn.SiLU , nn.Linear 128, 128 , nn.SiLU , nn.Linear 128, 128 , nn.SiLU , nn.Linear 128, 2 , def forward self, xt, t, y : y embedded = self.embedding y.squeeze -1 return self.net torch.cat xt, t, y embedded , dim=-1 Sampling then just fixes the guidance label and integrates the field with Euler steps. With guide label=0 we pull noise toward the dogs, with guide label=1 toward the cats: Classifier-free guidance classifier-free-guidance In this simple example, the vanilla approach above already works reasonably well. In practice, though, samples from such a model often don’t follow the prompt or label very strongly. The guided probability path is \ p t x\,|\,y \ . By Bayes’ theorem: \ p t x\,|\,y = \frac{p t y\,|\,x \, p t x }{p t y }\ Taking the log of both sides and then the gradient with respect to \ x\ the \ p t y \ term drops out, since it doesn’t depend on \ x\ : \ \nabla \log p t x\,|\,y = \nabla \log p t y\,|\,x + \nabla \log p t x \ \ \nabla \log\ is the score function . For a Gaussian probability path, the vector field relates to the score by \ v t x = \alpha t \nabla \log p t x + \beta t x t\ . Substituting gives a formula that connects the two vector fields: The extra term \ \nabla \log p t y\,|\,x \ is what contributes the guidance. A natural way to make generation follow the guidance more strongly is to scale this term up. Estimating \ p t y\,|\,x \ is essentially a classification problem — given a noisy sample \ x\ , predict its label — so this approach is called classifier guidance : it requires training one network for the vector field and a separate classifier. A more elegant approach eliminates the need for a separate classifier. Given the guided vector field \ v t x\,|\,y \ and the unguided vector field \ v t x \ , the influence of guidance is \ v t x\,|\,y - v t x \ . We scale it up by a weight \ w 1\ and use the following in forward sampling: \ 1-w \, v t x + w\, v t x\,|\,y \ To get both \ v t x \ and \ v t x\,|\,y \ from a single network, we treat the unguided vector field as a guided one with a special null prompt, which we inject at random into a fraction of the training samples: python def get batch dist fn, batch size, inject rate=0.1 : z, y = dist fn batch size z = torch.from numpy z .float y = torch.from numpy y .long .unsqueeze 1 x0 = torch.randn batch size, 2 t = torch.rand batch size, 1 xt = 1 - t x0 + t z vf = z - x0 Randomly replace some labels with the special 'null' label 2 num inject = int batch size inject rate if num inject 0: inject indices = torch.randperm batch size :num inject y inject indices = 2 return xt, t, y, vf At sampling time we run the network twice per step — once with the real label, once with the null label — and combine them with the guidance weight: python def cfg sample model, nsample, steps=100, guide label=0, weight=2.0 : null label = 2 x = torch.randn nsample, 2 dt = 1.0 / steps y guide = torch.full nsample, 1 , guide label, dtype=torch.long y null = torch.full nsample, 1 , null label, dtype=torch.long for step in range steps : t = step / steps torch.ones nsample, 1 vf guided = model x, t, y guide vf unguided = model x, t, y null x = x + 1 - weight vf unguided + weight vf guided dt return x Sweeping the guidance weight shows its effect directly. At \ w=1\ we recover plain conditional sampling. As \ w\ grows, the guidance term is amplified: samples cling more tightly to their target cluster and drift further from the opposite one — the samples follow the prompt more faithfully, at the cost of some diversity as they concentrate.