Learning to Read the Contextual Tokens in Diffusion Transformers

1 INTRODUCTION

• Contextual tokens are repeatedly updated through multimodal attention, yet their visual content and influence on generation remain unclear.Unlike visual tokens, they lack explicit spatial correspondence and dedicated supervision, making them difficult to interpret.Source: §1 · p. 2
• A lightweight bottleneck projects intermediate contextual tokens into a frozen LLM, enabling plain-text questions about the emerging image to be answered from hidden representations.The reader exposes how the model resolves prompt attributes left underspecified during generation.Source: §1 · p. 2
• Generation-specific semantics are readable surprisingly early, even with an empty prompt, and readability increases as finer details emerge over denoising.Generations with clearer contextual representations tend to receive higher human-preference scores.Source: §1 · p. 2
• Contextual Alignment reinforces visual-semantic information in contextual tokens and improves generation quality and distributional coverage.The paper presents contextual representations as both an interpretable view of generation and a training target beyond the standard diffusion objective.Source: §1 · p. 2

2 RELATED WORK

• MM-DiTs jointly update text and visual representations through multimodal attention, a design now used in image and video generation models.Earlier latent-representation studies focused on visual features or semantic information distributed across text tokens.Source: §2 · p. 3
• Prior work uses contextual tokens to localize concepts or intervene on generation, but does not directly reveal their semantic content or evolution over denoising.The paper instead decodes image-specific semantics into natural language, including without textual conditioning and before they are readily visible.Source: §2 · p. 3
• Contextual Alignment aligns implicitly encoded global visual semantics in text-conditioned contextual tokens, complementing methods that reinforce local spatial representations.Prior alignment work explored spatial structure and separately represented global semantics in class-conditional generation.Source: §2 · pp. 3–4

3 READING THE CONTEXTUAL SPACE

• The Contextual Reader answers visual questions about an emerging image using contextual representations from a selected generation timestep.Questions can probe properties not explicitly specified by the prompt.Source: §3 · p. 4

3.1 READER ARCHITECTURE

• A Bottleneck Network aggregates contextual tokens from multiple MM-DiT layers into a compact Contextual Descriptor, which a frozen LLM uses with the visual question to generate an answer.The network compresses the layer representations through a Q-Former and projects them into the frozen LLM’s input embedding space.Source: §3.1 · pp. 4–5

3.2 TRAINING THE CONTEXTUAL READER

• The reader is trained on single-subject images using 15 questions, with five seeds per prompt and VLM-generated answers as supervision.Training samples contextual representations at arbitrary denoising timesteps and optimizes only the Bottleneck Network while the MM-DiT and LLM remain frozen.Source: §3.2 · p. 5

3.3 SEMANTIC READABILITY OF THE CONTEXTUAL SPACE

• Contextual tokens expose seed-specific attributes left unspecified by a shared prompt, sometimes before those differences are clear in the visual prediction.For example, the reader distinguishes different backgrounds for the same prompt and answers “Unknown” when a detail falls outside the image.Source: §3.3 · p. 6
• Across architectures, conditional readability increases through denoising and peaks at 80%, while empty-prompt readability rises rapidly and eventually surpasses conditional performance for FLUX.2.For SD3.5, empty-prompt readability approaches conditional performance; the final 100% conditional timestep shows a small decrease as generation refines high-frequency details.Source: §3.3 · p. 7
• Across architectures and conditioning settings, seeds ranked higher by reader scores have consistently positive ΔHPS compared with bottom-ranked seeds.The top- and bottom-ranked generations often differ in visual fidelity and prompt alignment.Source: §3.3 · p. 7

4.1 CONTEXTUAL ALIGNMENT

• CoAl aligns contextual tokens with a global semantic representation of the clean image, adding the alignment loss to diffusion training and jointly optimizing the MM-DiT and Aligner Network.The Aligner Network is used only during training, leaving inference cost unchanged.Source: §4 · p. 8, §4.1 · p. 8
• Null-Condition Weighting upweights alignment loss when text is dropped, encouraging semantic information to flow from image tokens into contextual tokens.It is applied in later full-training stages or throughout fine-tuning.Source: §4.1 · p. 8

4.2 EXPERIMENTAL SETUP AND BASELINES

• CoAl is evaluated both on a dual-stream MM-DiT trained from scratch on MS-COCO and on pretrained SD3 fine-tuned with Fine-T2I.Fine-tuning evaluation focuses on target-distribution fit and coverage using FID, KID, and DINOv2 Precision and Recall.Source: §4.2 · pp. 8–9

4.3 MAIN RESULTS

• Contextual Alignment improves all three alignment baselines across evaluated metrics in full training; in fine-tuning, it leads in FID and Recall and remains competitive in Precision.REPA, HASTE, and SRA each achieve worse FID than vanilla fine-tuning.Source: Table 1 · p. 9, §4.3 · p. 9
• Qualitative comparisons report improvements in image quality, sharpness, subject anatomy, compositional coherence, and prompt following.Additional comparisons are reported in Appendix B.2.4.Source: §4.3 · p. 9

