Diffusion Transformers (DiTs) for image generation suffer from outlier tokens—high-norm tokens that attract excessive attention while carrying limited local information. This phenomenon, previously observed in Vision Transformers, appears in both the encoder and denoiser of modern Representation Autoencoder (RAE)-DiT pipelines. The authors show that simple masking of high-norm tokens does not improve performance, suggesting the issue stems from corrupted local patch semantics rather than extreme values alone.

To address this, they introduce Dual-Stage Registers (DSR), a register-based intervention targeting both components. DSR uses trained registers when available, recursive test-time registers otherwise, and diffusion registers for the denoiser. The method is evaluated on ImageNet class-conditional generation and large-scale text-to-image generation. Results consistently show reduced outlier artifacts and improved generation quality. The authors highlight that outlier-token control is a crucial but previously overlooked component in building effective DiTs.