Learning to Read the Contextual Tokens in Diffusion Transformers

AI summary

Overview

• Research area: Computer vision / generative modeling — specifically the internal representations of Multimodal Diffusion Transformers (MM-DiTs) used for text-to-image generation.
• Technical level: Advanced. The paper assumes familiarity with diffusion models, transformer attention, classifier-free guidance, representation alignment (REPA-family losses), and LLM-style vision-language interfaces such as LLaVA and Q-Formers.
• Scope: The paper introduces a method for decoding the "contextual tokens" of MM-DiTs into natural language, analyzes what those tokens encode over the denoising trajectory, and turns that finding into a training objective called Contextual Alignment.

What This Paper Is About

In MM-DiTs, text tokens are not fixed conditioning vectors: they are repeatedly updated through multimodal attention and become what prior work (Dahary et al., 2026) calls contextual tokens. Unlike visual tokens, these contextual tokens have no spatial correspondence and receive no dedicated supervision, so it has been unclear what visual information they actually carry. This paper builds a "Contextual Reader" that lets a frozen LLM answer plain-language questions about the emerging image using only the contextual tokens, and then uses what it learns to improve generation via a new training technique.

Key Contributions

1. A framework for reading the contextual space in natural language. A lightweight bottleneck network (a 2-layer Q-Former with 16 learned queries) compresses contextual tokens from all MM-DiT layers into a "Contextual Descriptor," which is fed as 16 continuous prefix embeddings to a frozen Qwen2.5-7B-Instruct LLM alongside a visual question. Only the bottleneck is trained; the MM-DiT and LLM stay frozen.

2. An empirical characterization of what contextual tokens encode. The reader shows that contextual tokens contain a rich, global representation of the emerging scene, that generation-specific semantics (including attributes the prompt leaves underspecified) are decodable surprisingly early in denoising, and that this holds even when the MM-DiT receives an empty prompt.

3. A link between readability and generation quality. Generations whose contextual representations are more readable tend to receive higher human-preference (HPS) scores, measured as a consistently positive ΔHPS between top- and bottom-ranked seeds per prompt.

4. Contextual Alignment (CoAl), a training technique. A "Aligner Network" (a single-layer Q-Former with one learned query) maps contextual tokens from a single layer to a predicted embedding that is matched by negative cosine similarity to the pooled SigLIP image embedding of the clean image. This is combined with the diffusion loss and, in full model training, complements REPA, HASTE and SRA; in fine-tuning it improves FID and Recall where those baselines do not.

Main Findings

• Readability emerges early and grows over denoising. VLM-as-a-judge scoring across the trajectory shows contextual readability increases as generation progresses. With the input prompt it is already high from the start; with an empty prompt it starts substantially lower and rises rapidly.

• Readability peaks before the end of denoising. The highest readability value is reached at 80% of the trajectory, with a small decrease at the final 100% timestep, which the authors attribute to a shift toward refining high-frequency details rather than the semantic properties probed by the questions.

• Image-specific semantics survive without any prompt. In the unconditional setting, readability eventually surpasses the conditional reader for FLUX.2 and approaches conditional performance for SD3.5. Figure 4 shows that with no prompt, the reader identifies that the person is skiing by 8% of the trajectory and that they are wearing goggles by 20%, before those details are clearly observable in the visual prediction.

• Seed-specific, prompt-unspecified details are recoverable. For the prompt "A person holding a scarf," the reader reports "green bushes and trees" for one seed and "a plain white wall" for the other. For "A person looking into a mirror," it identifies a smartphone in one seed and answers "Unknown" in the other, where the hands fall outside the image.

• Two architectures behave similarly. The analysis covers SD3.5 (separate image and text streams exchanging information through multimodal attention) and FLUX.2 (which additionally includes single-stream blocks jointly processing concatenated image and text tokens).

• Higher readability tracks higher HPS. Ranking the five generations per prompt by average reader score, ΔHPS — the HPS of the top-ranked generation minus that of the bottom-ranked — is consistently positive across architectures and conditioning settings. Qualitative examples show bottom-ranked samples with artifacts (a malformed toothbrush or tennis racket) or prompt misinterpretation (riding rather than walking behind a horse; two women on a couch rather than one).

