[Paper] 驯服扩散Transformer中的异常Token

Source: arXiv - 2605.05206v1

概述

论文 “Taming Outlier Tokens in Diffusion Transformers” 揭示了使用视觉Transformer（ViTs）的现代扩散式图像生成器中隐藏的缺陷。研究表明，编码器和去噪Transformer都可能产生少量“异常”token——这些向量幅度异常大，主导注意力却几乎不携带有用的视觉信息。作者通过引入一种轻量级的基于寄存器的修复方法——Dual‑Stage Registers (DSR)，显著降低了此类伪影，并提升了在ImageNet以及大规模文本到图像模型上的生成质量。

关键贡献

• 识别扩散Transformer（DiTs）中的异常标记。 证明高范数标记不仅出现在预训练的ViT编码器中，也会在扩散去噪的内部出现，尤其在中间层。
• 展示朴素掩码失效。 简单地将高范数标记置零并不能提升结果，表明问题在于语义腐败，而非仅仅是极端数值。
• 提出双阶段寄存器（DSR）。 一种两阶段、基于寄存器的干预方法：
  1. 训练时寄存器，在模型训练期间学习替换或纠正异常标记。
  2. 递归测试时寄存器，在推理过程中实时检测并替换异常标记，并为去噪器提供专用的扩散寄存器。
• 广泛的实证验证。 在标准的ImageNet生成任务和大规模文本到图像基准上，DSR始终能够降低视觉伪影并提升FID/IS分数。
• 开启新的研究方向。 强调异常标记控制是构建稳健的基于扩散的生成模型的关键且此前被忽视的组成部分。

方法论

1. 诊断问题

   • 作者首先分析 Representation Autoencoder‑DiT (RAE‑DiT) 编码器‑解码器流水线中的 token 范数。
   • 他们可视化注意力图，发现少数 token 主导了注意力分布，同时对应的是模糊或噪声的图块。

2. 基线实验

   • 采用简单的掩码（将超过范数阈值的 token 置零）和范数裁剪，结果显示对生成质量影响微乎其微，甚至出现负面效果。

3. 双阶段寄存器 (Dual‑Stage Registers, DSR)

   • 训练阶段寄存器： 小的可学习向量（即“寄存器”）被附加到 token 序列中。训练时，门控网络学习何时用寄存器条目替换异常 token，从而“修复”被破坏的语义。
   • 测试时寄存器： 推理阶段，递归检测模块扫描每一层的高范数 token，用最合适的寄存器条目替换它们，并将校正后的序列重新输入后续层。
   • 扩散寄存器： 专门为去噪 transformer 训练的一组寄存器，使其能够纠正由随机扩散过程本身产生的异常 token。

4. 评估

   • 在无条件 ImageNet 生成（256×256）以及大型文本到图像模型（如类似 Stable Diffusion 的架构）上测试该流水线。
   • 使用标准指标（FID、IS、CLIP‑Score）以及定性视觉检查来评估改进效果。

结果与发现

• 视觉质量： 使用 DSR 生成的样本出现的“斑块”或“棋盘格”伪影更少，这些伪影此前被追溯到异常 token。
• 注意力分布： DSR 后的注意力图更加平衡，在各 patch 之间的分布更平滑，证实寄存器成功稀释了异常 token 的主导作用。
• 效率： 寄存器模块在推理时间上仅增加 < 2 % 的开销，使其在实际部署中具有实用性。

Practical Implications

• Cleaner outputs for production‑grade generators. Companies building AI‑powered image creation tools (e.g., design assistants, content‑generation platforms) can integrate DSR to reduce glitchy artifacts without retraining the entire model.
• Improved downstream tasks. Better‑quality latent representations translate to higher fidelity in downstream pipelines such as image editing, in‑painting, or style transfer that rely on diffusion models.
• Low‑cost upgrade path. Since DSR works as a plug‑in (registers can be trained on top of an existing checkpoint), developers can retrofit legacy diffusion models with minimal compute budget.
• More stable fine‑tuning. When adapting a large diffusion model to a new domain (e.g., medical imaging), DSR can mitigate the emergence of outlier tokens that often cause training instability.

实际影响

• 为生产级生成器提供更干净的输出。 构建 AI 驱动图像创作工具的公司（例如设计助理、内容生成平台）可以集成 DSR，在无需重新训练整个模型的情况下减少故障伪影。
• 提升下游任务表现。 更高质量的潜在表示能够在依赖扩散模型的下游流水线（如图像编辑、图像修补或风格迁移）中实现更高的保真度。
• 低成本升级路径。 由于 DSR 作为插件工作（可在已有检查点之上训练寄存器），开发者可以以最小的计算预算为旧版扩散模型进行改造。
• 更稳定的微调。 在将大型扩散模型适配到新领域（例如医学影像）时，DSR 能缓解常导致训练不稳定的异常 token 的出现。

局限性与未来工作

• Scope of token types. 本研究聚焦于视觉 token；将分析扩展到多模态扩散模型（例如文本‑图像或视频）仍是未解之题。
• Register capacity. 固定的少量寄存器在极大或高度多样化的数据集上可能最终饱和；可以探索自适应或层次化的寄存器方案。
• Theoretical understanding. 虽然实证结果强劲，但对扩散动力学中为何会出现异常 token 的更深层理论解释仍然缺乏。
• Real‑time constraints. 尽管开销适中，超低延迟的应用场景（如移动端推理）可能需要进一步优化递归检测步骤。

Bottom line: 通过揭示扩散 transformer 中一个细微却普遍存在的问题——异常 token，本文为开发者提供了实用工具（DSR），帮助生成模型更可靠、视觉效果更佳，为更高质量的 AI 驱动内容创作铺平道路。

作者

• Xiaoyu Wu
• Yifei Wang
• Tsu-Jui Fu
• Liang-Chieh Chen
• Zhe Gan
• Chen Wei

论文信息

• arXiv ID: 2605.05206v1
• 类别: cs.CV, cs.AI, cs.LG
• 发布日期: 2026年5月6日
• PDF: 下载 PDF