Diffusion Transformers (DiT)

#What the Denoiser Has to Do

Start with the job, before any new part. At every step of generation, the network gets three things. It gets the current noisy latent, a number that says how noisy it is (the time), and a description of what you want (a class label or a prompt). It returns a prediction with exactly the same shape as the latent. The sampler then uses that prediction to take one small step toward a clean picture.

So the network is a function from "grid of numbers, plus time, plus a wish" to "grid of numbers of the same size". Nothing in that job says it must use convolutions. Any network that keeps the shape and can mix information across the whole grid will do.

#Why Swap the U-Net for a Transformer

A U-Net is a stack of convolution layers that shrinks the grid, then grows it back, with shortcuts between matching sizes. A convolution looks only at a small neighbourhood. To let far-apart cells talk, the U-Net shrinks the grid, and Stable Diffusion's U-Net also adds attention layers at some of its sizes.

A transformer works differently. Every block lets every token look at every other token. The same block is repeated, say 28 times, and nothing else is needed. Three reasons are given for the swap.

There is a price. A convolution assumes that nearby cells matter most, and that guess is usually right for pictures. A transformer has no such guess built in. It must learn from data that nearby patches are related. Its attention also costs more as the token count grows, which the cost section shows.

#Step 1: Cut the Latent into Patches

A transformer reads a list of tokens, but a latent is a grid. The fix is simple: cut the grid into small squares and treat each square as one token. This is called patchifying. Vision transformers do the same to ordinary pictures.

Picture first. Take a latent that is 8 cells wide and 8 tall, with 4 numbers per cell (4 channels). Cut it into squares of 2 by 2 cells. The grid then has 4 squares across and 4 down, so 16 squares in all. Each square holds 2 x 2 x 4 = 16 numbers. Flatten each square into a list, and you have 16 tokens, each a list of 16 numbers.

Now a real size. A 256 by 256 colour image goes through the VAE of the latent diffusion lesson, which shrinks each side by 8. The latent is 32 by 32 with 4 channels. Cut that into 2 by 2 squares, and you get 16 across and 16 down: 256 tokens.

The rule is the latent side divided by the patch side, then squared. For a 32 by 32 latent, patch sizes 2, 4 and 8 give 256, 64 and 16 tokens. Halving the patch side quadruples the token count. The patch size is a dial: small patches keep fine detail and cost much more.

Each flattened square is then mapped to the model width by one learned linear layer. That layer turns a list of 16 numbers (in the small example) into a list of width d numbers. After it, every token has the same size, and the transformer can start.

The DiT paper names its models by size and patch size together. "DiT-XL/2" means the extra-large network with a patch size of 2. "DiT-XL/8" would be the same network with patch size 8, which has far fewer tokens and therefore a far smaller cost per call. The name is a reminder that two separate dials set the cost: the network size and the patch size.

Why not use a patch of one cell, so that no detail is lost? The token count would reach 1,024 for a 32 by 32 latent, four times the p = 2 count, and the attention cost would grow even faster. The patch size trades detail per token against the number of tokens, and every design picks a point on that trade.

• Do: Step through the five stages and switch the position overlay on at the fourth stage.
• Notice: An 8 by 8 picture cut into 2 by 2 squares gives 16 tokens, the same count as the small latent above. Each token gets its own colour once position is added.
• Why: Cutting, flattening and projecting are what patchify means. Drawn stand-in, not a trained model: the picture is a hand-made digit, the token colours come from how filled each square is, and the extra class token at the front is a vision transformer habit that a DiT does not need.

A last point about patches. A token holds only the numbers of its own square, so at the start it knows nothing about its neighbours. Knowledge about neighbours must come from attention. In block one, a token can already compare itself with all the others. A convolution would need several layers before a far-away cell could have any influence. This is the main structural difference between the two designs.

#Step 2: Tell Each Token Where It Is

