驯服扩散 Transformer 中的离群标记

论文信息

标题: Taming Outlier Tokens in Diffusion Transformers

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

发布日期: 2026-05-06

PDF 链接: 下载 PDF

引言：当生成模型遭遇注意力异常

扩散模型已经成为图像生成领域的主导力量，特别是基于 Transformer 架构的扩散 Transformer（DiT）将自注意力的可扩展性与扩散过程的逐步去噪相结合，在 ImageNet 类条件生成和大规模文本到图像合成中取得了令人瞩目的成绩。然而，一个被长期忽略的细节正在悄然影响这些模型的生成质量 —— 异常 token（outlier tokens）。早在 Vision Transformer（ViT）用于图像分类时，研究者就发现网络中会自发产生极少数高范数的 token，它们在自注意力中攫取不成比例的注意力权重，自身却只包含非常有限的局部信息，甚至沦为一种“背景注册器”。这类异常 token 在判别模型中被证明会损害特征表达和迁移能力。然而，在生成式 DiT 的情境下，异常 token 的表现形式、成因以及影响尚缺乏系统探索。论文《Taming Outlier Tokens in Diffusion Transformers》正是填补这一空白的开创性工作，它揭示了高范数异常 token 在现代 RAE-DiT 图像生成管线中的普遍性，并指出性能下降的根源远非几个极端值那么简单，而是与局部块语义的系统性损坏紧密相关。本文将从动机、方法、实验与实践建议等角度对这项研究进行深入解析，为量化交易等技术背景的读者架起一座理解生成模型内部行为的桥梁。

异常 Token 现象：从 ViT 到 DiT 的延伸

在 ViT 中，输入图像被切分为若干固定大小的 patch，经线性投影后形成 token 序列，附加上一个可学习的分类 token [CLS]，共同输入 Transformer 编码器。研究早已发现，中间层的 [CLS] token 和某些 patch token 的范数会异常增大，远远超出其他 token。在自注意力计算中，这些高范数 token 的注意力分数极具统治力，使得信息聚合严重偏向它们，而真正的局部结构信息却得不到有效传递。不幸的是，这种“注意力垄断”在生成模型里也毫无保留地复现了。

论文观察到，当前主流的图像生成管线普遍采用 Representation Autoencoder（RAE）与 DiT 的组合：先用预训练的 ViT 编码器将图像压缩到低维潜在表示，再用 DiT 在该潜在空间进行去噪生成。这里的 ViT 编码器（例如在 Stable Diffusion 3 等模型中常见的架构）同样表现出了明显的异常 token 现象，它在输出端会产生少量高范数的潜在 token。更关键的是，后续 DiT 的去噪过程中，内部也会“二次孕育”出新的异常 token，尤其集中在中间层。这些 token 不再只是编码器引入的遗留问题，而是模型自身固有的倾向。也就是说，异常 token 犹如一条贯穿生成链条的暗线，从编码到去噪持续存在并相互影响。

不仅仅是极端值：语义损坏才是关键

一个直观的应对策略是直接检测并掩盖这些高范数 token，比如在推理时将其标量置零。但实验结果表明，简单的掩码不仅无法提升生成质量，甚至使 FID 等指标变差。这强烈暗示，异常 token 的存在并非仅仅由于数值尺度异常，不能用粗暴“剪枝”的方式解决。更深入地分析发现，问题实质是：异常 token 所对应的 patch 特征已经不再准确地反映原始图像局部块的语义信息。正是因为这种语义扭曲，注意力图被带偏，生成结果出现伪影、细节模糊或结构紊乱。因此，核心挑战在于如何在不丢失关键信息的前提下，修复这些被污染的局部表示，而非简单地压制或删除。

双阶段寄存器（DSR）：驯服编码器与去噪器

针对上述双重来源的异常 token，论文提出了一种优雅的解决方案 —— 双阶段寄存器（Dual-Stage Registers, DSR）。其核心思路受到 ViT 研究中“寄存器 token”（registers）的启发：在 ViT 训练时人为插入若干个可学习的寄存器 token，它们能吸收异常 token 的过大范数和冗余注意力，从而释放普通 patch token 的表达能力。DSR 将这一思想分别应用到编码器和去噪器两个阶段，但在编码器上区分了有无训练权限的场景。

训练好的寄存器与测试时递归寄存器

当可以访问并重新训练 ViT 编码器时，直接在编码器中添加一组可学习的寄存器 token，并与原始图像 patch token 一同训练。这些寄存器会被优化成为“异常吸收器”，有效抑制高范数 token 的出现，使编码器输出更健康的潜在表示。这适用于从头训练或微调编码器的场景。

然而，在多数工业级生成系统中，ViT 编码器通常是已预训练好且冻结的庞大模型，重新训练成本极高。针对这种更现实的设定，论文设计了测试时递归寄存器（Recursive Test-time Registers）。具体做法是：在编码器推理过程中，每一层都从层输入中动态构造一个寄存器 token，并让其参与该层的注意力计算。寄存器的值并非随意给定，而是通过一个递归公式从前一层传递而来，并结合当前层的统计信息更新。这个递归的寄存器 token 完全不需要任何梯度或额外的训练数据，它就像在每一层放置一个“归一化探针”，实时感知并中和异常偏置。这种巧妙的设计使得冻结的编码器也能输出相对稳定的潜在 token，为下游 DiT 提供更干净的起点。

