研究 / 官方
苹果研究团队提出双阶段寄存器方法，有效抑制扩散Transformer中的异常令牌
苹果机器学习研究团队发表论文，系统研究了扩散Transformer（DiT）图像生成模型中的异常令牌现象。研究发现，在基于表示自编码器（RAE）与DiT组合的现代生成管线中，预训练视觉Transformer编码器会产生高范数异常表示，DiT自身在中间层也会生成内部异常令牌。这些异常令牌吸引了模型不成比例的注意力，却携带有限的局部语义信息，从而损害生成图像的质量。值得注意的是，简单屏蔽高范数令牌并不能改善效果，说明问题根源在于局部图像块语义的损坏，而非单纯的极端数值。为此，研究团队提出双阶段寄存器（DSR）方法，分别针对编码器和去噪器进行干预：在有预训练寄存器时使用训练寄存器，否则采用递归测试时寄存器，并为去噪器引入扩散寄存器。在ImageNet及大规模文本到图像生成任务上的实验均验证了该方法的有效性。
视觉Transformer（ViT）中存在一类特殊的高范数令牌，它们在注意力机制中占据主导地位，却几乎不携带有效的局部图像信息。这一现象此前已在判别式视觉模型中被观察到，但其在生成模型中的作用与影响长期缺乏系统研究。苹果机器学习研究团队的这项工作首次将目光聚焦于扩散Transformer的生成场景，填补了这一研究空白。
研究团队发现，异常令牌问题在现代RAE-DiT生成管线的两个关键环节均有出现。一方面，用于将图像压缩为潜在表示的预训练ViT编码器会产生异常的高范数表示；另一方面，DiT去噪网络本身在中间层也会自发形成内部异常令牌。这意味着该问题贯穿整个生成流程，需要在多个阶段同时加以干预。
一个直觉上的解决思路是直接屏蔽或过滤掉这些高范数令牌，但实验结果表明这种做法并不奏效。研究团队指出，问题的本质不在于少数极端数值本身，而在于这些令牌所对应的局部图像块语义已经遭到破坏。因此，有效的解决方案需要从根本上修复语义表示，而非简单地移除异常值。
为此，研究团队提出了双阶段寄存器（Dual-Stage Registers，DSR）方法。该方法借鉴了寄存器令牌的思路，分别为编码器和去噪器设计了针对性的干预策略：当存在可用的预训练寄存器时，直接使用训练好的寄存器；在测试阶段无法获取时，则采用递归方式动态生成测试时寄存器；针对DiT去噪器，则专门引入扩散寄存器机制。
在ImageNet图像生成基准以及大规模文本到图像生成任务上，DSR方法均取得了一致的性能提升，有效减少了生成图像中的异常伪影，提高了整体生成质量。研究团队认为，对异常令牌的有效控制是构建更强大扩散Transformer模型的重要组成部分，这一发现有望为后续生成模型架构设计提供新的优化方向。
要点
- 异常高范数令牌问题在RAE-DiT生成管线的编码器和去噪器两个环节均存在，需要端到端的系统性解决方案
- 直接屏蔽高范数令牌无法改善生成质量，问题根源在于局部图像块语义损坏而非极端数值本身
- 双阶段寄存器（DSR）方法通过训练寄存器、测试时递归寄存器和扩散寄存器三种机制，分场景灵活干预异常令牌
- 该研究将异常令牌控制确立为提升扩散Transformer生成质量的关键技术要素，对生成模型架构设计具有重要参考价值
原始标题：Taming Outlier Tokens in Diffusion Transformers
本文由 DataHub 基于公开来源整理，用于信息发现与摘要阅读；具体事实、数据和后续更新以原始来源为准。