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
cs.CV
2026-07-21
文生图 DiT 里那些结构模板 token 在编码器出口几乎不含提示语义,却成为图到文注意力的主汇点和物体身份的隐式寄存器;据此裁掉最关注提示的注意力头,省 20% 注意力 FLOPs 只掉 1.4 分。
文生图扩散 Transformer(DiT)同时处理文本和图像 token,但去噪过程中内部到底怎么算的,一直不清楚。这篇做了一套因果可解释性框架:把注意力矩阵分解成文-文、文-图、图-文、图-图四块,再对 token 段、注意力头、层做有针对性的干预(交换段、跨轨迹移植头、逐层因果屏蔽)。主要研究对象是 Qwen-Image 系列,也验证了 FLUX.2、Krea-2-Turbo、Qwen-Image-Edit。
核心发现是那些结构模板 token(比如对话分隔符、<|imend|> 这类)的两面性。在编码器出口,它们几乎不含提示特有的语义:跨提示交换模板段,生成结果基本不变;一个跨所有提示通用的、无内容的模板就能渲染任何物体。但在 DiT 内部,同样的模板 token 变成了图到文注意力的主汇点,吃掉 76%-92% 的图到文注意力,并因果性地维持物体身份,扮演「隐式语义寄存器」。
机制是间接的:语义先从提示 token 注入图像 latent(在第一个 Transformer block 内完成),模板 token 再从图像流读回身份,而不是直接从提示 token 拿。证据是,屏蔽模板到提示的注意力,物体几乎不变;屏蔽模板到图像的注意力,物体在前两层就崩了。
据此设计免训练裁剪:最关注提示 token 的头其实是可有可无的。按图到文语义注意力从高到低排序,裁掉 top-K 头,在最后 80% 去噪步生效。在 Qwen-Image-2512 上,裁掉 360 个头(占 1440 个头的 25%)省 20% 联合注意力 FLOPs,GenEval 从 76.1 掉到 74.7,只掉 1.4 分;更优配置 K=216 省 12% FLOPs 只掉 0.3 分。反过来按结构汇点排序裁,GenEval 掉到 69.6,差得多,证明裁的是对的那批头。
对做 DiT 的人,最直接的用处是那条免训练裁剪规则,省两成注意力 FLOPs 几乎不伤质量。更深的意义是它推翻了一个直觉:输入端编码语义的 token,不一定是生成时维持语义的 token。提示语义走的是「提示→图像→模板」的间接路径,模板 token 才是生成过程中物体身份的真正载体。这套「寄存器头管身份、渲染头管像素」、深度上「早层定身份、中层传递、晚层精修」的分工,和 LLM 推理的阶段划分惊人地相似。
作者承认:这篇主要是分析性的,免训练裁剪只是对分析结果最朴素的一次应用,设计感知汇点的稀疏注意力本可拿更大收益;「为什么这些无内容的结构 token 会变成寄存器」本身仍是开放问题。读下来还有一点:结论主要建立在 Qwen-Image 系列上,FLUX.2、Krea-2 等只做了扩展验证,寄存器机制是否在所有 DiT 架构上同样成立、同样可裁,还不够充分。