Diffusion Transformers (DiT / Flux / SD3)

Four patches, flattened into a sequence -- exactly like four words. Every row is a query patch attending over every key patch, the same scaled dot-product softmax a text transformer runs, just fed patch content instead of word embeddings.

A Diffusion Transformer (DiT) starts exactly like a Vision Transformer: cut the (noisy) image into patches, flatten the grid into a sequence, and run ordinary self-attention over it. The heatmap above is the same scaled dot-product softmax from the attention chapters — nothing about it knows these tokens came from pixels instead of words.

The one new ingredient is how the model learns which noise level it's looking at. Instead of a per-resolution encoder path, DiT folds the timestep into every patch with a simplified adaptive layer norm (adaLN) — a scale and shift applied identically to every patch, before attention runs:

• $t$ — the diffusion timestep this forward pass is conditioned on; $T$ is the total number of timesteps.
• $c$ — the timestep normalized to $[0,1]$: the only input the conditioning below depends on.
• $\gamma(t), \beta(t)$ — the per-timestep scale and shift, shared identically across every patch at a given $t$.
• $v$ — a patch's raw content value; $\tilde v$ is that value after adaLN conditioning, the value attention actually sees.

1. t=0 is an ordinary, unconditioned ViT block

   At $t=0$, $c=0$, so $\gamma=1$ and $\beta=0$: $\tilde v = v$, exactly the plain patch embedding a normal Vision Transformer would use.

2. Rising t stretches and shifts every patch the same way

   As $t$ grows, $\gamma$ and $\beta$ move linearly away from the identity, scaling every patch's content by the same factor and shifting it by the same amount — the timestep never singles out one patch over another.

3. Conditioned content flows into ordinary self-attention

   Attention then runs on the conditioned values $\tilde v$ exactly as it would on any transformer's token embeddings — same queries, keys, values, same scaled dot product, same softmax.

Drag the timestep from $0$ to $4$. Because the conditioning applies the same scale and shift to every patch, it doesn't reorder which patch p0 attends to most — it just stretches the gap between p0's own (largest) content and everyone else's, so p0's self-attention weight on itself climbs as $t$ rises.

Four patches with raw content $[3, 1, 0, 2]$ for $p_0, p_1, p_2, p_3$, $T=4$, $w_\gamma=0.5$, $w_\beta=-1$:

1. Condition each patch at t=2

   $c = 2/4 = 0.5$, so $\gamma = 1+0.5(0.5) = 1.25$, $\beta = -1(0.5) = -0.5$. Conditioned values:

   • $\tilde v_0 = 1.25(3) - 0.5 = 3.25$
   • $\tilde v_1 = 1.25(1) - 0.5 = 0.75$
   • $\tilde v_2 = 1.25(0) - 0.5 = -0.5$
   • $\tilde v_3 = 1.25(2) - 0.5 = 2.0$
2. Score p0's query against every key

   Each patch's full embedding is $(\tilde v,\ 0.1\times\text{grid index})$:

   • $p_0=(3.25,0)$
   • $p_1=(0.75,0.1)$
   • $p_2=(-0.5,0.2)$
   • $p_3=(2.0,0.3)$

   Scoring query $p_0$ against each (dot product, then divide by $\sqrt2$):

   • $p_0\cdot p_0 = 3.25^2 = 10.5625$, scaled: $10.5625/\sqrt2\approx7.469$
   • $p_0\cdot p_1 = 3.25(0.75)+0(0.1) = 2.4375$, scaled: $\approx1.724$
   • $p_0\cdot p_2 = 3.25(-0.5)+0(0.2) = -1.625$, scaled: $\approx-1.149$
   • $p_0\cdot p_3 = 3.25(2)+0(0.3) = 6.5$, scaled: $\approx4.596$

   Softmax (subtracting the max, $7.469$, for stability):

   • self: $e^{0}=1$
   • $p_1$: $e^{1.724-7.469}\approx0.0032$
   • $p_2$: $e^{-1.149-7.469}\approx0.00018$
   • $p_3$: $e^{4.596-7.469}\approx0.0565$

   Sum $\approx1.0599$. Normalizing gives weights $\approx [0.943, 0.003, 0.0002, 0.053]$ — p0 attends overwhelmingly to itself, then to $p_3$, the next closest content value.

3. Attention output mixes values by those weights

   $p_0$'s attention output is $\approx 0.943(3.25) + 0.003(0.75) + 0.0002(-0.5) + 0.053(2) \approx 3.175$ — pulled almost entirely from its own conditioned content, with a small nudge from $p_3$.

4. Self-attention sharpens further as t keeps rising

   At $t=3$: $c=3/4=0.75$, $\gamma=1.375$, $\beta=-0.75$, so $\tilde v_0 = 1.375(3)-0.75 = 3.375$ (versus $3.25$ at $t=2$) — the gap between p0's content and the others widens further. Scaled scores against each patch:

   • $p_0$ (self): $3.375^2/\sqrt2\approx8.055$
   • $p_1$: $\approx1.492$
   • $p_2$: $\approx-1.790$
   • $p_3$: $\approx4.773$

   Running the same softmax gives a self-weight $\approx0.963$, past $0.95$ — while at $t=2$ it was still $\approx0.943$, below that threshold. The conditioning needs to stretch the gap a bit further before self-attention all but ignores every other patch.