Research from Apple and Rice University investigates outlier tokens in Diffusion Transformers (DiTs) used for image generation. These high-norm tokens attract disproportionate attention while carrying limited local information — a known issue in Vision Transformers (ViTs) that has been underexplored in generative models. The study finds outlier tokens appear in both the encoder and denoiser of RAE-DiT pipelines, and that simply masking them doesn't help because the root cause is corrupted local patch semantics. To address this, the authors propose Dual-Stage Registers (DSR): trained registers, recursive test-time registers, and diffusion registers for the denoiser. Experiments on ImageNet and large-scale text-to-image generation show consistent reductions in outlier artifacts and improved generation quality.

1m read timeFrom machinelearning.apple.com
Post cover image
3.4K Impressions