Once the squares are in a list, the network has lost something. A list has an order, but attention itself does not use it. If you shuffled the tokens, plain attention would not notice. That is a problem for pictures, because the top-left square and the bottom-right square are not interchangeable.

The fix is to add position numbers to each token before the first block. Picture a chessboard where every square has a name, such as "row 3, column 5". The position numbers are that name, written as a list of numbers.

A tiny count shows the idea. In the small example, the 16 tokens come from a 4 by 4 grid of squares. Token 0 sits at row 0, column 0. Token 5 sits at row 1, column 1. Token 15 sits at row 3, column 3. The numbers differ for each, so the network can tell them apart after attention mixes everything.

Other models choose other position schemes. Some learn the table, and some rotate the queries and keys by the position. The reason does not change: without it, a grid becomes a bag of squares.

#Step 3: Tell Every Block the Time and the Class

Position says where. The network also needs to know when and what. Early in sampling the latent is almost pure noise, and the right move is bold. Late in sampling it is nearly clean, and the right move is a small touch. The block must behave differently at each noise level. It must also know the class label (such as "dog") or the prompt.

First, combine the two into one vector, called c. The time is turned into a list of sine and cosine numbers, the same trick as position. The class label is looked up in a table that gives a learned list of numbers. Adding them gives c, one vector for the whole picture.

Now the question is how to push c into every block. The DiT paper compared four designs: putting c in as extra tokens, using cross-attention, adaptive layer norm, and adaptive layer norm with a zero-start gate. In the original experiments the last one worked well.

Take one number to see it. Suppose a normalised value is 0.5. At a noisy time, c says scale = 1.0 and shift = 0.2. The output is 0.5 x (1 + 1.0) + 0.2 = 1.2. At a clean time, c says scale = -0.5 and shift = 0.0. The output is 0.5 x (1 - 0.5) = 0.25. Same block, same input, different behaviour.

This zero start is called adaLN-Zero. The derivation below shows why training still gets going.

Here is how the whole conditioning story looks in one sentence. The short facts, time and class, enter by modulation, which changes how every token is normalised. The long facts, the words of a prompt, enter as tokens and are read by attention. Position, the third kind of fact, is added once at the start. Three kinds of information use three routes, and each route fits its kind.

#Step 4: Let Text Join the Attention (MM-DiT)

A class label is one small vector, so adaptive layer norm carries it well. A prompt is different. It is a list of tokens, one per word piece, and each carries its own meaning. Squeezing it into one vector would lose detail. So text needs a second route.

The MM-DiT idea, described in the Stable Diffusion 3 paper, is different. It puts the image tokens and the text tokens side by side in one long list, and attention runs over the whole list. Every image token can look at every text token, and every text token can look at every image token. Information flows both ways, and the text tokens are updated in each block too.

The two kinds of token are not alike, so they do not share weights. The image stream has its own projection matrices and its own MLP. The text stream has its own as well. Only the attention step is shared: queries, keys and values from both streams are joined, one softmax is taken, and the result is split back into two parts. This is why it is called two streams with one attention.

The paper also feeds a single summary vector of the prompt, together with the time, into the modulation, in the same way as the label above. So both routes are used: modulation for the short summary, and joint attention for the full text.

A small count shows the price. Say there are 256 image tokens and, as an example, 100 text tokens. Image-only attention scores 256 x 256 = 65,536 pairs. Joint attention scores (256 + 100)^2 = 126,736 pairs, nearly twice as many. The text length here is an example, not the length any real model uses.

#The Whole Network in One Picture

Here is the path from a noisy latent to a prediction. The block in the middle is repeated.

The last two stages undo the first. A final layer norm and a linear layer turn each token of width d back into p x p x C numbers. Unpatchifying then puts every square back where it came from. The output has the shape of the input latent, as the job description demanded.

#Build One Tiny Block

