PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2607.19139v2cs.CV

TL;DR

It remains unclear how DiTs incorporate and maintain prompt semantics when conditioning mixes prompt tokens with chat-template tokens. Using causal interventions, the paper finds that template tokens act as semantic registers through an indirect S → I → R circuit rather than repeated prompt reading.

  • Problem

    It remains unclear which conditioning tokens incorporate and maintain semantics during DiT generation, especially when prompts include chat-template tokens.

  • Method

    The paper applies causal interpretability interventions to separate prompt and template tokens and trace information flow through joint text-image attention.

  • Results

    Template tokens with little prompt-specific encoder information become dominant attention sinks and causally maintain object identity through the indirect S → I → R circuit.

  • Takeaways & Limitations

    DiT semantics need not remain attached to input-encoding tokens, revealing structured internal carriers and supporting training-free pruning of causally inert prompt-reading heads.

  • Takeaways & Limitations

    The study is primarily analytical, and its naive training-free pruning rule may leave larger efficiency gains unrealized.

Abstract

from arXiv · show

Modern text-to-image diffusion transformers (DiTs) generate images through joint attention, in which text and image tokens interact directly within a single sequence. In large-scale DiTs, the conditioning input contains not only the user prompt but also chat-template tokens introduced by LLM-based text encoders. Yet how these tokens participate in the denoising computation remains poorly understood. To probe this, we introduce a causal interpretability framework. Using it to separate prompt-content tokens from chat-template tokens, we find that the template tokens carry little prompt-specific information at the encoder output. Yet surprisingly, they emerge as dominant image-to-text attention sinks and causally maintain object identity inside the DiT, acting as implicit semantic registers. We show that they acquire this identity indirectly. Rather than reading the prompt tokens, they draw the identity from the image latents into which the prompt semantics have already been injected at the very first layer. We further reveal a division of labor across heads and depth in DiTs, where distinct heads route semantics or render visual structure, and identity is committed in early blocks, carried by middle blocks, and refined in late ones. As a practical payoff, this analysis yields a training-free pruning rule that removes the causally inert prompt-reading heads and cuts $20\%$ of joint-attention FLOPs at a $1.4$-point cost in GenEval accuracy. Overall, our work not only reveals that the tokens encoding semantics at the input need not be those that maintain them during generation, but also provides a causal view of internal mechanisms in diffusion transformers.

1 Introduction

The paper shows that semantic conditioning in diffusion transformers is indirect: prompt semantics enter image latents early, then chat-template tokens recover and maintain object identity as implicit semantic registers. A causal framework further reveals head- and depth-specific roles and enables training-free pruning of causally inert prompt-reading heads.

  • 1 Introduction: Template tokens contain little prompt-specific information initially but become dominant image-to-text attention sinks that maintain semantic conditioning inside the DiT.Image-token attention concentrates on template tokens rather than semantic prompt tokens, making them implicit semantic registers.
  • 1 Introduction: Prompt semantics enter the image stream in the first block, template tokens read that information from image latents in the second block, and deeper layers repeatedly retrieve it from them.The template tokens acquire semantics indirectly rather than by directly attending to the prompt.
  • 1 Introduction: Register heads preserve object identity, image-to-image heads render visual structure, and depth assigns identity commitment to early blocks, propagation to middle blocks, and refinement to late blocks.Prompt-reading heads are largely causally inert in this division of labor.
  • 1 Introduction: The proposed causal interpretability framework traces how textual information is injected, routed, stored, and transformed inside text-to-image DiTs.It uses token decomposition, span interventions, head transplantation, and layer-wise masking.
  • 1 Introduction: 20% of joint-attention FLOPs are removed by pruning causally inert prompt-reading heads, with a 1.4-point reduction in GenEval accuracy.The pruning rule is prompt-independent, training-free, and derived from the mechanistic account.

2 Related Work

Prior work identifies attention sinks as content-light tokens that absorb disproportionate attention and examines dedicated registers as buffers for attention and computation. These phenomena span autoregressive language models, multimodal models, and vision transformers.

  • The section reviews seminal studies on attention sinks and registers, with an extended discussion deferred to Appendix C.
  • Attention Sinks in Generative Models: Attention sinks absorb disproportionate attention mass despite carrying little semantic content, first identified in autoregressive LLMs as stable early-token anchors for streaming inference.Subsequent work links sinks to massive activations, depth-wise mixing control, and training dynamics, and finds they emerge during pre-training rather than initialization.
  • Registers as Attention and Computation Buffers: Dedicated content-free register slots provide a complementary way to control sink location and buffer attention or internal computation.In vision transformers, registers absorb high-norm artifacts and support cleaner feature or attention maps; language-model sink tokens stabilize streaming inference.

