Diffusion transformers (DiT)

The network's job: noisy tokens and a time in, one velocity per token out

Picture a cloud of 1,024 points, part way from noise to a shape, at time t. In flow matching, the network must return 1,024 arrows: the velocity of each point. Each arrow depends on where the other points are. A point at the tip of a half-formed branch should move differently from one in the middle of a cluster.

A DiT handles this in three stages. First, a linear layer turns each point's 3 coordinates into a vector of D numbers, its token (see embeddings). For an image, a token is a small square patch instead of a point. Second, a stack of blocks updates all the tokens. Third, a final layer turns each token into that token's output.

Inside a block: attention and an MLP, each steered by the time

A DiT block is a transformer block with two changes, both from time conditioning. First, each layer normalization takes its scale and shift from the conditioning vector c. Second, each branch's output is multiplied by a gate before it is added back. A small layer (a SiLU activation, then a linear layer) maps c to six vectors of D numbers: γ₁, β₁, α₁ for the attention branch and γ₂, β₂, α₂ for the MLP branch. With LN for layer normalization and ⊙ for feature-by-feature products, the block computes:

z = LN(h) ⊙ (1 + γ₁) + β₁
h ← h + α₁ ⊙ Attention(z)

z = LN(h) ⊙ (1 + γ₂) + β₂
h ← h + α₂ ⊙ MLP(z)

The final layer is one more adaLN step (a scale and a shift, no gate), then a linear layer from D numbers to the output size of each token: 3 for a velocity in 3D. In the DiT paper, the layers that make the six vectors and the final linear layer all start at zero. So the untrained network passes tokens through every block unchanged and outputs zero for every token.

Worked example: the sizes for a point cloud

Your turn. With 2,048 points per cloud, how many token pairs does attention compare per head and per block? And how many parameters does each block's conditioning layer have when D = 768?

dit_block.py (PyTorch)

import torch
import torch.nn as nn

class DiTBlock(nn.Module):
    # One transformer block whose normalization is set by the time (adaLN-Zero).
    def __init__(self, d, heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(d, elementwise_affine=False)
        # Attention with several heads. batch_first=True puts the batch first, as in h.
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.norm2 = nn.LayerNorm(d, elementwise_affine=False)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(d, 6 * d))
        nn.init.zeros_(self.ada[1].weight)    # "Zero": every gate starts at 0
        nn.init.zeros_(self.ada[1].bias)

    def forward(self, h, c):                  # h: (batch, tokens, d); c: (batch, d)
        # [:, None] adds an axis of length 1, so each of the 6 parts acts on all tokens.
        g1, b1, a1, g2, b2, a2 = self.ada(c)[:, None].chunk(6, dim=-1)
        z = self.norm1(h) * (1 + g1) + b1
        # z, z, z: the queries, keys and values all come from z. [0] keeps the output.
        h = h + a1 * self.attn(z, z, z, need_weights=False)[0]
        z = self.norm2(h) * (1 + g2) + b2
        return h + a2 * self.mlp(z)

torch.manual_seed(0)
block = DiTBlock(d=384, heads=6)
h = torch.randn(2, 1024, 384)                 # 2 examples, 1024 tokens each
c = torch.randn(2, 384)                       # time (plus condition) embedding
out = block(h, c)
print(out.shape, "unchanged at the start:", torch.equal(out, h))
print("parameters in the time layer:", sum([p.numel() for p in block.ada.parameters()]))

torch.Size([2, 1024, 384]) unchanged at the start: True
parameters in the time layer: 887040

Without position labels, a DiT treats points as a set

An image DiT adds a position embedding to each patch token, because a patch's place in the grid matters. A point token already carries its position: its coordinates. If nothing else marks an order, attention and the per-token MLP treat all tokens alike. Shuffle the input points, and the output velocities come out shuffled the same way. This property is called permutation equivariance, and it suits point clouds, which have no order (see permutation invariance).

A condition can enter in two ways. A short one, such as a class label, is added to c. A long one, such as the known points of a shape to complete or a sentence, becomes its own set of tokens that the blocks read with cross-attention, or joins the data tokens in self-attention.

The name says diffusion, but the network fits flow matching too

The DiT paper (Peebles and Xie, 2023) trained the network as a diffusion model that predicts the added noise, on image patches in a compressed form of the image. Nothing in the network depends on that choice. To use it for flow matching, change only the training target, to the velocity x₁ − x₀.