文本模板Token是扩散Transformer中的隐式语义寄存器 Text Template Tokens Are Implicit Semantic Registers in Diffusion Transformers

方法步骤完整描述如下。第一步，条件序列构造：prompt $p$ 被包装为 $c_{raw}=[r_{prefix};\ p;\ r_{suffix}]$ 送入 Qwen2.5-VL，取末层隐藏态并丢弃固定的 34 个模板前缀位置，得到 $c=(c_1,\dots,c_{n_{cond}})$，按位置切分为语义 span $S$ 与模板 span $R$（如 "An apple" 为 $|S|=2$、$|R|=5$）。第二步，注意力质量度量：对每个 block $l$、head $h$ 的 post-softmax 联合注意力矩阵 $A^{(l,h)} \in \mathbb{R}^{(n_{cond}+n_{img})\times(n_{cond}+n_{img})}$，定义 $m^{(l,h)}_K(Q)=\frac{1}{|Q|}\sum_{i\in Q}\sum_{j\in K}A^{(l,h)}_{i,j}$，聚焦图像 query $Q=I$ 对条件 key $T=S\cup R$ 的 I2T 块，沿全部 $L\times H\times T$ 聚合得 $\bar{m}^R, \bar{m}^S$ 与 per-token 比值。第三步，span 干预：(a) 跨 prompt 交换 $[S_A; R_B]$——固定 seed 与采样器，保留 A 的内容词、换成 B 的模板 token，用 DINOv3 [CLS] 余弦相似度验证物体不变；(b) 平均模板替换——把 100 个 GenEval prompt 的 R 平均成单个 prompt 无关的 $\bar{R}$，替换 40 个 held-out prompt 的 R 构成 $[S; \bar{R}]$。第四步，跨轨迹头移植：同时跑 A="An apple" 与 B="A banana" 两条轨迹（2 步蒸馏 LoRA、只跑单步以固定时间步），第 $n$ 轮把 A 中按语义读取注意力 $m^S$ 排序后的前 $n$ 个头的 (q,k,v) 投影换成 B 的对应头，共 $L\times H=1440$ 头，观察图像何时翻转；另做 per-span 整段交换作为对照。第五步，逐层因果 masking：在寄存器 query $R$ 处 softmax 前屏蔽其对语义 token $S$ 或对图像潜变量的注意力，按前 $k$ 个 block 累积施加，比较物体何时塌缩；用'仅屏蔽第一块无效'这一事实推断 $S\to I$ 注入在第一块完成、$I\to R$ 在第二块完成。第六步，剪枝规则：按 $\bar{m}^S$ 降序对 1440 头排序，剪除 top-$K$ 头的全部 q/k/v 投影、输出投影与注意力计算，只在最后 80% 去噪步生效（$f=0.2$，前 20% 步保留），joint-attention FLOPs 减少比例为 $(1-f)K/1440$。

技术新颖性体现在五个层面。(1) 框架层面：这是首个面向大规模文生图 DiT 的系统性因果可解释性框架，把 token 分解、span 干预、头移植、逐层 masking 四种干预组合成可复用的流水线，并在 60 块 × 24 头 = 1440 头的 20B 级模型上实际可行（作者开源了代码）。(2) 概念层面：明确区分'注意力模式'与'因果功能'——模板 token 的 sink 模式（描述性）与寄存器功能（因果性）由不同的实验分别确立，这种严谨性在视觉生成可解释性工作中少见。(3) 机制层面：$S \to I \to R$ 间接电路及'第一块完成语义注入、第二块完成读回'的时序定位是全新发现；深度剖面（模板 per-token 注意力在第一块接近零、block 11 才达半平台、block 20 达峰值）进一步排除了'输入嵌入伪影'解释，支持这是学到的、深度局域化的机制。(4) 泛化层面：同一现象在 Qwen-Image/2512、FLUX.2（32B、Mistral-3 编码器、sink 分散于系统块）、Krea-2-Turbo（12B 单流）、2 步蒸馏模型、指令编辑模型、中英文 prompt 上均成立，证明这是 LLM 编码器 + 联合注意力这一架构范式的普遍性质。(5) 应用层面：把'高 $m^S$ 头因果惰性'这一机制发现直接转成 prompt 无关、免训练、无需逐 prompt 剖析的剪枝规则，完成了从机制到系统的闭环。