3 Preliminaries

The paper studies text-to-image diffusion transformers as iterative latent denoisers with joint text–image attention, focusing on Qwen-Image’s chat-formatted conditioning. It separates prompt-content and trailing template tokens and analyzes their influence through the image-to-text attention pathway.

  • Model setup: Diffusion transformers iteratively denoise image latent tokens under flow matching, while a VAE decodes the final clean latent into pixels.The transformer predicts a flow from image latents, timestep, and text conditioning; a sampler integrates this flow from noise.
  • Model setup: Qwen-Image uses a Qwen2.5-VL encoder on chat-formatted prompts within a dual-stream MMDiT whose text and image tokens interact through joint attention.Each modality has separate per-modality parameters inside blocks, but both modalities participate in one joint attention stream.
  • Conditioning structure: The retained conditioning sequence splits into a semantic span containing prompt-content tokens and a trailing template span containing fixed chat-format delimiters without prompt-specific content.The template span consists of tokens produced solely by chat-template conditioning and has no counterpart in CLIP/T5 text encoders.
  • Attention analysis: The analysis focuses on image queries attending to text keys in the I2T block, the cross-modal pathway through which conditioning acts on image latents.Attention mass is measured from image queries onto semantic and template spans using the post-softmax joint attention matrix; template-specific statistics include mR, mS, mtokS, and Pr[mR > mS].

4 Text Template Tokens as Implicit Semantic Registers

The causal analysis identifies chat-template tokens as implicit semantic registers: although they carry little prompt-specific information at the encoder output, they become dominant image-to-text attention sinks and maintain object identity inside the DiT. Semantics enter these registers indirectly from image latents, while specialized heads and depths divide identity transfer from visual rendering, enabling training-free pruning of causally inert heads.

  • 4.1 Template Tokens Are Dominant Attention Sinks: Template tokens absorb substantially more image-query attention than prompt-content tokens across simple, complex, and multilingual prompts.For GenEval, the five template tokens receive 0.19 of image-query attention versus 0.046 for content tokens, with a 6.4× per-token ratio; the ratio reaches 7.2× on DPG-Bench and 6.2× on Qwen-Image-Bench.
  • 4.2 Template Tokens Carry Little Semantics: A prompt-agnostic average template can replace genuine template representations with little change to the intended object, showing that template outputs carry weak prompt-specific semantics.Replacing the template span changes its representation by relative L2 0.36 and perturbs the first-step prediction by 1.8%, while median own-vs-average image similarity is 0.99 on Qwen-Image.
  • 4.2 Template Tokens Carry Little Semantics: Template-token substitutions can nevertheless cause object, count, or attribute failures, although classifier-free guidance largely suppresses these deviations and recovers the intended object.The failures often collapse to the unconditional default of a generic human portrait, and stripping the “a photo of” prefix leaves the regenerated output unchanged.
  • 4.3 Semantics Reside in Template Registers Implicitly: Inside the DiT, heads that read prompt-content tokens are poor predictors of causal identity transfer, whereas swapping heads that barely attend to them changes the generated object.In reverse transplantation order, approximately 18% of the 1,440 heads suffice to flip an apple into a banana, decoupling prompt reading from object carrying.
  • 4.3 Semantics Reside in Template Registers Implicitly: Register queries acquire object content from image latents rather than directly reading semantic prompt tokens.Cumulatively masking register attention to the semantic span over the first k blocks leaves the object essentially unchanged, while the analysis identifies image latents as the implicit source.
  • 4.4 Distinct Mechanisms at Different Heads and Different Depths: Semantic register heads preserve object identity across denoising, while rendering heads concentrate on image-to-image processing and construct spatially coherent visuals.Ablating even a small fraction of rendering heads collapses the output into incoherent patch mosaics, although some heads blend both roles.
  • 4.4 Distinct Mechanisms at Different Heads and Different Depths: Identity-transfer heads are distributed across early and late blocks, exposing a depth-dependent division of labor.Among approximately 270 heads, 89 lie in L1–10, 85 in L51–60, and 96 across L11–50; restricting swaps by group reveals distinct identity-processing roles.
  • 4.5 Semantic Registers Guide Training-Free Head Pruning: 20% joint-attention FLOPs can be removed by pruning 360 heads while GenEval accuracy drops only 1.4 points, from 76.1 to 74.7.The reported trade-off is head-selection dependent: perceptual quality declines more quickly, and K=216 gives the best trade-off.

