
驯服扩散Transformer中的异常标记

在图像生成领域,扩散Transformer(DiT)是主流架构。但之前研究发现,Vision Transformer(ViT)会产生少量高范数(high-norm)的异常标记,这些标记吸引过多注意力,却携带有限局部信息。本文系统研究了这一现象在生成模型中的表现。
作者发现,在RAE-DiT流水线中,预训练ViT编码器和DiT去噪器都会产生异常标记,尤其在中间层。简单掩码高范数标记并不能改善,说明问题根源在于受损的局部补丁语义,而非少数极端值。
为此,提出Dual-Stage Registers(DSR),一种双阶段寄存器干预:对编码器使用训练好的寄存器或递归测试时寄存器,对去噪器使用扩散寄存器。在ImageNet和大规模文生图任务上,DSR一致减少了异常伪影,提升了生成质量。
该工作表明,控制异常标记是构建更强DiT的关键。


