HyperAIHyperAI

Command Palette

Search for a command to run...

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

摘要

文本到图像扩散 Transformer(DiT)联合处理文本和图像 token,但其去噪过程中的内部计算机制仍不清晰。我们为现代大规模 DiT 引入了一个因果可解释性框架,该框架将注意力分解与跨 token 跨度、注意力头和层的定向干预相结合。利用该框架将提示内容 token 与结构模板 token 分离,我们发现结构 token 在编码器输出端几乎不携带提示特定信息。然而,令人惊讶的是,它们成为主要的图像到文本注意力汇聚点,并在 DiT 内部因果性地维持对象身份,充当隐式语义寄存器。我们表明,它们间接获取这一身份:提示语义首先被注入图像潜变量,然后被读回模板 token,而非直接从提示 token 传递。受上述发现启发,我们为 DiT 设计了一种无需训练的剪枝规则。最强烈关注提示 token 的注意力头是可舍弃的,剪枝这些头可移除 20% 的注意力 FLOPs,而 GenEval 指标仅下降 1.4 点。我们进一步揭示了 DiT 中的生成计算如何跨注意力头和深度进行组织,将语义路由与视觉合成分离,并从身份形成推进到传播与细化。我们的工作不仅揭示了在输入端编码语义的 token 未必是在生成过程中维持语义的 token,还为 DiT 的内部机制提供了一种因果视角。

一句话总结

由南京大学、阿里巴巴集团和浙江大学的研究者开发的一种结合注意力分解与定向干预的因果可解释性框架,揭示了扩散 Transformer 中的文本模板 token 充当隐式语义寄存器,通过图像隐变量间接获取身份信息,并推导出一种无需训练的剪枝规则,在 GenEval 上仅下降 1.4 点的情况下移除 20% 的注意力 FLOPs。

核心贡献

  • 一种针对扩散 Transformer 的因果可解释性框架,结合注意力分解与跨 token 片段、注意力头和层的定向干预,追踪去噪过程中文本-图像 token 的交互。
  • 聊天模板结构 token 在编码器输出端几乎不携带提示特定信息,却充当隐式语义寄存器。它们通过接收来自图像隐变量的提示语义,并将该身份信息经由图像到文本注意力进行路由,从而因果性地维持物体身份。
  • 一种无需训练的头剪枝规则,通过剪枝那些对提示 token 关注最强的注意力头,移除 20% 的注意力 FLOPs,GenEval 仅下降 1.4 点。该框架还揭示,生成计算在注意力头和深度上分离了语义路由与视觉合成,从身份形成逐步推进到传播和细化。

引言

现代文本到图像扩散 Transformer 越来越依赖带有聊天模板的大型语言模型编码器,这些模板在用户提示之外引入了系统分隔符和用户分隔符等结构 token。这些 token 通常被视为格式残留并被忽略。同时,对扩散 Transformer 的机制理解远落后于语言模型:不清楚在去噪过程中哪些层或注意力头实际路由和存储条件信息。作者引入一种因果可解释性框架,利用 token 级注意力分解、片段干预、头移植和逐层掩码来追踪信息流。将该框架应用于带有聊天模板的扩散 Transformer 后,他们发现结构 token 形成主导的文本侧注意力汇,并充当语义寄存器,从不断演变的图像隐变量中提取物体身份,而非直接从提示 token 获取。这一发现引出一条无需训练的头剪枝规则,移除因果惰性的读取提示的注意力头,削减 20% 的注意力 FLOPs,而准确率下降很小。

方法

作者研究 Qwen-Image 系列,这是一种双流多模态扩散 Transformer(MMDiT),通过流匹配目标迭代地对隐变量 token 去噪来合成图像。设 zRnimg×dz \in \mathbb{R}^{n_{img} \times d}zRnimg×d 为图像隐变量 token 序列,cRncond×dc \in \mathbb{R}^{n_{cond} \times d}cRncond×d 表示文本条件。一个包含多个块的 Transformer 预测流,采样器将其从纯噪声积分到干净的隐变量,再由 VAE 解码为像素。与依赖 CLIP 或 T5 编码器的模型不同,Qwen-Image 的条件来自应用于聊天格式提示的大型视觉语言模型。用户提示被包裹在聊天模板中,保留的条件序列从编码器最后一层隐藏状态读取,丢弃固定的模板前缀。该序列分为两个不相交的部分:语义片段包含提示内容 token,结构片段包含尾部的聊天模板 token,用于界定对话格式。

在每个 MMDiT 块中,文本和图像 token 通过一个单一的自注意力机制处理,该机制作用于条件和隐变量 token 的拼接。将联合注意力矩阵的行和列划分为文本和图像段,得到四个不同的块。

这些块分别表示文本查询对文本键和图像键的关注,以及图像查询对文本键和图像键的关注。图像到文本块作为跨模态通路,文本条件通过该通路作用于图像隐变量。为了量化每个 token 接收的注意力大小,作者使用完整的 softmax 后联合注意力矩阵,并定义从查询组到键片段的注意力质量。

作者引入一个因果可解释性框架,结合 token 级注意力分解与片段级条件干预,以追踪去噪过程中条件信息在何处被读取、路由和存储。他们发现结构聊天模板 token 充当主导的注意力汇。在整个去噪轨迹中,图像查询对无内容模板片段的关注比提示内容高出一个数量级。

注意力表面在结构片段上形成尖锐的脊,而图像 token 之间的注意力保持低位。尽管这些模板 token 几乎不携带明确的提示特定语义,它们在扩散 Transformer 内部充当隐式语义寄存器。作者追踪信息流,发现它是隐式的:语义首先从语义 token 传递到图像 token,然后才从图像 token 传递到模板 token。屏蔽寄存器查询对图像隐变量的注意力会在前几个块内使物体崩溃,证实寄存器通过早期关注图像流来组装其内容。

基于这一分析,作者推导出一条无需训练的加速规则。由于图像到语义注意力高的头对于传递物体身份是因果惰性的,他们按该指标对所有注意力头进行排序,并剪枝排名靠前的头的计算。这种剪枝仅在去噪后期步骤应用,因为早期步骤更敏感并确定物体身份。由于联合注意力块的每个组件随头数线性扩展,该方法在很好地保留物体正确性的同时,按比例移除了计算成本。

实验

跨提示交换和平均实验表明,文本模板 token 仅携带弱语义,其移除或替换仅在引导强度低时导致模型坍缩为通用人像,揭示该人像是无条件默认输出。利用这一洞察,一种无需训练的头剪枝方法根据图像到语义注意力分数对 1,440 个联合注意力头进行排序,并在最后 80% 的去噪步骤中剪枝排名靠前的头。该方法移除高达 20% 的联合注意力 FLOPs,GenEval 准确率仅小幅下降,并证实基于语义寄存器角色选择要剪枝的头对于保留物体正确性和感知质量至关重要。

短提示中,尾部的结构模板 token 充当主导注意力汇,吸收的图像查询注意力远多于语义内容 token。随着提示长度增加,语义 token 获得注意力份额,而结构注意力下降,但每个结构 token 的注意力在所有基准测试中仍比每个语义 token 高出数倍。结构 token 仅携带弱语义,移除它们会导致生成回退到无条件默认输出。对于短提示(平均 7.9 个 token),结构 token 的平均注意力质量为 0.19,语义 token 为 0.04,且在 89% 的注意力位点上结构质量超过语义质量。对于长提示(平均 82.1 个 token),语义质量升至 0.17,结构质量降至 0.08,结构质量占优的位点比例降至 12%。每 token 注意力质量比在所有基准测试中始终保持高位(6.2–7.5),因此即使对于长提示,每个结构 token 单独吸引的注意力仍多于每个语义 token。结构模板 token 仅编码弱语义内容,当它们被完全移除时,生成坍缩为模型的无条件默认结果。

基于图像到语义注意力 m_S 的无需训练头剪枝在减少联合注意力 FLOPs 的同时保留了物体正确性。剪枝四分之一的所有头(K=360)削减 20% 的 FLOPs,GenEval 仅下降 1.4 点,但感知质量下降更快,因此 K=216 的较轻预算提供了最佳权衡。头的选择至关重要:所提出的排序显著优于结构汇和随机头选择。在 K=360 时,20% 的联合注意力 FLOPs 被移除,GenEval 准确率从 76.1 降至 74.7。感知质量(LPIPS)比物体正确性下降更快,使得 K=216 成为最佳平衡点。在所提出的排序下,K=288 达到 GenEval 75.5;切换到结构汇头则降至 69.6,随机选择则崩溃至 51.3。HPSv3 偏好分数在 K=216 时保持高位(9.47),但在 K=360 时降至 8.76,K=400 时降至 8.44。剪枝仅在最后 80% 的去噪步骤中应用,保留早期身份确定步骤不受影响。该方法与提示无关,无需额外训练或在线优化。

第一个实验分析文本到图像生成中的注意力分布,揭示结构模板 token 在短提示中充当主导注意力汇,而语义 token 在较长提示中获得注意力份额;然而,每 token 结构注意力始终保持更高,移除这些 token 导致生成回退到无条件默认输出,证实其弱语义内容。第二个实验验证了一种无需训练的头剪枝方法,该方法根据图像到语义注意力对注意力头排序,实现联合注意力 FLOPs 减少 20%,物体正确性损失极小,但感知质量下降更快,且所提出的排序显著优于结构汇或随机头选择。总体而言,这些研究突显了结构 token 的注意力汇行为,并证明有依据的头剪枝可以有效减少计算量,同时基本保持生成质量。


用 AI 构建 AI

从创意到上线——通过免费 AI 协同编码、开箱即用的环境和最优惠的 GPU 价格,加速您的 AI 开发。

AI 协同编码
开箱即用的 GPU
最优定价

HyperAI Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供