以下是独立于作者的弱点分析，每条附改进方向。(1) 感知质量是剪枝的主要代价：K=360 时 LPIPS 0.39、HPSv3 从 9.56 降至 8.76，且 Fig 18 的定性结果显示细节纹理随预算增大逐渐糊化。改进方向：sink-aware 稀疏注意力——保留寄存器位置作为 key（作者在附录 A 也承认这一点），只剪除其对应的 query 侧计算或 value/output 投影，把'维护语义所需的读取'与'渲染细节所需的计算'分离，理论上能保住更多感知质量。(2) 排序是静态且 prompt 无关的：$\bar{m}^S$ 是跨数据集聚合量，特定 prompt 下被剪的头可能恰好参与关键语义（例如多物体、属性绑定的 prompt）。改进方向：训练一个轻量 per-prompt 头选择器，或按去噪步动态路由——前 20% 步全保留的设定已经暗示时序自适应有价值。(3) 对涌现机制沉默：论文无法回答'为什么内容为空的 token 会变成寄存器'，也就无法预测新架构是否自动获得该性质、训练配方如何促进它。改进方向：沿预训练 checkpoint 序列追踪 sink 的出现时间与深度剖面（类比 LLM 中 sink 在预训练中涌现的研究）。(4) 因果实验在单步蒸馏设置下完成，多步原生轨迹上的重复验证（例如滑动窗口 masking）仍是空白。(5) CFG=1 下的失败尾部（Fig 4b 中相似度低至 0.35~0.56 的个例，物体被整体替换或数量改变）说明对模板表示的粗暴操作在高保真场景有风险。改进方向：两阶段采样——先用高 CFG 锚定语义再放宽。(6) 编辑模型分析只给了 per-token 注意力统计，寄存器与 vision token、参考图 latent 之间的竞争关系缺乏因果干预实验。

作者明确提出的方向有两个：其一是把剪枝升级为 sink-aware 稀疏注意力，保留寄存器位置为 key 而剪除其余，预期收益远超当前 20%；其二是解释内容为空的模板 token 为何会演化成语义寄存器，这需要训练动力学层面的追踪。基于论文成果还可以延伸：(3) 显式寄存器架构设计——训练时就加入专用语义寄存器 token 并施加结构约束（如只允许在指定层读写），把偶然涌现变成可控设计，可能同时提升效率与可控性；(4) 寄存器作为编辑接口——既然物体身份存于寄存器状态，直接改写寄存器表示有望实现免训练的属性替换、概念擦除或风格迁移，比修改 prompt 或微调 LoRA 更精准；(5) 蒸馏对齐——少步蒸馏时对齐 $S\to I\to R$ 电路的层间时序，或能缓解蒸馏常见的语义漂移；(6) 系统工程落地——把 20% FLOPs 削减与 KV cache 压缩、头并行策略结合，转化为实际延迟与显存收益；(7) 利用 early commit / middle carry / late refine 的三段分工做分块缓存或步间特征复用，middle 40 层几乎不重写身份，是缓存的理想目标；(8) 把分析扩展到视频生成 DiT 与多模态条件（编辑、identity 保持）下的寄存器竞争格局。

复现条件总体良好。代码已开源（GitHub: Met4physics/DiT-Interpretability），覆盖注意力统计、span 交换、头移植、masking 与剪枝全流程。涉及模型全部公开可得：Qwen-Image、Qwen-Image-2512、Qwen-Image-Edit-2511（HuggingFace 开源权重）、FLUX.2、Krea-2-Turbo，以及社区发布的 2 步蒸馏 LoRA（Wuli-art）。评估数据均为标准公开基准：GenEval（553 条）、DPG-Bench（1,065 条）、Qwen-Image-Bench（1,000 条中文）、GEdit-Bench。作者报告的环境为 PyTorch 2.12.1 + DiffSynth-Studio v1.1.7，NVIDIA H20（1,024GB 显存）+ 192 核 CPU，固定随机种子 42，保证逐图可复现。这是纯分析性工作，无需任何训练，主要算力开销在于推理时缓存并检查 1440 个头的注意力矩阵与 q/k/v 投影；头移植和 masking 用标准 PyTorch forward hook 即可实现。难度评估：中等偏高——难点不在算法（所有干预都概念简单），而在资源：H20 1TB 显存配置少见，在 20B 级模型上缓存全量 per-head 注意力矩阵需要大显存或 offload 策略；若只在 Qwen-Image 上以 bf16、单张 80GB A100/H100、适当分辨率做核心实验（注意力统计 + 单步头移植）是可行的，FLUX.2 32B 的复现则需要多卡。总体属于可复现性较好的机制分析论文。