发现并解决扩散模型中异常激活的令牌问题,提升图像生成质量。
Taming Outlier Tokens in Diffusion Transformers

- 引入双阶段寄存器干预机制,动态修正异常令牌
- 在ImageNet和文生图任务上显著减少生成伪影
- 适合关注生成模型稳定性和细节质量的研究者
我们研究了扩散Transformer(DiT)在图像生成中异常令牌的现象。已有研究表明视觉Transformer(ViT)会产生少量高范数令牌,吸引过多注意力但携带信息有限,但在生成模型中的作用仍不明确。本文发现这一现象出现在现代表示自编码器(RAE)-DiT流水线的编码器与去噪器中:预训练的ViT编码器会产生异常表示,而DiT自身在中间层也会发展出内部异常令牌。单纯遮蔽高范数令牌无法提升性能,说明问题并非仅由极端值引起,更与局部补丁语义损坏有关。为此,我们提出双阶段寄存器(DSR),针对两个组件分别采用训练好的寄存器或测试时递归生成的寄存器,并为去噪器设计扩散寄存器。在ImageNet及大规模文生图任务上,该方法持续降低异常伪影并提升生成质量。结果表明,控制异常令牌是构建更强DiT的重要一环。
原文摘要 · Abstract (English)
We study outlier tokens in Diffusion Transformers (DiTs) for image generation. Prior work has shown that Vision Transformers (ViTs) can produce a small number of high-norm tokens that attract disproportionate attention while carrying limited local information, but their role in generative models remains underexplored. We show that this phenomenon appears in both the encoder and denoiser of modern Representation Autoencoder (RAE)-DiT pipelines: pretrained ViT encoders can produce outlier representations, and DiTs themselves can develop internal outlier tokens, especially in intermediate layers. Moreover, simply masking high-norm tokens does not improve performance, indicating that the problem is not only caused by a few extreme values, but is more closely related to corrupted local patch semantics. To address this issue, we introduce Dual-Stage Registers (DSR), a register-based intervention for both components: trained registers when available, recursive test-time registers otherwise, and diffusion registers for the denoiser. Across ImageNet and large-scale text-to-image generation, these interventions consistently reduce outlier artifacts and improve generation quality. Our results highlight outlier-token control as an important ingredient in building stronger DiTs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。