扩散寄存器：为去噪过程量身定制

即便编码器端得到了治理，DiT 在内部多层去噪时仍可能自发产生异常 token。为此，DSR 在 DiT 侧引入了扩散寄存器（Diffusion Registers）。这些寄存器是一组可学习的 token，与带噪声的潜在 token 一同送入 DiT 的每一层。在训练时，扩散寄存器与模型其他参数同时端到端优化，学会在扩散步骤中持续吸收异常能量，防止去噪网络本身积累高范数 token。这种设计轻量且易于整合，不需要修改核心 DiT 架构，额外增加的参数量和计算开销都极小，却能在整个扩散链路上维持 token 表示的均衡性。

实验验证：从类条件到文本到图像生成

论文在 ImageNet 类条件图像生成和大规模文本到图像生成两个基准上全面评估了 DSR 的效果。在 ImageNet $256\times256$ 下，将 DSR 应用于 DiT-XL 模型，FID（Fréchet Inception Distance）等生成质量指标获得了一致改善，尤其是掩盖高范数 token 导致失真的案例在引入寄存器后得到了明显修复。生成图像中的结构性伪影减少，细节更锐利自然。

在大规模文本到图像生成实验中，采用了类似 Stable Diffusion 3 的 RAE-DiT 架构。当使用测试时递归寄存器处理冻结编码器、并用扩散寄存器优化去噪器后，无论是在图像清晰度、文本对齐度还是整体美学评分上，均观察到了稳定提升。论文还通过可视化展示了异常 token 的范数分布：原始模型在部分层存在极端高值的“尖峰”，而 DSR 版模型则将这些范数拉回到合理范围。更重要的是，注意力图也从高度集中于少数 token 变为更均匀地分布在真正有语义的 patch 区域，定性验证了语义修复的有效性。

对量化和生成系统的实践启示

这篇论文对相关技术领域的从业者具有直接指导意义。对于量化交易从业者而言，虽然直接操作的是金融时间序列，但可以类比地审视 Transformer 在时序预测或因子挖掘中是否存在类似的“异常 token”现象。高范数的时序嵌入可能对应着某些具有极端波动或异常跳跃的日期，它们同样可能垄断注意力并扭曲整个序列的表达。我们可以借鉴 DSR 的思想，在时序 Transformer 中加入可学习的寄存器 tokens，引导模型将异常波动吸收进寄存器，从而保护规整的时序模式。测试时递归寄存器的方法尤其适合无法重新训练预训练金融编码器的场景，只需在推理时轻量地插入递归寄存器，即可平稳抑制异常注意力。

对于人工智能工程和生成模型部署，DSR 提供了一套低成本、高收益的模型升级策略。当前工业级扩散模型大多依赖昂贵的冻结编码器，直接应用递归寄存器可以在不重新训练编码器的前提下，显著缓解生成伪影和结构失真。扩散寄存器则像一块即插即用的“健康插件”，在训练下一个版本的 DiT 时就可以无缝嵌入，几乎不需要额外工程适配。因此，在文生图、视频生成、3D 资产合成等产品的迭代中，DSR 可以作为一项标准组件被纳入模型设计蓝图。

此外，DSR 还为量化压缩和推理加速带来了间接利好。当 token 范数分布更均衡后，后量化（post-training quantization）和稀疏注意力机制遇到的离群值干扰会大幅降低。已有研究表明，ViT 中的异常 token 极大地恶化了量化误差，而 DSR 有意无意地解决了这一难题，可能使 DiT 模型更易于进行 INT8 甚至更低位宽的量化部署，这对资源受限的终端设备意义重大。

未来方向与挑战

尽管 DSR 取得了令人信服的效果，但仍有多个方向值得深入探索。第一，递归寄存器的设计目前是启发式的，是否存在理论最优的递归形式和初始化策略？能否通过元学习或自动搜索进一步提升其适应能力？第二，异常 token 的成因为它们提供了一个有趣的逆向视角：是否可以利用这类 token 进行模型水印、对抗攻击检测或生成控制？反常能量吸收机制或可转化为“可控异常”的创意工具。第三，DSR 目前主要针对 RAE-DiT 管线设计，在纯扩散模型、流匹配模型或多模态大模型中，异常 token 现象是否同样普遍？其跨架构的迁移研究将拓展方法的适用范围。最后，将寄存器技术与动态网络、专家混合（MoE）等范式结合，或能实现更智能的注意力资源调度，让生成模型在计算效率和生成质量之间达到新的平衡。

总结

《Taming Outlier Tokens in Diffusion Transformers》以其细致的现象剖析和实用的双阶段干预方案，为扩散 Transformer 的优化打开了一个全新的维度。它告诉我们，生成模型的“瑕疵”有时并非源于架构设计不足或训练数据匮乏，而可能仅仅是由于注意力机制内部小小的“失衡”累积而成。通过借鉴 ViT 领域的寄存器思想，并将其创造性地推广到编码器冻结场景和去噪器内部，DSR 用简洁优雅的方式实现了“四两拨千斤”的效果。这项工作不仅提高了图像生成的质量基准，更为所有关注 Transformer 表示健康度的研究者提供了一个可操作的范例。在迈向更强大、更可控的生成式人工智能的征途中，驯服异常 token 或许正是其中一个关键却长期被低估的环节。