5 Conclusions · Appendix

The conclusions revise the conventional view of text conditioning in diffusion transformers: prompt semantics need not remain attached to the tokens that initially encode them. The paper presents its interpretability framework as a systematic way to trace this internal information flow, with the appendix documenting the framework.

  • 5 Conclusions: Table 2 evaluates training-free late-step head pruning on Qwen-Image-2512 using GenEval.K denotes the number of silenced heads out of 1,440, and K=0 is the unpruned baseline.
  • 5 Conclusions: K=0 serves as the unpruned baseline in the late-step head-pruning evaluation.The table defines K as the number of silenced heads out of 1,440.
  • 5 Conclusions: Prompt semantics do not simply remain attached to the tokens that encode them during diffusion-transformer generation.This finding challenges the conventional picture of text conditioning.
  • 5 Conclusions: The model’s internal representation is more complex than current understanding suggests.The conclusion frames this complexity as a consequence of semantics moving beyond their original encoding tokens.
  • 5 Conclusions: The proposed interpretability framework systematically traces information during diffusion-transformer generation.The supplied conclusion introduces the framework as a way to analyze how models read, route, and transform information.
  • 5 Conclusions: The framework addresses an open question about how diffusion transformers read, route, and transform information during generation.The conclusion identifies this internal computation as remaining poorly understood despite rapid progress in the field.
  • Appendix: The appendix includes a causal interpretability framework for diffusion transformers.The supplied appendix passage identifies this framework as its topic without providing further procedural detail.
  • Appendix: Figure 6 provides an overview of the causal interpretability framework.The supplied figure passage identifies the figure as an overview.

A Limitations and Future Work … C.2 Attention Sinks in Sparse Attention

The paper frames its pruning rule as an initial analytical result, while leaving register formation and stronger sink-aware sparsity open. It also documents Qwen-Image’s chat-templated conditioning and situates attention sinks within the evolution of joint-attention diffusion transformers and sparse-attention research.

  • A Limitations and Future Work: The training-free head-pruning rule is a naive first cut, motivating sink-aware sparse attention that preserves register positions while pruning other keys.The authors suggest this design could yield substantially larger gains.
  • A Limitations and Future Work: Why content-free template tokens become semantic registers remains unresolved and is left for future investigation.The study is primarily analytical rather than a complete account of register formation.
  • B Qwen-Image Chat Templates and Retained Spans: The study examines the original Qwen-Image, Qwen-Image-2512, and instruction-based Qwen-Image-Edit-2511 checkpoints, whose templates differ across model families.The first two checkpoints share the same text-to-image chat template, whereas Qwen-Image-Edit-2511 uses a different template for image editing.
  • B Qwen-Image Chat Templates and Retained Spans: Qwen-Image’s template embeds a system description, user prompt, and assistant turn, with preprocessing retaining the prompt and selected chat markers before DiT input.The depicted template includes system and user messages followed by an assistant segment.
  • B Qwen-Image Chat Templates and Retained Spans: Qwen-Image-Edit-2511 adds image descriptions and visual placeholders, including one Picture i: block per input image before the edit instruction.The ellipsis in the template abbreviates repeated image-placeholder blocks rather than denoting a literal token.
  • C.1 Evolution of Diffusion Models: Diffusion modeling progressed from denoising and score-based formulations through DDIM, learned variances, guidance, and transformer-based image generators.The broader DiT line includes ViT-style tokenization, masked-token training, text-to-image scaling, and architectural refinements.
  • C.1 Evolution of Diffusion Models: Flow-matching and rectified-flow objectives increasingly accompany DiTs, which now use multimodal joint-attention streams where prompt, template, and image tokens compete directly.This joint stream makes attention-sink analysis relevant because conditioning is not merely side information injected through cross-attention.
  • C.2 Attention Sinks in Sparse Attention: Sink-aware sparse attention preserves high-mass sink positions while evicting low-mass tokens, supporting stable streaming or compressed KV caches.Recent systems also associate local or streaming heads with recent tokens and sinks, while retrieval heads retain long-range context.

D Experimental Setup · E Extended Experiments

The experiments use Qwen-Image-2512 under fixed computational settings, while the extended-experiments section adds results and experiments beyond the main text.

  • D Experimental Setup: D Experimental Setup: Qwen-Image-2512 is the default model for all experiments unless otherwise specified.
  • D Experimental Setup: D Experimental Setup: Experiments use a fixed random seed of 42.
  • D Experimental Setup: D Experimental Setup: Implementations use PyTorch 2.12.1 and DiffSynth-Studio v1.1.7.
  • D Experimental Setup: D Experimental Setup: Runs execute on NVIDIA H20 GPUs with 1,024GB of GPU memory.
  • D Experimental Setup: D Experimental Setup: The compute host has a 192-core CPU and runs CentOS 7.
  • E Extended Experiments: E Extended Experiments: This section reports additional results for experiments presented in the main text.
  • E Extended Experiments: E Extended Experiments: It also presents further experiments not included in the main text.

E.1 Extending the Analysis to FLUX.2

On FLUX.2, the full Mistral chat-template sequence is retained as conditioning, and template tokens again dominate image-to-text attention. Across 553 GenEval prompts, this dominance is nearly universal across attention sites, while progressive head swapping supports the same conclusion.

  • Conditioning setup: FLUX.2 uses a 32B rectified-flow DiT conditioned by a Mistral-3 24B vision–language model, with the complete serialized chat-template sequence retained.The template contains system and user roles but no assistant turn; unlike Qwen-Image, no fixed prefix is discarded.
  • Attention sinks: 92% of image-to-text attention lands on 33 retained template tokens, versus 8.1 mean content tokens, with a 3.3× per-token ratio.Aggregated over all 553 GenEval prompts, template tokens absorb m̄_R = 0.22 of image-query attention, compared with m̄_S = 0.016 for content tokens.
  • Attention sinks: At 99.8% of (t, l, h) sites, the template span outweighs the entire semantic span in image-to-text attention.This establishes near-total site-by-site dominance of the template span.
  • Causal validation: The progressive head-swap experiment on FLUX.2 likewise supports the paper’s conclusion.The experiment is reported in Fig. 9.

E.2 Extending the Analysis to Krea-2-Turbo

On Krea-2-Turbo, the chat-template attention sink recurs despite a different encoder and single-stream architecture. Its template tokens dominate image-to-text attention, and head-swap results remain consistent with the paper’s conclusion.

  • Model and setup: Krea-2-Turbo is a 12B single-stream flow-matching DiT conditioned by a 4B Qwen3-VL vision–language model, using the same template as Qwen-Image.The model and conditioning design are described as close to Qwen-Image.
  • Attention sink: 67% of all I2T attention lands on the template span, compared with 0.11 template-span mass versus 0.053 prompt-content mass.These values aggregate across blocks, heads, eight sampling steps, and 553 GenEval prompts; the per-token ratio is 3.1×.
  • Attention sink: The chat-template sink carries over to Krea-2-Turbo despite its different encoder and architecture.The dominant sinks are again all template tokens.
  • Causal validation: The progressive head-swap experiment on Krea-2-Turbo produces an outcome consistent with the paper’s conclusion.The experiment is reported in Fig. 10.

E.3 Extending the Analysis to Distilled and Edit Models … E.8 Per-Span Progressive Swap