• CoAl improves full-model training over all three baselines. Training a dual-stream MM-DiT from random initialization on MS-COCO 2014 (Table 1a): REPA goes from FID 4.95 to 4.29, CLIP ×10² from 24.20 to 24.40, HPS ×10² from 21.38 to 21.60, Precision from 0.591 to 0.628, Recall from 0.338 to 0.407. HASTE goes from 4.87 to 4.44 FID, 24.21 to 24.32 CLIP, 21.42 to 21.56 HPS, 0.609 to 0.623 Precision, 0.361 to 0.400 Recall. SRA goes from 5.89 to 4.36 FID, 23.49 to 24.28 CLIP, 20.37 to 21.42 HPS, 0.515 to 0.616 Precision, 0.261 to 0.398 Recall.

• CoAl works better than alignment baselines in fine-tuning. Fine-tuning SD3 on a curated subset of Fine-T2I (Table 1b): the base model scores FID 19.30, Precision 0.993, Recall 0.939; vanilla flow matching reaches FID 17.18, Precision 0.993, Recall 0.966; REPA 17.45 / 0.993 / 0.966; HASTE 17.53 / 0.994 / 0.963; SRA 17.48 / 0.991 / 0.964; CoAl reaches the best FID of 16.75 and best Recall of 0.969 at Precision 0.993. REPA, HASTE and SRA all produce FID worse than vanilla fine-tuning in this regime.

• Null-condition weighting matters. Because the prompt provides an easy shortcut for the alignment loss, the authors upweight the loss on null-condition (classifier-free-guidance dropout) iterations, forcing alignment to flow from the image tokens into the contextual tokens through interaction. This upweighting is applied only in the later stages of full model training, or throughout fine-tuning.

• Design choices were validated empirically. The choice of pooled SigLIP image embedding as teacher and the choice of aligning a single early layer are both confirmed in Appendix B.2.1 (the appendix content is not included in the provided text).

Methodology in Plain English

The authors treat "what does the model already know about the image it is drawing?" as a visual question answering problem.

1. Capture the text side. They run a frozen MM-DiT and, at a chosen denoising timestep, pull out the text-token representations produced by the attention operator in every transformer block, before the block's later processing steps. Only text positions are kept; image tokens are excluded. Each token sequence is tagged with a learned layer embedding and a timestep embedding (a 256-dimensional sinusoidal encoding passed through a 256→256→D MLP with SiLU), so the reader knows where and when each representation came from.

2. Compress into a descriptor. A small Q-Former — two blocks, eight heads, 16 learned queries, hidden dimension matching the diffusion backbone, feed-forward D→2048→D with GELU, cross-attention applied first — compresses the concatenated multi-layer tokens into 16 query vectors. A learned affine projection maps each into the frozen LLM's embedding space, L2-normalized and rescaled to the mean norm of the LLM vocabulary embeddings.

3. Ask the LLM. The descriptor is prepended as 16 continuous prefix embeddings to a question and answer-format instruction for Qwen2.5-7B-Instruct. The original generation prompt is deliberately withheld. Only the bottleneck is trained, using autoregressive cross-entropy over the answer tokens for 5,000 iterations at batch size 256.

4. Build supervision. From the MS-COCO 2014 training split, they keep images with exactly one non-crowd person, use Qwen2.5-7B-Instruct to filter captions to those describing exactly one unambiguous real person (62,069 eligible captions), then take 2,500 prompts from distinct source images — 500 held out, 2,000 for training, with no source-image overlap. Each prompt is rendered with five seeds (no seed shared between prompts), producing 10,000 training and 2,500 held-out generations per backbone at 1024×1024. SD3.5 uses 28 FlowMatch Euler steps at guidance scale 7.0; FLUX.2 uses 50 steps at guidance scale 4.0. InternVL3-8B answers 15 hand-curated questions per generated image, seeing only the image and the question, never the prompt. Because prompt, seed and initial noise are stored, the contextual representations at any timestep can be reconstructed during training.

5. Turn the finding into a loss. For Contextual Alignment, a frozen semantic encoder produces a pooled SigLIP embedding of the clean image. A minimal Aligner Network — one Q-Former layer, one learned query — maps contextual tokens from a single early layer to a predicted embedding, and a negative cosine similarity loss pulls the two together alongside the flow-matching objective. The MM-DiT and Aligner Network are optimized jointly, and the Aligner Network exists only during training, so inference cost is unchanged. They evaluate in two regimes: training a dual-stream MM-DiT from random initialization on MS-COCO 2014, and fine-tuning pretrained SD3 on a curated Fine-T2I subset, against REPA, HASTE and SRA. Full training is scored on FID, CLIP, HPS, and DINOv2 Precision and Recall; fine-tuning omits CLIP and HPS (the authors note the ground-truth photographs score below the base model on those metrics, making the comparison ill-posed) and uses FID, KID, and DINOv2 Precision and Recall instead.