5 CONCLUSIONS

• Contextual tokens can reveal final-image semantics before those details become perceptually apparent, including when the model receives no text prompt.The authors state that how these representations form, evolve, and shape the generated image remains largely unexplored.Source: §5 · pp. 9–10
• Generations with clearer internal representations tend to produce better images, and Contextual Alignment is reported to improve generations.Source: §5 · p. 10

A.1.1 CONTEXTUAL FEATURE EXTRACTION AND ENCODING

• Reader inputs are text-side attention outputs before later block updates, with layer and timestep embeddings added; image-token positions are excluded.The timestep embedding is added to every contextual token, while layer-token pairs remain separately accessible to the Q-Former.Source: §A.1.1 · pp. 15–16

A.1.2 BOTTLENECK NETWORK AND LANGUAGE-MODEL INTERFACE

• A two-block Q-Former maps contextual features into 16 projected prefix embeddings for frozen Qwen, which receives the question and answer-format instruction but not the original prompt.Training uses teacher-forced answer and EOS targets, with cross-entropy computed only over answer tokens.Source: §A.1.2 · p. 16

• Dataset construction starts from COCO images with exactly one non-crowd person and retains captions classified as describing exactly one unambiguous real person, yielding 62,069 eligible captions.Source: §A.1.3 · p. 16
• From distinct source images, the authors select 2,000 training prompts and 500 held-out prompts, rendering five seeds per prompt for 10,000 training and 2,500 held-out generations per backbone.No source image overlaps between the two sets.Source: §A.1.3 · p. 16
• Images are generated at 1024 × 1024 with backbone-specific denoising steps and guidance scales, and each image is labeled using 15 manually curated visual questions.InternVL3-8B answers from the image and question alone; invalid answers are retried up to twice before being replaced with “unknown.”Source: §A.1.3 · pp. 16–17

A.1.4 READER TRAINING

• Reader training freezes the diffusion backbone, text encoders, VAE, and LLM while optimizing the Q-Former and projection on contextual features sampled across denoising noise levels.Training uses autoregressive cross-entropy; validation loss falls steeply until about 2,000 steps, then decreases more modestly.Source: §A.1.4 · p. 17

A.1.5 VLM-AS-A-JUDGE DETAILS

• Prometheus-Vision 13B judges each reader prediction against its reference answer given the visual question and final generated image, scoring correctness, completeness, and image consistency on integer 1–5 scores normalized to [0,1].At each denoising timestep, readability is the mean normalized score across all validation images and questions.Source: §A.1.5 · pp. 17–18

A.2.1 FEATURE EXTRACTION AND ALIGNER NETWORK ARCHITECTURE

• For CoAl, a single-layer, single-query Q-Former and MLP map contextual features to an embedding trained against the teacher with negative cosine similarity loss.The contextual features are pooled after the shared attention block without positional embeddings.Source: §A.2.1 · p. 18

A.2.2 NULL-CONDITION WEIGHTING

• Null-Condition Weighting increases CoAl loss for dropped-text samples and reduces it for text-guided samples, with different weights for full training and fine-tuning.Source: §A.2.2 · p. 18

A.2.3 TEACHER EMBEDDINGS

• The default teacher is the pooled 768-dimensional embedding from SigLIP base; ablations also consider CLIP, DINOv2, and Qwen3-VL.Source: §A.2.3 · p. 19

A.2.4 TRAINING CONFIGURATIONS

• Full training uses a scratch SD3-based 24-layer model on MS-COCO at 256 × 256; fine-tuning updates all SD3-Medium weights on curated Fine-T2I at 1024 × 1024.Source: §A.2.4 · p. 19

A.2.5 EVALUATION

• FID, KID, CLIP similarity, DINO precision and recall, VQAScore, and HPS are the reported generation metrics.FID and KID use InceptionV3 pool3 features; the passage also specifies the packages or feature encoders used for the other metrics.Source: §A.2.5 · p. 19
• Full-model metrics use 40,192 generated samples, except VQAScore, which uses 4,096; fine-tuning metrics use the full 5,000-image held-out validation set.The full-model samples are evaluated against the MS-COCO 2014 validation set.Source: §A.2.5 · p. 19
• Prometheus-Vision scores reader predictions from 1 to 5 using the visual question, prediction, reference answer, and generated image.The figure describes the VLM-as-a-judge prompt and scoring rubric.Source: Figure 10 · p. 20
• Contextual readers provide higher readability during early denoising, while x0+VLM achieves the highest readability from approximately 40% onward.Figure 11 compares conditional and unconditional Contextual Readers, a Latent Reader, and x0+VLM across two diffusion backbones.Source: Figure 11 · p. 20

B.1.1 COMPARING SEMANTIC READABILITY ACROSS REPRESENTATIONS

