Text Template Tokens Are Implicit Semantic Registers in Diffusion Transformers
Maohua Li, Qirui Li, Yanke Zhou, Yiduo Li, Zhaosheng Chi, Chao Xu, Cuifeng Shen, Yixuan Xu, Hanlin Tang, Kan Liu, Tao Lan, Lin Qu, Shao-Qun Zhang
cs.CV
2026-07-21
In text-to-image DiTs, structural template tokens carry almost no prompt semantics at the encoder output yet become the dominant image-to-text attention sinks and implicit registers of object identity; pruning the heads that attend most to the prompt cuts 20% of attention FLOPs for a 1.4-point GenEval drop.
Text-to-image diffusion Transformers (DiTs) process text and image tokens together, yet how they compute internally during denoising stays opaque. This paper builds a causal interpretability framework: decompose the attention matrix into text-to-text, text-to-image, image-to-text, and image-to-image blocks, then run targeted interventions across token spans, heads, and layers (span swaps, cross-trajectory head transplants, layer-wise causal masking). The main subject is the Qwen-Image family, with extensions to FLUX.2, Krea-2-Turbo, and Qwen-Image-Edit.
The core finding is the two-faced nature of structural template tokens (chat delimiters, tokens like an end-of-turn marker). At the encoder output they carry almost no prompt-specific semantics: cross-prompt span swaps leave generation nearly unchanged, and a single content-free template shared across all prompts suffices to render any object. Inside the DiT, the same template tokens become the dominant image-to-text attention sinks, absorbing 76-92% of image-to-text attention, and they causally maintain object identity as implicit semantic registers.
The mechanism is indirect: semantics first flow from the prompt tokens into the image latents, completing within the first Transformer block, and the template tokens read identity back from the image stream rather than taking it directly from the prompt. The evidence: blocking template-to-prompt attention leaves the object almost unchanged, while blocking template-to-image attention collapses the object within the first two blocks.
This yields a training-free pruning rule: heads that attend most strongly to the prompt are dispensable. Rank heads by descending image-to-text semantic attention, prune the top K, applied over the last 80% of denoising steps. On Qwen-Image-2512, pruning 360 heads (25% of 1,440) cuts 20% of joint-attention FLOPs and drops GenEval from 76.1 to 74.7, just 1.4 points; a better setting of K=216 cuts 12% of FLOPs for 0.3 points. Pruning by structural-sink ranking instead drops GenEval to 69.6, far worse, confirming the right heads were removed.
For DiT builders the most direct payoff is the training-free pruning rule, cutting a fifth of attention FLOPs with almost no quality loss. The deeper point overturns an intuition: the tokens that encode semantics at the input need not be the ones that maintain it during generation. Prompt semantics travel an indirect prompt-to-image-to-template path, and template tokens are the real carriers of object identity during generation. The split between register heads handling identity and rendering heads handling pixels, with depth organizing into early identity commitment, middle propagation, and late refinement, maps strikingly onto LLM inference stages.
The authors concede the paper is primarily analytical, that training-free pruning is only the most naive application of the analysis, that sink-aware sparse attention could yield larger gains, and that why these content-free tokens become registers in the first place remains open. One gap: the conclusions rest mainly on the Qwen-Image family, with FLUX.2 and Krea-2 as extensions, so whether the register mechanism holds and prunes equally across all DiT architectures is not yet established.