学习读取扩散 Transformer 中的上下文 Token Learning to Read the Contextual Tokens in Diffusion Transformers
用轻量读取器将 MM-DiT 上下文 token 解码为自然语言，并提出上下文对齐训练以提升生成质量
前置知识
MM-DiT 与多模态注意力
MM-DiT（如 SD3、FLUX、Qwen-Image 的骨干）是扩散 Transformer 的一种架构：文本 token 与图像 token 各自维护独立的数据流，但在每个 Transformer 块中通过共享的多模态注意力相互关注、彼此更新，而非传统 U-Net 那样仅用单向 cross-attention 注入文本条件。SD3.5 采用双流结构，FLUX.2 还额外包含图文 token 拼接后联合处理的单流块。
本文研究对象正是文本数据流在多模态注意力中逐层演化出的中间表示，不理解双流联合注意力就无法定位“上下文 token”从何而来。
上下文 Token（contextual tokens）
在 MM-DiT 中，输入的文本嵌入（含 padding token）并不保持固定，而是在每一层被多模态注意力反复改写，持续吸收来自图像 token 的信息。这些逐层演化后的文本侧表示被称为上下文 token。它们没有空间结构、也不受去噪目标直接监督，此前只有零散的功能性证据（Padding Tone、ConceptAttention 等）。
本文的全部工作——读取（Contextual Reader）与监督（CoAl）——都围绕这一表示展开，它是论文的核心对象。
Q-Former 与瓶颈网络
Q-Former 是带少量可学习查询 token 的 Transformer 模块：查询先经 cross-attention 读取输入序列，再经自注意力与 FFN 精炼，把成百上千个输入 token 压缩成极少量输出（$Q \ll L\cdot M$），随后线性投影到目标模型嵌入空间。BLIP-2 用它连接视觉编码器与 LLM，LLaVA 则用更简单的线性投影实现类似接口。
Contextual Reader 的瓶颈网络和 CoAl 的对齐器都是 Q-Former 变体；压缩能力是“LLM 只能依赖上下文 token 答题”这一构造成立的关键。
表征对齐（REPA 路线）
REPA 在扩散训练中加入辅助损失，把中间视觉特征与预训练视觉编码器（如 DINOv2）的表征对齐，可大幅加速收敛并提升生成质量；后续 HASTE 对齐注意力特征、SRA 用模型自身表征、REG 用全局语义 token。iREPA 进一步指出这类对齐的收益主要来自保留空间结构而非全局语义。
CoAl 是把“对齐”从图像流搬到文本流的对照工作，只有理解 REPA 一族对齐什么、用什么教师，才能看出本文“全局语义教师 + 上下文 token”的互补定位。
Classifier-Free Guidance（CFG）与空条件
CFG 在训练时以一定概率（如 15%）把文本 prompt 替换为空串（null condition），推理时用条件预测与无条件预测的外推 $\tilde{\epsilon} = \epsilon_u + s(\epsilon_c - \epsilon_u)$ 放大提示影响。论文的“空 prompt”实验与 CoAl 的 Null-Condition Weighting 都建立在这一机制之上。
空条件设置是本文最惊人发现（无 prompt 也能读出语义）的来源，也是 CoAl 防止模型走“照抄 prompt”捷径的手段。
VLM-as-a-judge
用视觉语言模型代替人工评估开放式生成结果：把问题、被评答案、参考答案和图像一起交给评判 VLM（本文用 Prometheus-Vision 13B），按 1–5 分量表就语义正确性、完整性与图像一致性打分，再归一化到 $[0,1]$，是对自由文本答案的常用自动评估方案。
论文中所有“可读性”分数均由该协议产生，理解其打分口径才能正确解读 Figure 5、11、12、13 的曲线含义。
研究动机
在现代 MM-DiT（如 SD3、FLUX、Qwen-Image）中，文本 token 并不是固定的条件向量：它们在每一层都被多模态注意力反复改写，持续吸收来自图像 token 的信息，演化成所谓“上下文 token”。与视觉 token 不同，上下文 token 没有空间对应关系，也不受去噪目标的显式监督——视觉 token 的中间特征可以直接可视化、被引导和编辑，而文本侧这条数据流的内部状态长期处于黑箱。已有工作只提供了零散证据：ConceptAttention 用其注意力图定位生成图像中的概念，Padding Tone 证明文本条件下 padding token 能携带额外视觉信息，Contextual Repulsion 把上下文空间当作提升多样性的干预算子。但这些干预式分析都无法回答两个更基本的问题：上下文空间里到底编码了关于图像的哪些语义？这些语义如何随去噪过程演化、又如何最终影响生成结果？
本文的目标是本文目标分两步。第一步是“读”：构建 Contextual Reader，把 MM-DiT 中间层的上下文 token 投影到冻结 LLM 的输入空间，使人能用自然语言直接就“正在生成的图像”提问并获得回答，从而系统刻画上下文空间的语义内容——包括带 prompt 的条件设置与空 prompt 的无条件设置、去噪早期到晚期的时间演化，以及在 SD3.5 与 FLUX.2 两种架构上的一致性；并检验这种可读性与最终生成质量（HPS）是否相关。第二步是“用”：既然上下文 token 编码了丰富的全局视觉语义、且其清晰度与生成质量正相关，就应把它变成显式训练信号。为此作者提出 Contextual Alignment（CoAl）正则项，在从头训练与预训练模型微调两种场景下改进生成质量与分布覆盖，并验证其与现有对齐方法的互补性。
与已有工作不同的是，独特切入角度有三点。其一，本文不满足于注意力热图或因果干预等间接手段，而是把上下文空间直接解码成自然语言，实现对生成模型内部意图的“审讯”，语义内容以开放问答粒度可读。其二，作者设计了空 prompt 的无条件设置：此时上下文 token 无法从提示中获得任何信息，任何可读语义必然来自图像 token 经多模态注意力的回流——这把“图像到文本”的信息流隔离成可测量的现象，是此前工作未曾给出的证据。其三，在训练侧，已有对齐方法（REPA/HASTE/SRA）都作用于空间化的视觉特征，iREPA 更指出其收益来自空间结构；REG 的全局语义 token 只适用于无文本分支的类条件 DiT，直接搬到 MM-DiT 上反而劣于普通训练。本文据此提出用全局语义教师（SigLIP 池化嵌入）监督文本侧上下文 token，与现有局部空间对齐形成正交互补。
核心方法
方法分“读”与“训”两部分，共享同一直觉：上下文 token 是一个信息丰富但需要接口才能访问的语义寄存器。读取侧：在选定去噪时刻 $t$，从冻结 MM-DiT 全部 $L$ 个块提取文本侧注意力输出 $c^{(l)}(t) \in \mathbb{R}^{M\times d}$，加上可学习的层嵌入与时间步嵌入后拼接为 $C^{(t)} \in \mathbb{R}^{LM \times d}$；一个 2 层、16 查询的 Q-Former 瓶颈网络把它压缩成 16 个“上下文描述子”token，经线性投影与 L2 归一化送入冻结的 Qwen2.5-7B-Instruct，连同 15 个视觉问题之一自回归生成答案，仅用答案 token 上的交叉熵训练瓶颈网络。监督信号来自 InternVL3-8B 对每骨干 12,500 张生成图像的作答。训练侧：CoAl 用一个仅 1 层、1 查询的对齐器把第 4 层上下文 token 映射到 SigLIP 池化图像嵌入，以负余弦相似度损失与去噪目标联合优化；推理时丢弃对齐器，不增加任何开销。
核心创新在于把“可解释性”与“训练信号”接到同一个表示上。与以往研究注意力图、做因果干预不同，Reader 用“LLM 只能依赖上下文 token 答题”的构造性约束反推出表示内容：LLM 看不到原始 prompt，瓶颈若不把图像语义压进 16 个描述子，答题便无从谈起。由此得到两个此前未知的事实：上下文 token 在去噪极早期（约 8%）就已编码 prompt 未指明、随种子变化的生成特定语义；即使不给任何 prompt，语义也能凭去噪目标本身从视觉流回流进文本流，且在 FLUX.2 上晚期可读性反超条件设置。第二个创新是 CoAl：与 REPA 族对齐空间视觉特征不同，它用全局语义教师监督文本侧 token，与局部空间对齐鲁异且互补；Null-Condition Weighting 则封死“从 prompt 抄答案”的捷径，强制语义从图像 token 流入上下文 token——这恰好呼应了无条件阅读的发现。
方法步骤详情
第一步，构建读取器数据：从 MS-COCO 2014 过滤出单真人图像的 62,069 条候选说明，分层采样 2,000 个训练与 500 个留出 prompt，各生成 5 个种子，得到每个骨干 10,000 张训练图与 2,500 张留出图（1024×1024；SD3.5 用 28 步 FlowMatch Euler、CFG 7.0，FLUX.2 用 50 步、CFG 4.0）；由 InternVL3-8B 在不接触原 prompt 的前提下，对每张图回答人工设计的 15 个视觉问题（外观、姿势、动作、背景、相机角度等），答案作为监督。第二步，训练读取器：每次迭代采样 $\sigma \sim U[0,1]$，由存储的 $x_0$ 与 $\epsilon$ 重构 $z_t = (1-\sigma)x_0 + \sigma\epsilon$，前向冻结模型抽取全部层在注意力算子输出处（门控/残差/MLP 之前）的上下文 token，加层嵌入与 256 维正弦时间嵌入；Q-Former（2 块、8 头、16 查询、cross-first）压缩后线性投影进 LLM 嵌入空间并按词表嵌入均值范数缩放；LLM 仅收到 16 个连续前缀嵌入与问题，对答案 token 计算自回归交叉熵。AdamW、lr $10^{-4}$、batch 256、5,000 步，8×A100 约 4–8 小时。第三步，评估：Prometheus-Vision 13B 按 1–5 分比较预测与参考答案，归一化后对验证集取均值得到可读性。第四步，CoAl：抽取第 4 层上下文 token，经 1 层 24 头单查询 Q-Former 与 MLP（$D \to 1024 \to 768$）得到预测嵌入，以 $\mathcal{L}_{\text{CoAl}} = 1 - \cos\big(E_{\text{SigLIP}}(x_0),\, A(c^{(4)}(t))\big)$ 加入流匹配目标；条件被丢弃的迭代将损失乘 $w_{\text{null}}$（完整训练 2.0、微调 6.0），有条件迭代乘 $w_{\text{text}}$（1.0 / 0.25），整体权重 $w_{\text{CoAl}}$ 为 1.3 / 0.75。
技术新颖性
新颖性体现在四个层面。一是表示读取范式的转移：从注意力热图、探针分类器到“冻结 LLM + 瓶颈投影”的自然语言接口，把对扩散模型内部状态的分析提升到开放问答粒度，与 ConceptAttention、Padding Tone、Contextual Repulsion 等干预式工作互不重叠。二是实验设计的归因能力：Prompt Reader、Latent Reader、$x_0$+VLM 三个对照把“信息来自 prompt、LLM 先验还是视觉流”分离干净，使“无条件可读性反超”这类结论站得住脚。三是把观测转成训练目标：CoAl 是首个针对 MM-DiT 文本侧上下文 token 的对齐正则，以全局语义教师（SigLIP 池化嵌入）补充 REPA 族的空间结构教师，消融显示两者叠加后所有指标列同时最优。四是 Null-Condition Weighting 巧妙利用 CFG 的条件丢弃机制构造信息流压力测试，这一技巧本身可迁移到其他多模态架构的训练中。
实验结果
读取侧发现：条件设置下可读性从一开始就高并持续上升，在 80% 处达峰、100% 略降；空 prompt 时起点低但快速攀升，FLUX.2 上最终反超条件设置、SD3.5 上逼近之。归因实验：条件 Reader 全程优于只看 prompt 的 Prompt Reader（20%–80% 差距最大），无条件版在 FLUX.2 上也反超后者；条件 Reader 全程优于读潜变量输出的 Latent Reader，无条件版从 10% 起占优；直接看预测干净图的 $x_0$+VLM 仅在约 40% 之后领先——更早阶段，不看图像的上下文 Reader 可读性反而更高。同一 prompt 的 5 个种子按 Reader 分数排序，首尾的 HPS 差 $\Delta$HPS 在四种设置（2 架构 × 2 条件）下一致为正（约 0.05–0.30），低可读性样本常伴畸形物体或曲解 prompt；留出问题泛化实验中条件 Reader 平均分仍约 0.50。训练侧：从头训练（COCO 256²，150k 步）中 CoAl 使 REPA 的 FID 4.95→4.29（-13.4%）、Recall 0.338→0.407（+20.5%），SRA 5.89→4.36（-25.9%，Recall +52.5%），HASTE 4.87→4.44，Vanilla 6.52→6.23，REPA+CoAl 在所有指标列均为全场最佳；微调 SD3-Medium（Fine-T2I，30k 步）时 REPA/HASTE/SRA（FID 17.45/17.53/17.48）皆劣于 Vanilla（17.18），仅 Vanilla+CoAl 取得最佳 FID 16.75 与 Recall 0.969，各组合 KID 改善 5.6%–20.7%。消融：SigLIP 教师最优（CLIP 4.71、DINOv2 4.67、Qwen3-VL 5.03）；第 4 层对齐最优（第 20 层 FID 恶化至 17.70）；去掉空条件加权使微调 FID 从 16.75 恶化到 17.60；CFG 1–9 全程扫描中 CoAl 始终最强。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 从头训练 MM-DiT（MS-COCO 2014，256×256，150k 步） | FID ↓ | 4.29（REPA+CoAl） | 4.95（REPA）；6.52（Vanilla） | 较 REPA 相对降低 13.4%，较 Vanilla 降低 34.2%；且所有指标列全场最佳 |
| 从头训练 MM-DiT（弱基线组合） | FID ↓ | 4.36（SRA+CoAl） | 5.89（SRA） | 相对降低 25.9%；KID -38.7%，Recall +52.5% |
| SD3-Medium 微调（Fine-T2I curated，30k 步） | FID ↓ | 16.75（Vanilla+CoAl） | 17.18（Vanilla）；REPA/HASTE/SRA 为 17.45/17.53/17.48 | 相对降低 2.5%；KID -18.4%，Recall 0.966→0.969 |
| 从头训练 MM-DiT（分布覆盖） | Recall ↑ | 0.407（REPA+CoAl） | 0.338（REPA） | +20.5%（HASTE+CoAl 0.400、SRA+CoAl 0.398 亦全面超过各自基线） |
| 可读性与生成质量的相关性（同 prompt 5 种子） | ΔHPS（Reader 排序首尾名次的 HPS 差） | 四种设置一致为正，约 0.05–0.30 | 0（假设排序无区分度） | Reader 高可读性种子系统性获得更高 HPS |
局限与改进
作者承认的局限：读取器的监督与评估都建立在 VLM 之上——训练答案来自 InternVL3-8B，评分由 Prometheus-Vision 完成，而 $x_0$+VLM 基线恰好也用 InternVL3，这种同源对齐可能系统性偏爱图像基线；15 个问题全部围绕 COCO 单人图像设计，问题集与领域都较窄；可读性与 HPS 的关联只是相关性证据（$\Delta$HPS），并未证明因果。我的补充观察：其一，Reader 的训练与使用需存储每个生成的初始噪声并重跑扩散前向，工程上不易做成在线工具；其二，LLM（Qwen2.5-7B）的先验可能“脑补”答案，Prompt Reader 对照缓解但不能消除这一混杂因素；其三，CoAl 的实证规模有限——从头训练只是 24 层、768 维、256×256 的缩小版模型，微调也仅限 SD3-Medium，能否迁移到 FLUX.2 级别的大模型与视频 MM-DiT 尚未知；其四，$w_{\text{null}}=6.0$ 与 $i_{\text{start}}$ 等超参在两种设置间差异很大，暗示需按场景仔细调参；其五，论文未报告人工评测，生成质量改善主要由 FID/KID 等分布级指标支撑。
独立分析的弱点
独立分析出的弱点与改进方向：（1）问题集封闭——只有 15 个固定问题，无法探询任意感兴趣的属性；可改为开放式字幕或让 LLM 自由生成问题，把 Reader 升级为真正的“任意审讯”接口。（2）VLM 答案天花板——Reader 学到的是 InternVL3-8B 的视角，其幻觉与盲区会被直接继承；可用多 VLM 投票或人工校验子集缓解，并对 “unknown” 答案做更细粒度建模。（3）相关非因果——$\Delta$HPS 为正不等于“读取更清楚所以图像更好”；可在训练时对上下文清晰度做扰动（加噪、投影）做因果检验，这同时能验证 CoAl 的作用机制。（4）超参脆弱——微调场景 $w_{\text{text}}=0.25$、$w_{\text{null}}=6.0$ 的量级差异说明损失易被有条件捷径主导；可设计随训练进程自适应的权重调度。（5）架构外推性——结论全部来自双流/单流 MM-DiT，对 U-Net+cross-attention（文本侧仅注入一次、不逐层更新）不适用，应明确界定适用边界。（6）单模态局限——视频、音频 MM-DiT（如 CogVideoX、HunyuanVideo）的上下文 token 是否同样可读、可监督，是显而易见却未被验证的空白。
未来方向
作者提出的方向：理解这些内部表示如何形成、演化并最终塑造生成图像——“模型很早就知道自己要生成什么”这一事实对生成过程本身提出了大量新问题，值得继续挖掘控制与改进的机会。基于成果可延伸的研究：其一，把 Reader 变成在线监控器，在去噪早期依据可读性对坏样本做早停或拒绝采样，几乎零额外训练成本；其二，反向操作——不仅“读”还“写”，通过编辑上下文 token 实现语义级引导（如强制或擦除某属性），与 Contextual Repulsion 的多样性干预一脉相承；其三，将 CoAl 推广到视频与音频 MM-DiT，检验跨模态上下文寄存器的通用性；其四，把可读性作为预训练过程的无监督诊断信号，替代昂贵的人工抽检；其五，探索更强或更多样的语义教师（如 VLM 字幕嵌入、分层语义表征），并研究“对齐早层”与大 CFG 鲁棒性（图 19 中 CoAl 在 CFG 7–9 下明显更稳，暗示条件/无条件预测更一致）之间的机理联系。
复现评估
复现条件总体友好，组件全部开源可得：SD3.5、FLUX.2、SD3-Medium 权重，Qwen2.5-7B-Instruct、InternVL3-8B、Prometheus-Vision 13B、SigLIP 等模型，以及 MS-COCO 2014 与 Fine-T2I（huggingface.co/datasets/ma-xu/fine-t2i）数据集均公开；附录给出了几乎全部超参（学习率、batch、EMA 0.9999、$w_{\text{CoAl}}/w_{\text{null}}/w_{\text{text}}$、对齐层选择、特征抽取位置等）。算力门槛适中：训练一个 Reader 只需 8×A100 约 4–8 小时，瓶颈网络参数量极小；主要开销在于数据生成（每骨干 12,500 张 1024×1024 图像）与训练时按需重跑冻结扩散前向。完整训练复现（150k 步、batch 256、256²）在 8×B200 上约 12.5 小时，SD3-Medium 微调约 5.6 小时，属可行范围。论文未提及官方代码发布，需自行实现 Q-Former 与特征抽取，但描述足够细致。综合评估：读取部分复现难度中等偏低，CoAl 训练部分中等。
论文图表
六个不同 prompt（做饭、打网球、滑板、跳舞、弹吉他、骑马）的生成在去噪早期视觉预测仍模糊时，用空 prompt 下的 Reader 从上下文 token 中解码出的活动描述，全部回答正确。
一图点题：即使模型没有收到任何 prompt，极早期的上下文 token 已编码了最终图像的活动语义，是全文最核心论断的直观展示。
消融语义教师、对齐层与空条件加权：SigLIP（FID 4.29）优于 CLIP 4.71、DINOv2 4.67、Qwen3-VL 5.03；第 4 层最优，越深越差（L20 5.30/17.70）；去掉空条件加权使完整训练 FID 4.44、微调 17.60，均劣于完整方案。
验证方法的三个关键设计选择，说明“早层对齐 + 全局语义教师 + 空条件加权”缺一不可，尤其微调场景对后者最敏感。