DiTs Route Semantics via Template Tokens, Cut 20% FLOPs
TL;DR
- Structural template tokens act as implicit semantic registers in text-to-image DiTs, causally maintaining object identity during denoising.
- Prompt semantics are first injected into image latents and then read back into template tokens rather than passed directly between text tokens.
- A training-free pruning rule removes 20% of attention FLOPs with only a 1.4-point drop on GenEval by cutting heads that attend most to prompt tokens.
A new arXiv paper works out where text-to-image diffusion transformers actually keep their semantics during denoising, and the answer flips a common assumption. The prompt tokens you type are not the tokens that carry your prompt through generation. The structural template tokens around the prompt do most of that work.
The paper by Maohua Li and colleagues introduces a causal interpretability framework combining attention decomposition with targeted interventions across token spans, heads, and layers. The authors report that structural tokens carry little prompt-specific information at the encoder output, yet 'emerge as dominant image-to-text attention sinks and causally maintain object identity' inside the DiT, acting as implicit semantic registers. They argue this identity is acquired indirectly, with prompt semantics first injected into the image latents and then read back into the template tokens rather than transferred directly from the prompt tokens.
The practical hook is a training-free pruning rule. Heads that attend most strongly to prompt tokens turn out to be doing less than you would guess, and the authors report that 'pruning them removes 20% of attention FLOPs with only a 1.4-point drop on GenEval.' That is the kind of number that gets forwarded around inference teams, because DiT serving cost scales with attention compute and shaving a fifth of it off cheaply is meaningful for anyone hosting text-to-image models at scale.
The honest caveats: the abstract does not spell out which specific DiT checkpoints were tested, whether the 1.4-point drop concentrates on a particular GenEval sub-skill like counting or spatial relations, or how the pattern holds at higher resolutions and longer prompts. Take the FLOP number as evidence the lever exists, not a promise for every workload.
If the paper's framing that 'the tokens encoding semantics at input need not be those that maintain it during generation' travels, the win is twofold. Compression teams get a reusable, retraining-free recipe, and interpretability researchers get a causal view of how DiTs separate semantic routing from visual synthesis across heads and depth.
Originally reported by paper
Read the original article →Original headline: DiT Interpretability: Template Tokens Are Secret Semantic Registers—20% FLOP Pruning for 1.4pt Cost