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

Chat-template structural tokens in diffusion transformers act as implicit semantic registers that causally maintain object identity.

How do text-to-image diffusion transformers use chat-template tokens during denoising, and can we use this to interpret or prune the model?

Modern diffusion transformers treat chat-template structural tokens as inert formatting residue, yet these tokens compete for attention in the model's joint-attention stream. The authors show that these tokens function as dominant attention sinks and implicit semantic registers, acquiring object identity indirectly from image latents rather than directly from prompt tokens. This mechanism allows for a training-free pruning rule that removes 20% of attention FLOPs with only a 1.4-point drop in GenEval accuracy.

Paper Primer

In joint-attention diffusion transformers, all conditioning tokens—both the user's prompt and the chat template's structural delimiters—compete as keys for the image stream. The authors demonstrate that while structural tokens carry negligible prompt semantics at the encoder output, they emerge as the primary sites where the model routes and maintains object identity during denoising.

The method uses a causal interpretability framework to trace information flow: it combines token-level attention decomposition, span-level interventions, and cross-trajectory head transplantation. The core move is identifying that semantics are injected from prompt tokens into image latents first, and then read back into the structural "register" tokens, decoupling the tokens that encode the prompt from those that maintain it.

Structural template tokens act as dominant attention sinks and semantic registers.

Across benchmarks, structural tokens absorb 6–7× more attention per token than semantic prompt tokens, and transplanting these "register" heads flips object identity in the generated output. Structural tokens account for up to 78% of all image-to-text attention.

Prompt-reading heads are causally inert for object identity.

Heads that attend most strongly to prompt tokens fail to transfer object identity during head-transplantation experiments, whereas heads that attend to structural registers successfully flip the generated object. Pruning the top 25% of semantic-reading heads removes 20% of attention FLOPs with a minimal 1.4-point accuracy drop.

Why do these structural tokens become registers if they carry no prompt content?

The model learns to route semantics from the image latents into these content-free slots early in the denoising process, effectively using them as stable "registers" to hold object identity throughout the generation trajectory.

Does this register behavior depend on the specific model architecture?

No; the authors observe this sink behavior across multiple architectures, including dual-stream MMDiTs like Qwen-Image, single-stream flow-matching models like Krea-2-Turbo, and large-scale rectified-flow models like FLUX.2.

This paper shifts the understanding of diffusion transformers from "prompt-to-image" conditioning to a "register-mediated" mechanism, where the tokens that define the prompt at input are not the ones that maintain it during generation.

Researchers can now prune diffusion transformers by targeting "semantic-reading" heads that are causally inert, enabling significant compute savings without retraining.

Introduction: The Black Box of DiTs

We expose hidden semantic registers in diffusion transformers and prune redundant heads without retraining.

Modern text‑to‑image diffusion transformers (Diffusion Transformers, DiTs) jointly attend over text and image tokens, yet we lack a mechanistic view of how conditioning information flows during denoising. This gap is especially acute for the chat‑template tokens that appear only as formatting but may become hidden computational registers.

We need a causal lens on diffusion transformers because seemingly inert formatting tokens can silently store and route semantic content, and because heads that merely attend to prompt tokens may be computationally wasteful.

**Figure 1.** Overview of proposed causal interpretability framework.

The need for causal interpretability in large‑scale diffusion models is paramount: it reveals hidden semantic registers and guides training‑free efficiency gains.

Architecture and Attention Dynamics

Defines the MMDiT backbone, token partitions, and joint attention mass used for analysis.

This section grounds the core MMDiT building blocks, the token split used by the chat template, and the attention‑mass metric that underpins the later pruning analysis.

The model stacks transformer blocks that each process image latents and text conditioning together, but keep separate per‑modality parameters while sharing a single joint self‑attention.

How does MMDiT differ from a vanilla transformer that simply concatenates text and image tokens?

In a vanilla transformer the same linear projections are applied to all tokens, so there is no modality‑specific parameterization. MMDiT keeps separate projection matrices for text and image within each block, preserving distinct feature spaces before they meet in the shared attention matrix.

