驯服扩散 Transformer 中的离群词元

Taming Outlier Tokens in Diffusion Transformers

arXiv: 2605.05206v1

论文信息

标题: Taming Outlier Tokens in Diffusion Transformers

作者: Xiaoyu Wu, Yifei Wang, Tsu-Jui Fu, et al.

发布日期: 2026-05-06

arXiv ID: 2605.05206v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题：扩散 Transformer（DiT）图像生成模型中出现高范数异常令牌，破坏局部语义，影响生成质量，但该问题在生成模型中长期被忽视。
  • 核心方法：提出双阶段寄存器（Dual‑Stage Registers, DSR），在视觉编码器与扩散去噪器两端分别插入寄存器令牌：编码器端使用训练好的寄存器或递归测试时寄存器，扩散器端加入可训练的扩散寄存器。
  • 关键结果：在 ImageNet-256 上，RAE‑DiT + SigLIP2‑B 的 FID 从 5.89 降至 4.58（表 3 和表 6），GenEval 从 42.6 提升至 46.6（表 8），且收敛速度快 4 倍（图 7）。
  • 主要局限：DSR 依赖编码器端提前获取的寄存器（训练或测试时），对未公开训练注册的编码器需要递归测试时策略，额外计算开销约增加 10% FLOPs（表 9）。在大规模文本‑图像任务中，扩散损失改善稳定，但语言模型损失出现尖峰，可能影响复杂指令遵循（附录 D）。
  • 适合读者：从事扩散模型、视觉 Transformer、生成质量优化的研究者与工程师，尤其关注 DiT 训练不稳定性和语义保持的人群。

论文背景和研究动机

视觉 Transformer 在图像理解中占据主导地位，近年更是进入生成领域，衍生出 DiT、RAE 等架构，逐步取代卷积 VAE 和 UNet。然而 Transformer 的全局注意力机制带来一个隐性问题：少量令牌出现异常高的 ℓ2\ell_2 范数，吸引大量注意力，却携带极少的局部语义。该现象在 DINO、CLIP 等视觉模型中已被证实，引入额外可学习寄存器令牌可以有效缓解，但在扩散生成模型中尚无系统研究。

论文观察到，RAE‑DiT 流水线中，无论预训练视觉编码器（如 SigLIP2 或 DINOv2）还是扩散去噪器本身，都会产生异常令牌，且其分布不同于理解模型——扩散中的异常令牌出现在中间层，而非最后几层（图 2）。更关键的是，简单的损失掩码（去掉范数过大的令牌的损失监督）并不能提升生成质量（表 1），说明问题并非仅由几个极端损失值驱动，而是深层局部语义被破坏。这推动作者提出：异常令牌是局部块语义退化的症状，而非单纯极端值。

这一见解促使作者不再追求极端值抑制，而是通过寄存器介入恢复块级语义，从而设计出覆盖编码器和扩散去噪器的双阶段寄存器方案。

核心方法和技术细节

异常令牌的实质：损坏的局部语义

在 RAE‑DiT 中，训练目标为 vv‑预测损失：

Lv(θ)=Ez0,ϵ,t[∑i=1N∥v^θ(zt,t,c)i−vi∥22]\mathcal{L}_v(\theta)=\mathbb{E}_{z_0,\epsilon,t}\left[\sum_{i=1}^N \| \hat{v}_\theta(z_t,t,c)_i - v_i \|_2^2\right]

作者尝试按令牌范数进行掩码：

mi=1 ⁣[∥z0,i∥2≤τ]m_i = \mathbf{1}\!\left[\|z_{0,i}\|_2 \le \tau\right] Lvmask(θ)=E[1∑imi∑i=1Nmi∥v^θ(zt,t,c)i−vi∥22]\mathcal{L}_v^{\text{mask}} (\theta) = \mathbb{E}\left[\frac{1}{\sum_i m_i}\sum_{i=1}^N m_i \| \hat{v}_\theta(z_t,t,c)_i - v_i \|_2^2\right]

即使仅过滤掉 0.1% 的令牌，FID 反而从 5.89 升至 6.06（表 1），表明直接将异常令牌排除出损失无法恢复生成质量，因为这些令牌所处位置的局部信息已经缺失。这与先前发现一致：空间结构对生成质量至关重要，而高范数异常令牌往往导致密集预测任务失败。

双阶段寄存器 DSR

编码器端寄存器。若编码器在预训练时已包含训练过的寄存器（如 DINOv2 with registers），直接使用即可显著减少异常令牌并提升生成质量（表 2：FID 从 4.16 降至 3.95）。对于像 SigLIP2 这样没有训练寄存器的编码器，论文采用递归测试时寄存器（TTR）：先在第一阶段插入测试时令牌稳定输出，若仍存在第二簇异常，则再次应用 TTR，有效抑制两组分离的异常源（图 8，附录 A）。

扩散端寄存器。即使在编码器端加入寄存器，扩散 Transformer 内部仍会在中间层出现高范数令牌（图 5）。受注意力沉陷假设启发，作者在扩散模型的输入序列中加入少量可训练寄存器令牌（默认起始于第 8 块，数量为 36）。这些寄存器在训练时与模型参数联合学习，推理时丢弃，不改变输出形状。结果内部异常几乎消失（图 9），中间层语义在 PCA 可视化中更清晰（图 6）。

整个 DSR 框架统一将寄存器应用于编码器和扩散器，吸收异常行为，从而保护了块级语义和空间结构。