Why This Matters

Impact on research. The paper reframes the text stream of MM-DiTs from a passive conditioning signal into an interpretable, supervision-worthy representation. It shows that global visual semantics are implicitly encoded in contextual tokens, complementing the local, spatial representations targeted by existing alignment methods such as REPA, and it links an internal representational property — readability — to an external quality signal (human preference). It also adds a new lens to the small but growing literature on causal interventions in contextual tokens (ConceptAttention, Contextual Repulsion, Padding Tone, Vision-Language Binding).

Real-world applications:

• Quality triage and seed selection. Because higher readability correlates with higher HPS, a reader-style score could be used to rank candidate generations before showing them to a user, without a human in the loop.
• Prompt underspecification diagnosis. The reader exposes which attributes a prompt leaves open and how the model resolves them, which is useful for debugging prompt behavior in creative tools.
• Fine-tuning to a target domain. CoAl's strongest reported result is in the fine-tuning regime, which is the realistic workflow for adapting a pretrained generator to a specific dataset (Table 1b: FID 16.75 versus 17.18 for vanilla and 17.45–17.53 for REPA/HASTE/SRA).
• Model inspection for safety and auditing. Being able to ask "what is in the image" before the image is perceptually settled gives a window into a model's intent early in the generation process.

Industry relevance. The technique requires no changes to inference: the Aligner Network is discarded after training and CoAl adds no inference cost. That makes it plausible to bolt onto existing MM-DiT training pipelines alongside or instead of existing alignment losses. The fine-tuning finding — that REPA-family methods give no meaningful gains when adapting an already-trained model, while contextual supervision still does — is directly relevant to teams that adapt off-the-shelf generators rather than pretraining from scratch.

Future Directions

• How do these representations form and evolve? The paper's own conclusion states that how internal representations form, evolve, and ultimately shape the generated image remains largely unexplored.
• Why does meaningful information appear before it is perceptually apparent? The fact that semantics are decodable early — even before they are visible in the visual prediction, and even with an empty prompt — raises open questions about the generation process itself.
• Extending reading and supervision beyond single-subject, human-centric images. The reader's dataset and its 15 questions were deliberately restricted to prompts describing exactly one unambiguous person, so generalization to arbitrary scenes, multi-object compositions, and other content is untested.
• Broadening beyond the two studied backbones. The analysis covers SD3.5 and FLUX.2; whether the same readability dynamics and the same benefit from Contextual Alignment hold for other MM-DiT families — including the joint text–spatiotemporal video models the related work discusses — is not reported here.

Target Audience

Researchers and engineers working on diffusion transformers, text-to-image generation, and interpretability of generative models. It will be most valuable to readers already comfortable with MM-DiT architecture, multimodal attention, classifier-free guidance, and representation alignment methods, and to practitioners who adapt pretrained generators to custom datasets and want a training signal that still helps after large-scale pretraining. Readers looking for an introduction to diffusion models will find the paper assumes substantial background.

Authors’ abstract

Multimodal Diffusion Transformers (MM-DiTs) jointly process visual and textual representations throughout generation. These models repeatedly update the text tokens through multimodal attention, forming dynamic contextual tokens whose function is not well understood. In this work, we introduce a framework for reading this contextual space through natural-language interrogation. We train a lightweight bottleneck network that maps intermediate contextual tokens into the input space of a frozen Large Language Model (LLM), allowing the LLM to answer questions about the emerging image directly from these hidden representations. Our reader reveals that contextual tokens encode a rich, global representation of the emerging scene: generation-specific semantics, including attributes left underspecified by the prompt, are accessible surprisingly early in denoising, while increasingly fine-grained details become readable over time. Remarkably, this information remains decodable even when the MM-DiT receives an empty prompt, showing that contextual tokens accumulate substantial image-specific information from the evolving visual representation itself. We further find that generations with more readable contextual representations tend to receive higher human-preference scores. Building on these observations, we introduce Contextual Alignment, a training technique that explicitly reinforces the visual-semantic information encoded in the contextual tokens, improving generation quality and distributional coverage. Together, our results establish contextual tokens as both an interpretable view into the internal dynamics of MM-DiTs and an effective target for improving generative models.