One-sentence Summary

A causal interpretability framework combining attention decomposition and targeted interventions, developed by researchers from Nanjing University, Alibaba Group, and Zhejiang University, reveals that text template tokens in diffusion transformers act as implicit semantic registers that acquire identity indirectly from image latents, and derives a training-free pruning rule that removes $20\%$20% of attention FLOPs with only a $1.4$1.4-point drop on GenEval.

Key Contributions

• A causal interpretability framework for diffusion transformers combines attention decomposition with targeted interventions across token spans, heads, and layers to trace text-image token interactions during denoising.
• Chat-template structural tokens carry little prompt-specific information at the encoder output yet serve as implicit semantic registers. They causally maintain object identity by receiving prompt semantics via image latents and routing that identity through image-to-text attention.
• A training-free head-pruning rule removes 20% of attention FLOPs with only a 1.4-point GenEval drop by pruning heads that attend most strongly to prompt tokens. The framework also reveals that generative computation separates semantic routing from visual synthesis across heads and depth, progressing from identity formation to propagation and refinement.

Introduction

Modern text-to-image diffusion transformers increasingly rely on large language model encoders with chat templates that introduce structural tokens like system and user delimiters alongside a user’s prompt. These tokens are normally treated as formatting residue and ignored. At the same time, mechanistic understanding of diffusion transformers lags far behind that of language models: it was unclear which layers or attention heads actually route and store conditioning information during denoising. The authors introduce a causal interpretability framework that uses token-level attention decomposition, span interventions, head transplantation, and layer-wise masking to trace information flow. Applying this framework to chat-templated diffusion transformers, they find that structural tokens form dominant text-side attention sinks and act as semantic registers that extract object identity from evolving image latents rather than directly from prompt tokens. This finding leads to a training-free head-pruning rule that removes the causally inert prompt-reading heads, cutting 20% of attention FLOPs with only a small accuracy drop.

Method

The authors study the Qwen-Image family, a dual-stream multimodal diffusion transformer (MMDiT) that synthesizes images by iteratively denoising latent tokens under a flow-matching objective. Let $z \in \mathbb{R}^{n_{img} \times d}$z∈Rnimg​×d be the sequence of image latent tokens and $c \in \mathbb{R}^{n_{cond} \times d}$c∈Rncond​×d denote the text conditioning. A transformer with multiple blocks predicts the flow, which a sampler integrates from pure noise to a clean latent decoded to pixels by a VAE. Unlike models relying on CLIP or T5 encoders, Qwen-Image obtains its conditioning from a large vision-language model applied to a chat-formatted prompt. The user prompt is wrapped in a chat template, and the retained conditioning sequence is read from the encoder's last-layer hidden states after discarding the fixed template prefix. This sequence splits into two disjoint parts: a semantic span holding the prompt-content tokens and a structural span holding the trailing chat-template tokens that delimit the dialogue format.

Within each MMDiT block, text and image tokens are processed by a single self-attention over the concatenation of the conditioning and latent tokens. Partitioning the rows and columns of the joint attention matrix into their text and image segments yields four distinct blocks.

These blocks represent text queries attending to text and image keys, and image queries attending to text and image keys. The image-to-text block serves as the cross-modal pathway through which the text conditioning acts on the image latents. To quantify how much attention each token receives, the authors work with the full post-softmax joint attention matrix and define the attention mass on a key span from a query group.

The authors introduce a causal interpretability framework combining token-level attention decomposition with span-level conditioning interventions to trace where conditioning information is read, routed, and stored during denoising. They discover that the structural chat-template tokens act as dominant attention sinks. Across the denoising trajectory, image queries place an order of magnitude more attention on the content-free template span than on the prompt content.

The attention surface forms a sharp ridge over the structural span while attention among image tokens stays low. Despite carrying little explicit prompt-specific semantics, these template tokens function as implicit semantic registers inside the diffusion transformer. The authors trace the flow of information and find it is implicit: semantics pass first from the semantic tokens into the image tokens, and only then from the image tokens into the template tokens. Masking the attention of the register queries to the image latents collapses the object within the first few blocks, confirming that the register assembles its content by attending to the image stream early in the network.