创新点和贡献

  1. 首次揭示扩散 Transformer 中的异常令牌现象，并指出其分布（中间层）与理解模型（末层）不同。这一发现填补了生成模型领域的研究空白。
  2. 提供明确的证据表明异常令牌源于局部语义损坏，而不是极端损失值。损失掩码实验直接证明，仅仅过滤高范数令牌不能解决问题，需要恢复块语义信息。
  3. 提出双阶段寄存器 DSR，一种轻量、统一的干预措施。它涵盖编码器和扩散去噪器，并适配不同编码器（DINOv2、SigLIP2）和多种输入空间（像素、VAE 潜变量、表示空间），无需重新训练整个流水线即可减少异常令牌。
  4. 证明寄存器优于上下文条件化。与在上下文窗口中加入图像相关令牌的方案相比，扩散寄存器在语义表示空间中更为稳健（表 4），而上下文条件化在表示空间更结构化时收益下降。
  5. 在 ImageNet 和文本‑图像任务上验证了 DSR 的可扩展性。从 DiT‑S 到 XL，FID 均有改善，同时仅增加约 10% 的 GFLOPs（表 7）。在大规模 T2I 上，DSR 在训练仅 1 个 epoch、使用规模为原数据集四分之一的条件下，仍获得 46.6 的 GenEval 分数，超越基线 42.6。

实验结果分析

ImageNet 类条件生成

主实验使用 RAE‑DiT‑XL + SigLIP2‑B，训练 80 个 epoch。添加 DSR 后，无引导 FID 从 5.89 降至 4.58，IS 从 156.54 升至 165.99，精度提升而召回基本持平（表 3）。进一步加入 DDH 解码头后，FID 降至 3.56（80 epoch），并在 800 epoch 时达到 2.72（表 6）。收敛速度方面，DSR 在 80 epoch 就达到基线 320 epoch 的水平（图 7），说明寄存器介入加速了训练。

对扩散寄存器的消融（表 9）显示，起始块 8 且 36 个寄存器为最佳配置，过早或过晚插入，以及过少或过多的寄存器都会削弱效果，表明寄存器需要适度的容量和计算参与范围。

跨输入空间和模型规模

DSR 在多种 DiT 变体上均有效：VAE‑SiT‑XL 的 FID 从 16.05 降至 14.47；JIT‑H 从 30.34 降至 23.14（表 3）。缩放实验中，DiT‑B/XL 的改善幅度在 10‑30% 量级（表 7），证实干预的通用性。

文本‑图像生成

在 Scale‑RAE 上利用 Qwen2.5 1.5B + DiT 2.4B 进行训练，DSR 的扩散损失更稳定且整体更优，但语言模型损失出现一些尖峰（图 10），这解释了为什么 DSR 在复杂指令遵循的 DPG‑Bench 上相对提升较小（75.4 vs 74.3）。尽管如此，GenEval 上 DSR 明显占优（46.6 vs 42.6），证明其对对象级对齐有积极影响。

实践建议

当前主流扩散模型大都采用 Transformer 架构（如 DiT‑XL/DiT‑H），视觉编码器也广泛使用 ViT 或 SigLIP 提取特征。在实际部署或优化时，可考虑采用以下基于 DSR 的方法提升生成质量和训练效率：

  1. 编码器侧的寄存器添加

    • 若使用已自带训练寄存器的模型（如 DINOv2‑reg），直接使用其输出并丢弃寄存器令牌，即能降低异常令牌，无需额外成本。
    • 若使用 SigLIP2 或其他无寄存器模型，可在推理时添加测试时寄存器令牌，采用递归方法移除多个异常簇。注意，这会轻微增加推理时间，但论文表明无需重训解码器（附录 C），因为编码器分布偏移很小。

  2. 扩散器侧的寄存器训练

    • 对于 DiT 训练，从第 8 个 Transformer 块开始引入 36 个可学习寄存器令牌，训练后丢弃。调参时可围绕块数和寄存器数微调，但论文给出的默认值直接在多数场景适用。计算开销约增加 10% FLOPs（表 9），在训练阶段可接受，尤其考虑到收敛速度提升（4× 加速）带来的总体算力节省。
    • 在文本‑图像任务中，如果最终生成质量已满意但遵循复杂指令不足，可能需要单独优化语言分支；DSR 主要改善扩散损失，不直接强化语言指导。

  3. 与上下文条件化的权衡

    • 若已在用上下文条件化（输入额外上下文令牌），DSR 可与它叠加，但实验表明当输入表示偏语义（如 SigLIP2 特征）时，寄存器能带来更稳定的附加收益，而上下文条件化收益递减。因此，寄存器可用作基础增强，上下文条件化作为可选辅助。

  4. 监控异常令牌

    • 训练 DiT 时，可监控中间层的令牌范数分布（例如超过中位数 2 倍的比例），若发现异常比例高于 5-10%，可考虑采用 DSR，而不是简单丢弃或掩码损失。损失掩码无效性（表 1）说明必须从语义层面干预。

  5. 项目集成

    • 若已有基于 RAE‑DiT 或类似流水线的训练代码，集成 DSR 只需在编码器输出后和扩散输入前增加寄存器插入逻辑，并在扩散模块初始化时增加可学习寄存器参数。推理时移除即可。由于不改变主干结构，可快速验证改善效果。

综上，DSR 是一种简单而有效的技术，能显著提升基于 Transformer 的扩散模型的生成质量与训练效率，值得在工业级图像生成管线中尝试应用。