Below is a forward pass of one DiT-style block in numpy. It is a toy with random weights, so it learns nothing. Its purpose is to let you see every piece from this lesson as a few lines, and to check shapes. The latent is 8 by 8 by 4, the patch is 2 by 2, and the width is 16.

The printed lines match the comments in the solution box. The block here has one attention head and no layer-norm weights, so it is smaller than a real one. A real block has several heads, and a real network stacks many blocks. The structure is the same.

# Copy to Colab: needs torch and diffusers, and a GPU, so it cannot run on this page.
# Load a published class-conditional DiT (trained on ImageNet) and draw two pictures.
import torch
from diffusers import DiTPipeline, DPMSolverMultistepScheduler

pipe = DiTPipeline.from_pretrained("facebook/DiT-XL-2-256", torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

class_ids = pipe.get_label_ids(["white shark", "umbrella"])   # ImageNet class names
generator = torch.manual_seed(33)
images = pipe(class_labels=class_ids, num_inference_steps=25, generator=generator).images
images[0].save("shark.png")
images[1].save("umbrella.png")

#Compute Against Token Count

Cost is where the patch size bites. One block costs about 24 T d^2 + 4 T^2 d FLOPs, from the derivation above. The first term grows in step with the token count T. The second grows with T squared, because attention scores every pair.

Use the numbers from the DiT paper's largest network: width d = 1152 and 28 blocks. On a 32 by 32 latent, three patch sizes give these costs for one call of the network, counted with the formula.

Each halving of the patch side multiplies the cost by about four. At 256 tokens the count is 118.4 G multiply-adds, which is close to the 118.6 Gflops the paper reports for this network. You pay that price at every sampling step, so a 30-step run is about thirty times that.

Why does the second term matter? The video lesson works out the ratio of attention cost to the rest: it is T divided by 6d. At 256 tokens and d = 1152 the ratio is 0.037, so attention is a small part. At 3,840 tokens, the size of the video lesson's example clip, the ratio is 0.556, and one call costs 5.33e12 FLOPs. At 7,680 tokens the ratio passes one, and doubling the tokens raises the cost by 2.71 times, not two.

• Do: Pick image size 224, then switch the patch size from 32 to 16 to 8, and read the token count and the red cost.
• Notice: Each halving of the patch size quadruples the tokens, and the red global-attention cost grows about sixteen-fold.
• Why: Scoring every pair of tokens gives N squared entries, and N itself grows as one over the patch size squared. Exact counts on a toy pixel picture with 3 colour channels, not a latent and not a trained model. Ignore the green windowed line: a plain DiT uses global attention.

The numbers here count FLOPs, not seconds. Real speed depends on the machine, the numeric precision and the software. Measure on your own hardware before you budget.

#DiT and Flow Matching: Network Versus Target

The DiT paper used the usual diffusion recipe: the network predicts the noise, trained with the squared error from the diffusion lesson. The Stable Diffusion 3 paper trained its MM-DiT with rectified flow, where the network predicts a velocity along a straight path. Same family of network, different target.

Any backbone can be paired with any target, since they only meet at the input and the output of the network. Flow matching changes the loss and the path. A DiT changes the machine that reads the noisy latent. When someone says "SD 3 is a flow model" and "SD 3 is a transformer", both statements are true, and they describe two different choices.

#Key Takeaways

• The denoising network takes a noisy latent, a time and a condition, and returns a same-shaped prediction. A transformer can do that job
• Patchify cuts the latent into squares and treats each as a token. The count is the latent side divided by the patch side, squared. Halving the patch side quadruples the tokens
• Position numbers are added because attention alone ignores order. Time and class reach every block through adaptive layer norm, with a gate that starts at zero so the stack starts as the identity
• MM-DiT keeps two streams, image and text, each with its own weights, and joins them in one attention so information flows both ways
• Cost per block is about 24 T d^2 + 4 T^2 d. Attention starts to dominate when the token count passes 6 times the width
• DiT is the network. Diffusion or flow matching is the training target. They are separate choices

#Quick Check