Self‑attention is computed over the full $[c; z]$ sequence, then the resulting matrix is sliced into four quadrants that describe how each modality queries the other.

Why does the paper concentrate on the I2T block rather than the other three quadrants?

The I2T block is the only cross‑modal pathway that lets image latents attend to text keys; the other blocks either stay within a single modality (T2T, I2I) or describe how text queries image keys (T2I), which does not affect the image generation process.

The image shows a curly brace pointing downwards towards the text "dropped prefix".

**Figure 2** Joint attention partition in MMDiT. Rows are queries and columns are keys, yielding T2T, T2I, I2T, and I2I regions over concatenated text and image tokens. We mainly focus on the I2T block to measure the attention mass received by text keys, especially structural chat-template tokens.

Chat-Template Tokens as Semantic Registers

We expose how chat‑template tokens act as implicit semantic registers and enable training‑free head pruning.

We now detail the causal interpretability framework that reveals how the trailing chat‑template tokens become the model’s implicit semantic registers.

We treat the model as a plumbing system: token‑level attention decomposition measures flow through each pipe, span‑level conditioning interventions block or reroute whole sections, cross‑trajectory head transplantation swaps individual valves, and layer‑wise causal masking closes specific junctions to see what downstream water (semantic content) remains.

How does this differ from a standard attention‑heatmap inspection?

A heatmap only shows where attention lands; the framework actively intervenes (swaps, masks, replaces spans) and measures the downstream image change, turning a passive observation into a causal experiment.

Compute I2T mass for each head; the top head assigns $m_R$ = 0.998 to the template span.

Per‑token mass: template tokens receive ≈ 0.20 each, content tokens ≈ 0.018 each → 11× per‑token gap.

Swap the template span with an average template \bar{R} and observe a DINO similarity of 0.93, indicating minimal semantic change.

The example shows that even with a tiny prompt the model’s attention overwhelmingly sinks into the template tokens, yet those tokens contribute little semantic content.

Across the denoising trajectory the model’s image queries pour the vast majority of their attention into the trailing template span, treating it as a massive sink.

Why does a single delimiter token dominate attention instead of the semantic nouns?

Because the delimiter sits at the boundary of the template span, its position aligns with the model’s positional encoding that the attention mechanism has learned to treat as a high‑frequency anchor, making it an easy attractor for query vectors.

The highest‑attention head places 99.8 % of its image‑to‑template mass on the structural span.

Measured on a single‑step run for the prompt “An apple” (both model checkpoints).

Per‑token attention on template tokens exceeds that on content tokens by up to 13 times.

Mean per‑token mass $\bar{m}_R / \bar{m}_S$ = 13× for the minimal prompt; 6–7× for longer prompts.

**Figure 7** Per-span progressive swap. The $(q, k, v)$ of one span—semantic $\mathcal{S}$ (top) or structural $\mathcal{R}$ (bottom)—is copied from $B$ into $A$ over the first $r$ blocks. Swapping $\mathcal{S}$ preserves the apple; swapping $\mathcal{R}$ disrupts it early.

**Figure 3.** 3D surfaces of full query–key attention maps for the prompt “An apple” on Qwen-Image and Qwen-Image-2512. Columns show Top-n heads ranked by $m_{\mathcal{R}}$, the image-to-structural-token attention mass. Attention concentrates as a sharp ridge over the structural span $\mathcal{R}$ while attention among image tokens stays low. The text and image token ranges are displayed with a nonuniform axis scale for readability, and orange markers denote the structural token span $\mathcal{R}$.

We gradually replace the (q, k, v) projections of selected heads in one generation run with those from another run, observing when the generated object identity changes.

Why doesn’t swapping the heads that attend most to the semantic tokens change the object?

Those heads mainly pass the semantic tokens’ signals forward without writing them into the internal registers; the registers are updated by heads that attend elsewhere, so swapping the high‑semantic heads leaves the stored object unchanged.

We block attention from the template registers $R$ to either the semantic span $S$ or the image latents $I$ in a specific block and measure the effect on the generated image.