Across distilled, editing, swapping, averaging, and head-intervention tests, template tokens consistently preserve or control image semantics more than their prompt-content role suggests. Template-span interventions alter images early, while semantic-span or semantic-reader interventions have weaker identity-transfer effects without image-token projections.

  • E.3 Extending the Analysis to Distilled and Edit Models: In a 2-step distilled LoRA, the analysis recomputes image-to-text attention statistics after collapsing the 50-step sampler.Few-step distillation substantially alters denoising dynamics and tests whether the template attention signature remains computationally stable.
  • E.3 Extending the Analysis to Distilled and Edit Models: 7.7×, 13.8×, and 194×: each template token attracts more attention than a vision token, instruction token, or reference-image latent, respectively, in the editing model.The template span exceeds the per-token mass of all three conditioning spans in 100% of edits, with <|im_end|> again the dominant sink.
  • E.4 Extra Cases for Cross-Prompt Swaps: Across every cross-prompt pair, swapping template tokens preserves the semantics of prompt A while changing at most unspecified, prompt-irrelevant content.The intervention keeps A’s content tokens and substitutes B’s template tokens at a fixed seed and sampler.
  • E.5 Extra Cases for Template-Token Averaging: For most prompts, replacing the template span with a prompt-agnostic average leaves the object unchanged, with large semantic changes confined to rare low-similarity cases.The averaged span is grafted onto each prompt’s own semantics as [S; ¯R].
  • E.6 Seed Sweep for Template-Token Averaging: Across prompts, the semantic perturbation caused by the shared average varies markedly across random seeds.The seed sweep regenerates [S; ¯R] to distinguish the average template’s semantic effect from ordinary sampling variation.
  • E.7 Extra Cases for Progressive Head Swap: Swapping top semantic readers first fails to transfer identity, whereas reversing the order flips the object after only a small fraction of heads.This result confirms that reading the semantic span and carrying the object are decoupled.
  • E.8 Per-Span Progressive Swap: Swapping template-span projections changes the image drastically after a fraction of blocks, while swapping semantic-span projections leaves it nearly unchanged until near-complete depth.The per-span swap transfers full q, k, v projections progressively across blocks in the 2-step distilled LoRA at one denoising step.
  • E.8 Per-Span Progressive Swap: Even a full-depth text-side swap fails to reproduce the target banana and instead pushes the image toward the unconditional default when image tokens remain untouched.This differs from head swapping, which also transplants image-token projections.

E.9 Qualitative Results of Head Pruning · E.10 Distinct Mechanisms at Different Heads and Depths

Qualitative pruning results show that object identity persists as high-mS heads are silenced, while perceptual quality degrades gradually. Head ablations and depth localization distinguish identity-carrying semantic-register heads from rendering heads and support an early-commit, middle-carry, late-refinement division.

  • E.9 Qualitative Results of Head Pruning: Object identity is preserved across increasing pruning budgets K, while perceptual quality degrades gradually as more high-mS heads are silenced.This qualitative behavior is shown in Fig. 18.
  • E.9 Qualitative Results of Head Pruning: Increasing the pruning budget produces a gradual decline in perceptual quality rather than immediate loss of object identity.The result concerns silencing progressively more high-mS heads.
  • E.10 Distinct Mechanisms at Different Heads and Depths: Silencing even a small fraction of I2I-concentrated heads collapses outputs into incoherent patches.Fig. 17 reports these qualitative rendering-head ablations.
  • E.10 Distinct Mechanisms at Different Heads and Depths: I2I-concentrated heads therefore exhibit a rendering role distinct from the identity-carrying role of semantic-register heads.The distinction follows from the contrasting effects of rendering-head ablation and identity preservation under pruning.
  • E.10 Distinct Mechanisms at Different Heads and Depths: 270 heads ranked by reverse semantic attention show bimodal concentration in early and late layers.Fig. 16 localizes these heads across depth.
  • E.10 Distinct Mechanisms at Different Heads and Depths: The early- and late-layer concentration accompanies an early-commit, middle-carry, and late-refinement division across depth.This division is described in the main text and linked to the localization in Fig. 16.

E.11 Temporal and Depth Structure of Attention Sinks

Template tokens increasingly dominate image-to-text attention during denoising, while their sink behavior is assembled across depth rather than inherited from conditioning embeddings. The mechanism emerges in early-to-middle blocks and aligns with depth-localized effects on object identity.

  • Temporal structure: 5.3× to 7.7×: Template per-token image-to-text attention exceeds semantic-token mass throughout denoising, with the advantage increasing toward t=0.On Qwen-Image, template mass falls from about 0.050 to 0.035, while semantic-token mass falls from 0.009 to about 0.004.
  • Depth structure: Template-token attention is near zero in the first block and reaches its later-block plateau through a learned mechanism assembled by early-to-middle blocks.This depth-localized origin suggests the sink is not inherited from the conditioning embeddings.
  • Depth structure: The identity-transfer heads concentrate in early and late layers, matching the layer groups that most strongly affect apple-to-banana identity transfer.The layer distribution is reported for the first 270 heads ordered by reverse semantic-attention score.
Loading 2607.19139v2…