驯服离群点：DSR 机制如何提升 Diffusion Transformer 的生成质量

TL;DR

在 Diffusion Transformers (DiTs) 的演进中，我们往往关注规模（Scaling）与架构，却忽视了底层 Token 的“健康度”。本文揭示了 DiT 模型中普遍存在的 Outlier Tokens（离群点标记） 现象——这些高范数 Token 虽吸引大量注意力却破坏了局部语义。作者提出的 Dual-Stage Registers (DSR) 通过在编码器和生成器中嵌入“寄存器”标记，成功吸收了这些伪影，使得 ImageNet 生成精度大幅提升，收敛速度加快 4 倍。

痛点深挖：为何离群点是生成任务的毒药？

过去的研究发现，Vision Transformers 在处理图像时会产生极少数范数极大的 Token。这些 Token 就像黑洞，吸干了 Softmax 注意力图中的权重。

在理解任务（如分类）中，这种现象可能被视为一种特征压缩；但在生成任务中，它是一场灾难。作者发现：

• 污染表征空间：离群点通常出现在编码器的末层和 DiT 的中间层。
• 破坏局部语义：这些离群点并不携带有效的 Patch 信息，直接导致生成的图像块出现逻辑断层或伪影。
• 掩蔽无效：作者尝试直接在训练时屏蔽掉这些高范数 Token 的损失，结果发现性能并无提升。这说明问题不在于“数值太大”，而在于“原本该在那里的局部语义已经丢失了”。

核心内容：Dual-Stage Registers (DSR) 架构

作者的直觉非常明确：既然这些离群点是由于模型需要一个“地方”来存放全局或无关信息而产生的，那我们就主动给它提供这个“地方”。

1. 编码器端的“救治” (Registers in Encoders)

对于像 SigLIP2 这样没有预训练寄存器的模型，作者引入了 Recursive Test-Time Registers (TTR)。即便不重新训练编码器，只需在推理时插入额外的 Token，也能有效缓解输出表征的离群点现象。

2. 扩散器端的“防护” (Registers in Diffusion Transformers)

在扩散模型内部，作者加入了可学习的 Diffusion Registers。这些 Token 与图像 Patch 一起输入 Transformer，充当“注意力汇”（Attention Sink），专门吸收那些可能导致数值不稳定的全局干扰。

上图清楚地显示，如果不加寄存器，中间层（Middle Layers）会出现明显的离群点（深红色）；引入 DSR 后，范数分布变得异常均匀。

实验与结果

DSR 的效果在多种主流 DiT 变体上得到了验证：

• 性能暴涨：在 RAE-DiT 框架下，ImageNet-256 上的 FID 从 5.89 显著降至 4.58。
• 效率提升：达到同等生成水平所需的训练 Epoch 仅为基线的 1/4，极大节省了算力。
• 泛化性极强：无论是基于 VAE 的 DiT，还是直接基于像素空间的 JiT，甚至是最近大火的 Scale-RAE 文本生成模型，加入寄存器后均表现更佳。

深度洞察：为什么不是越早加越好？

一个有趣的消融实验显示，寄存器的插入位置有讲究。作者发现从第 8 个 Block 开始插入效果最佳（针对 XL 模型）。如果插入太早，可能会干扰模型对原始底层特征的提取；如果插入太晚，则离群点带来的破坏已经不可逆转。这提示我们，Transformer 的中间层是全局语义冲突最剧烈的“战场”。

总结与局限

Takeaway：DSR 是一种极其廉价（仅增加约 10% GFLOPs）但非常高效的“稳定剂”。它告诉我们，与其让模型自己去挣扎着处理那些不稳定的注意力分布，不如给它一个明确的缓存空间。

局限性：虽然 DSR 解决了特征层面的离群点，但在极大规模参数量下，Language Model 端（如文本 Prompt 编码器）的离群点问题依然存在，这可能是未来进一步提升文本-图像一致性的突破口。