像素空间扩散Transformer中的寄存器至关重要

摘要

已知视觉Transformer（ViTs）会出现高范数的异常块令牌，导致特征图质量下降，而注册令牌能有效缓解这一问题。随着扩散模型越来越多采用Transformer架构并转向像素空间训练，其形态日益接近ViT，这引出一个问题：注册令牌对扩散Transformer（DiTs）是否同样有用。本研究发现，DiTs与ViTs有一个关键差异：DiTs虽不产生异常块令牌，但仍能从注册令牌中受益。有趣的是，注册令牌在像素空间DiTs中的效果优于潜在空间DiTs。通过对中间表征的分析，发现注册令牌在高噪声水平下能生成更干净的特征图，这可能有助于提升像素空间生成的效果。进一步观察到，近期像素空间DiT架构已隐式纳入类似注册令牌的机制，这或许部分解释了其出色的实证表现。基于这些发现，我们提出"注册引导"技术，通过放大负责改善视觉结构与连贯性的注册令牌贡献，来提升生成质量。

English

Vision Transformers (ViTs) are known to exhibit high-norm patch-token outliers that degrade feature map quality, a problem effectively mitigated by register tokens. As diffusion models increasingly adopt transformer architectures and move toward pixel-space training, they become closer in form to ViTs, raising the question of whether register tokens are also useful for Diffusion Transformers (DiTs). In this work, we show that DiTs differ from ViTs in a key respect: they do not exhibit patch-token outliers but still benefit from registers. Interestingly, registers are more effective in pixel-space DiTs than in latent-space DiTs. By analyzing intermediate representations, we find that register tokens produce cleaner feature maps at high noise levels, which may contribute to their effectiveness in pixel-space generation. We further observe that recent pixel-space DiT architectures implicitly incorporate register-like mechanisms, which may partially account for their strong empirical performance. Motivated by these observations, we propose Register Guidance, a technique that amplifies the contribution of register tokens responsible for improving visual structure and coherence.