• Before approximately 40% denoising, both conditional and unconditional Contextual Readers outperform x0+VLM, and contextual tokens outperform latent representations throughout or from 10% onward, respectively.x0+VLM has the highest readability from approximately 40% onward; it answers from the predicted clean image and is more closely aligned with the evaluation target.Source: §B.1.1 · p. 21
• The conditional Contextual Reader exceeds the Prompt Reader throughout denoising; the unconditional reader surpasses it for FLUX.2 and approaches it for SD3.5 as denoising progresses.Prompt Reader uses the same architecture, questions, and target answers but receives only the input prompt.Source: §B.1.1 · p. 21, Figure 12 · p. 21

B.1.2 GENERALIZATION TO UNSEEN QUESTIONS

• For held-out questions, readability decreases but remains meaningful: conditional readers reach approximately 0.50, while unconditional readers approach conditional performance as denoising progresses.Each reader is trained on 14 of 15 questions and evaluated on the remaining question; this is repeated across all 15 questions.Source: §B.1.2 · pp. 21–22, Figure 13 · p. 22
• On held-out questions, SD3.5 unconditional readability reaches approximately 0.40, while FLUX.2 eventually surpasses its conditional reader later in denoising.The reported SD3.5 responses are generally relevant but include some omissions.Source: §B.1.2 · p. 22

B.1.3 ADDITIONAL QUALITATIVE RESULTS

• Final images resolve into the activities identified by the reader even when early visual predictions are ambiguous.These final images correspond to the generations shown in Figure 1.Source: §B.1.3 · p. 22, Figure 14 · p. 22
• For two seeds sharing a prompt, the reader recovers differently realized, underspecified attributes from contextual tokens at early timesteps.Figure 15 shows predicted clean images and reader predictions at 8% of denoising, alongside the final images.Source: §B.1.3 · p. 22, Figure 15 · p. 23
• Across unconditional examples, reader predictions progress from coarse semantic properties to more fine-grained answers as denoising advances.The results are presented for both FLUX.2 and SD3.5.Source: §B.1.3 · p. 22

B.2.1 ABLATION STUDIES

• Across full training and fine-tuning, CoAl is fairly robust to teacher choice, though SigLIP consistently edges out CLIP, DINOv2, and Qwen3-VL-2B.These ablations examine teacher, contextual alignment layer, and null-condition weighting.Source: §B.2.1 · p. 23
• Aligning earlier contextual layers consistently outperforms aligning later ones, likely because more of the network can use the aligned representation as it propagates forward.The comparison is reported across both full model training and fine-tuning settings.Source: §B.2.1 · pp. 23–24
• Null-condition weighting matters especially in fine-tuning, where upweighting text-dropped iterations addresses the loss of the prompt-alignment shortcut.The pretrained model already aligns reasonably well with the semantic teacher when the prompt is present.Source: §B.2.1 · p. 24

B.2.2 EXTENDED QUANTITATIVE COMPARISONS

• The extended comparisons add VQAScore and KID, and REG has worse FID and KID than plain diffusion training in MM-DiT.REG supervises a global semantic representation but is designed for class-conditional architectures without a text branch.Source: §B.2.2 · p. 24
• Without text prompts, reader predictions on FLUX.2 and SD3.5 recover increasingly specific semantics at later denoising points.The figures show predicted clean images at early trajectory points alongside the final generations and corresponding reader answers.Source: Figure 17 · p. 25, Figure 18 · p. 25
• Preference and alignment metrics require caution in fine-tuning: Reference scores below Base Model, reflecting metric preferences for training-distribution conventions rather than greater image fidelity.The reported comparison includes VQAScore alongside CLIP and HPS.Source: §B.2.2 · pp. 25–26
• Across full training, REPA + CoAl is best in every reported column; in fine-tuning, all three combinations improve KID, with only marginal precision/recall regressions.The full-training comparison also finds that each representation-alignment baseline improves on most metrics when combined with CoAl.Source: Table 3 · p. 26, §B.2.2 · p. 26, Table 4 · p. 27

B.2.3 PERFORMANCE ACROSS CFG SCALES

• Across CFG scales 1–9, CoAl is strongest among four methods; CFG 2 gives every method its best FID and KID, while CFG 7 is retained for balance.The authors describe CFG 7 as a better balance between prompt adherence and distributional fidelity.Source: §B.2.3 · pp. 26–27, Figure 19 · p. 27

B.2.4 ADDITIONAL QUALITATIVE RESULTS

• Across the additional qualitative comparisons, CoAl improves image quality and sharpness, subject anatomy, text rendering, compositional coherence, and prompt following.The reported observations span Figures 20, 21, and 22.Source: §B.2.4 · p. 27
• Figures 20 and 21 compare held-out references with Vanilla fine-tuning, REPA, HASTE, SRA, and CoAl on additional prompts.Both figures follow the comparison format of Figure 7.Source: §B.2.4 · p. 27, Figure 20 · p. 28, Figure 21 · p. 29
• Figure 22 pairs each baseline against the same baseline with CoAl, with five examples per pair and a distinct prompt set for each pair.The shown baselines are Vanilla fine-tuning, REPA, and HASTE.Source: §B.2.4 · p. 27, Figure 22 · p. 30