Building on this analysis, the authors derive a training-free acceleration rule. Because heads with high image-to-semantic attention are causally inert for transferring object identity, they rank all attention heads by this metric and prune the computations of the top-ranked heads. This pruning is applied only over the late denoising steps, as early steps are more sensitive and commit the object's identity. Since every component of the joint-attention block scales linearly with the number of heads, this approach removes a proportional fraction of the computational cost while well preserving object correctness.

Experiment

Cross-prompt swapping and averaging experiments demonstrate that text template tokens carry only weak semantics, with their removal or replacement causing the model to collapse to a generic human portrait only when guidance is low, revealing the portrait as the unconditional default. Leveraging this insight, a training-free head pruning approach ranks the 1,440 joint-attention heads by their image-to-semantic attention scores and prunes the top-ranked heads over the last 80% of denoising steps. This method removes up to 20% of joint-attention FLOPs with only a small drop in GenEval accuracy and confirms that selecting which heads to prune, based on their semantic register role, is crucial for preserving both object correctness and perceptual quality.

Short prompts see trailing structural template tokens act as dominant attention sinks, absorbing far more image-query attention than semantic content tokens. As prompt length grows, semantic tokens gain attention share while structural attention declines, yet per structural token attention remains several times higher than per semantic token across all benchmarks. The structural tokens carry only weak semantics, and their removal causes generation to revert to the unconditional default. For short prompts (mean 7.9 tokens), structural tokens receive a mean attention mass of 0.19 versus 0.04 for semantic tokens, and structural mass exceeds semantic mass at 89% of attention sites. For long prompts (mean 82.1 tokens), semantic mass rises to 0.17, structural mass falls to 0.08, and the fraction of sites where structural mass dominates drops to 12%. The ratio of per-token attention mass remains consistently high (6.2–7.5) across all benchmarks, so each structural token individually attracts more attention than each semantic token even for long prompts. The structural template tokens encode only weak semantic content, and when they are removed completely, generation collapses to the model's unconditional default outcome.

Training-free head pruning based on image-to-semantic attention m_S preserves object correctness while reducing joint-attention FLOPs. Pruning a quarter of all heads (K=360) cuts 20% of FLOPs with only a 1.4-point GenEval drop, but perceptual quality declines faster, so a lighter budget of K=216 offers the best trade-off. The choice of heads is critical: the proposed ranking substantially outperforms both structural-sink and random head selection. At K=360, 20% of joint-attention FLOPs are removed while GenEval accuracy drops from 76.1 to 74.7. Perceptual quality (LPIPS) degrades more quickly than object correctness, making K=216 the optimal balance point. Under the proposed ranking, K=288 achieves GenEval 75.5; switching to structural-sink heads drops to 69.6, and random selection collapses to 51.3. HPSv3 preference score remains high (9.47) at K=216 but falls to 8.76 at K=360 and 8.44 at K=400. Pruning is applied only over the final 80% of denoising steps, leaving early identity-committing steps untouched. The method is prompt-independent and requires no additional training or online optimization.

The first experiment analyzes attention distributions in text-to-image generation, revealing that structural template tokens serve as dominant attention sinks for short prompts, while semantic tokens gain attention share with longer prompts; however, per-token structural attention remains consistently higher, and removing these tokens causes generation to revert to an unconditional default, confirming their weak semantic content. The second experiment validates a training-free head pruning approach that ranks attention heads by image-to-semantic attention, achieving a 20% reduction in joint-attention FLOPs with minimal loss in object correctness, though perceptual quality declines faster, and the proposed ranking substantially outperforms structural-sink or random head selection. Overall, the studies highlight the attention-sink behavior of structural tokens and demonstrate that informed head pruning can efficiently reduce computation while largely preserving generation quality.