Taming Outlier Tokens in Diffusion Transformers

In image generation models based on Diffusion Transformer (DiT), it has been confirmed that an outlier token phenomenon occurs where some tokens have excessively large norms and attract excessive attention. These tokens fail to capture sufficient local patch information, potentially causing artifacts during the generation process.

Researchers analyzed the ViT encoder and DiT denoiser in the latest Representation Autoencoder (RAE)-DiT pipeline. They observed that outlier tokens form not only in the pre-trained ViT encoder but also inside the DiT, particularly in the middle layers.

Simply masking tokens with large norms did not improve performance. The analysis suggests that the core issue lies closer to damaged local patch semantic information rather than the extreme values themselves.

To address this, they proposed the Dual-Stage Registers (DSR) intervention technique. DSR applies the following registers to each component:

• Learned registers when available
• Registers generated recursively during testing
• Diffusion registers applied to the DiT denoiser

In ImageNet and large-scale text-to-image generation experiments, DSR intervention consistently reduced outlier artifacts and improved generation quality. The researchers explained that controlling outlier tokens is a crucial factor for building more robust DiTs.