文本模板标记是扩散中的隐式语义寄存器 Transformer

现代文本到图像的扩散 Transformer (DiTs) 通过联合注意力生成图像，其中文本和图像标记在单个序列中直接交互。在大规模 DiT 中，条件输入不仅包含用户提示，还包含基于 LLM 的文本编码器引入的聊天模板标记。然而，这些标记如何参与去噪计算仍然知之甚少。为了探讨这一点，我们引入了因果解释框架。使用它来将提示内容标记与聊天模板标记分开，我们发现模板标记在编码器输出中携带很少的特定于提示的信息。然而令人惊讶的是，它们作为图像到文本的主要注意力下沉而出现，并因果性地维持 DiT 内的对象身份，充当隐式语义寄存器。我们表明他们间接获得了这种身份。他们不是读取提示标记，而是从图像潜伏中提取身份，提示语义已经在第一层注入到该图像潜伏中。我们进一步揭示了 DiT 中不同头部和深度的劳动分工，其中不同的头部路由语义或呈现视觉结构，并且身份在早期块中提交，由中间块承载，并在后期块中细化。作为实际回报，该分析产生了 无需训练 剪枝规则，该规则消除了因果惰性的提示阅读头，并以 GenEval 准确度中 1.4$ 点的成本减少了 20\%$ 的联合注意力 FLOP。总的来说，我们的工作不仅揭示了输入处编码语义的标记不需要是在生成过程中维护它们的标记，而且还提供了扩散 Transformer 中内部机制的因果视图。