Why does disabling $Rarrow I$ attention erase the object while disabling $Rarrow S$ does not?

$R$ gathers its content from the image latents; cutting that pathway removes the stored object representation. The semantic span merely provides a cue that the registers ignore.

Pruning the 360 heads with highest $\bar{m}_S$ cuts 20 % of joint‑attention FLOPs while degrading GenEval accuracy by only 1.4 points.

Table 2 shows GenEval = 74.7 (vs 76.1 baseline) after pruning 360 heads (≈ 20 % FLOP reduction).

**Table 2.** Training-free late-step head pruning on Qwen-Image-2512, evaluated on GenEval. $K$ is the number of silenced heads out of 1,440; $K=0$ is the unpruned baseline.

**Table 1.** Attention statistics across prompt benchmarks. $\overline{|S|}$ is the mean number of semantic tokens per prompt. $m_{\mathcal{R}}$ denotes image-to-text attention mass assigned to trailing structural template tokens, and $m_{\mathcal{S}}$ denotes mass assigned to semantic prompt tokens. Bars denote means and tildes denote medians, averaged over prompt-level statistics. $\overline{m_{\mathcal{R}}^{\text{tok}}} / \overline{m_{\mathcal{S}}^{\text{tok}}}$ denotes the average mass per structural token divided by the average mass per semantic token. $\text{Pr}[m_{\mathcal{R}} > m_{\mathcal{S}}]$ is the fraction of step-layer-head sites where structural mass exceeds semantic mass.

**Figure 4** Structural template tokens carry semantics, but only weakly. (a) Cross-prompt swap substitutes another prompt’s template tokens, $[S^A; R^B]$, with little effect on the object. (b) Average on structural tokens replaces $R$ with a single prompt-agnostic average $\bar{R}$, $[S; \bar{R}]$, which can occasionally cause large semantic changes. Increasing the guidance scale (CFG) largely resolves these deviations. (c) Recomputing the average $\bar{R}$ from the same prompts with the “a photo of” prefix stripped leaves the result essentially unchanged from (b).

**Figure 5.** Template tokens act as the implicit semantic registers in the denoising process. (a) Progressive head swap, ordering heads by their semantic-span attention $m_s$. Swapping high-to-low ($m_s\downarrow$, bottom) barely transfers the identity and even collapses to the unconditional portrait; the reverse order ($m_s\uparrow$, top) flips apple to banana after only ~18% of the heads. (b) Progressive attention masking over layers. Blocking $\mathcal{R}arrow\mathcal{S}$ attention leaves the apple largely intact, while blocking $\mathcal{R}arrow\mathcal{I}$ attention quickly erases the object identity.

Together these probes reveal a three‑stage causal flow: early heads commit object identity into the template registers, middle heads carry it forward unchanged, and late heads refine visual details.

Limitations and Pruning Results

Limits of the analysis and directions for extending head‑pruning.

The paper showed that chat‑template tokens act as implicit semantic registers and can be used to prune attention heads without retraining.

A head is removed if its attention mass on the structural token span is low, because those heads contribute little to the image‑generation signal.

How does this pruning differ from standard magnitude‑based weight pruning?

Standard magnitude pruning removes parameters with small absolute values, regardless of their functional role. Training‑free head pruning instead looks at a functional signal—how much a head attends to the structural register—so it discards heads that are behaviorally irrelevant even if their weights are large.

The analysis is purely analytical; the pruning rule is a naive first‑cut that may leave substantial head‑capacity unused.

**Figure.** Qualitative comparison of image generation results across different pruning budgets and control methods. The left section shows results using the proposed ranking method with increasing pruning budgets (K=0 to K=400), while the right section shows control experiments with structured and random pruning at K=288.

Future work could explore sink‑aware sparse attention that retains register positions as keys, and investigate why structural tokens become registers in the first place.

Training‑free head pruning offers a smooth trade‑off: modest budgets preserve object identity while higher budgets mainly affect perceptual fidelity.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers