处理Diffusion Transformers异常值令牌方法

基本信息

导语

扩散Transformer在图像生成中取得显著进展，但其内部产生的离群令牌吸引过多注意却仅含局部语义，导致生成质量受损。作者提出双阶段寄存器（DSR），在有标签训练时通过专用寄存器抑制离群特征，在无监督或测试阶段采用递归轻量寄存器逐层校正令牌行为。实验表明DSR在不显著增加计算开销的情况下有效抑制离群令牌，为扩散模型生成质量的提升提供了轻量化的新思路。

摘要

背景

扩散Transformer（DiT）在图像生成任务中取得了显著进展，但仍有离群令牌现象影响生成质量。已有研究指出，视觉Transformer（ViT）会产生少量高范数令牌，吸引过多注意力却只携带局部信息，这一现象在生成模型中的作用尚未得到充分探索。

问题

在表征自编码器（RAE）-DiT 管线中，离群令牌同时出现在编码器和去噪器：预训练的ViT编码器会产生离群表示，而DiT自身在中间层也会生成内部离群令牌。实验表明，单纯掩蔽高范数令牌并不能提升性能，说明问题根源并非极端数值本身，而是与局部语义受损密切相关。

方法

为控制离群令牌，本文提出Dual‑Stage Registers（DSR）——一种基于寄存器的双阶段干预方案。具体包括：

  1. 训练寄存器：在有可用训练数据时，直接训练专用寄存器来抑制离群特征。
  2. 递归测试时寄存器：在无监督或测试阶段，采用递归方式在每层插入轻量寄存器，逐层修正离群行为。
  3. 扩散寄存器：针对去噪器，引入扩散过程感知的寄存器，以适配扩散模型的时序特性。

实验结果

在 ImageNet 类别图像生成和大规模文本‑到‑图像（T2I）任务上，DSR 均显著降低离群伪影，提升生成细节与整体质量。与基线 DiT 相比，FID、CLIP‑Score 等指标均有明显改善，验证了离群令牌控制对构建更强 DiT 的重要性。

结论：离群令牌的有效抑制是提升扩散Transformer生成质量的关键技术路径，DSR 提供了可行且通用的解决方案。

技术分析

研究背景

  • 摘要指出DiT在图像生成取得显著进展，但ViT产生的少量高范数令牌会吸引过多注意力并携带局部信息，离群令牌现象尚未被系统研究。（来源：摘要）
  • 推断离群令牌可能是导致生成细节出现伪影或局部不连贯的关键因素，间接影响FID等指标。

核心方法

  • 作者提出Dual‑Stage Registers (DSR)，包括训练寄存器、递归测试时寄存器和扩散寄存器，对编码器和去噪器分别进行双阶段干预。（来源：摘要）
  • 推断训练寄存器利用有标注数据学习抑制离群特征的映射；递归寄存器在无监督情况下通过逐层轻量插入修正离群行为；扩散寄存器针对扩散时序特性进行适配。

理论基础

  • 核心假设是离群令牌导致局部语义受损，而非单纯的极端数值，直接掩蔽高范数令牌未能提升性能。（来源：摘要）
  • 推断寄存器提供可学习的向量，在特征空间吸收或重新分配高范数分量，恢复局部语义的完整性。

实验与结果

  • 在ImageNet类别图像生成和大规模文本‑到‑图像任务上，DSR显著降低离群伪影，FID、CLIP‑Score等指标相对基线DiT有明显改善。（来源：摘要）
  • 推断改善幅度可能随模型规模和训练数据量的提升而进一步扩大，具体数值需在更大规模实验中验证。

应用前景

  • 推断DSR可迁移至其他基于ViT的生成模型（如自回归或混合扩散模型），提升细节保持和整体生成质量。

研究启示

  • 离群令牌的治理被强调为构建更强DiT的关键技术路径，为后续优化提供新方向。

相关工作对比

  • 已有研究尝试直接掩蔽高范数令牌，却未能提升性能，说明仅靠数值剪裁不足以解决问题。（来源：摘要）
  • 与之相比，DSR通过可学习的寄存器进行细粒度、层级特征校正，兼顾模型容量与局部语义保持。

关键假设、潜在失效与可证伪方式

  • 关键假设：离群令牌是导致局部语义受损的主因，寄存器能够有效吸收或重新分配其影响。（推断）
  • 潜在失效：若离群根本源于模型结构的根本缺陷或极端噪声，寄存器可能不足以根治；训练数据不足时训练寄存器效果受限。（推断）
  • 可证伪方式：删除所有寄存器或使用随机初始化寄存器，若生成质量仍保持或下降，则假设不成立；或在极端噪声条件下测试，观察寄存器是否能抑制伪影。（推断）

学习要点

  • 在 Diffusion Transformers 中，少数 token 的激活值异常巨大，导致训练不稳定和生成质量下降。
  • 这些异常 token 主要来源于深层自注意力和残差连接的高方差放大效应。
  • 论文提出基于统计阈值（如高阶矩或分位数）的轻量检测方法，可实时识别异常 token。
  • 为“驯服”异常 token，引入自适应 token 缩放/门控机制，在保持关键信息的同时抑制极端幅度。
  • 该方法仅需极小的额外计算量，即可无缝插入现有 Transformer 模块，几乎不增加延迟。
  • 实验表明，使用该技术后 FID、生成细节和采样步骤数均有显著提升，尤其在大规模模型上效果更明显。

引用

注：文中事实性信息以以上引用为准；观点与推断为 AI Stack 的分析。

站内链接

相关文章