Taming the Outliers: Dual-Stage Registers for Robust Diffusion Transformers
Taming Outlier Tokens in Diffusion Transformers
The paper introduces Dual-Stage Registers (DSR), a method to mitigate "outlier tokens"—abnormally high-norm tokens that degrade local patch semantics—in Diffusion Transformers (DiTs). By inserting register tokens into both the vision encoder (tokenizer) and the diffusion denoiser, the authors achieve SOTA-level improvements in image generation quality across ImageNet and large-scale T2I benchmarks.
TL;DR
High-norm "outlier tokens" are parasitic artifacts in Transformers that steal attention and corrupt local data. This paper identifies that these outliers plague both the encoders and the denoisers of modern Diffusion Transformers (DiTs). The solution is Dual-Stage Registers (DSR): adding dedicated "sink" tokens that soak up these artifacts, leading to much cleaner image generation and significantly faster training convergence.
The Problem: When Parasitic Tokens Corrupt Art
In the world of Vision Transformers (ViTs), "outlier tokens" are well-documented. They have massive numerical norms and absorb huge amounts of attention, yet they contain almost no useful semantic information.
The authors discovered a striking problem in RAE-DiT (Representation Autoencoder-Diffusion Transformer) pipelines:
- Encoder Corruption: Pretrained encoders (like SigLIP2) produce outlier representations that serve as "poisoned" input for the diffusion model.
- Denoiser Emergence: Even if the input is clean, the Diffusion Transformer itself develops outliers. Interestingly, unlike recognition ViTs where outliers appear at the end, in DiTs, they peak in the intermediate layers.
Crucially, simply masking these high-norm tokens during training doesn't work (see Table 1). This proves the problem isn't just about "big numbers"—it's about the loss of local spatial structure.
Methodology: The DSR Framework
To fix this, the authors propose Dual-Stage Registers (DSR), a two-pronged intervention:
Stage 1: Cleaning the Encoder
If the encoder has pre-trained registers (like DINOv2), they use them. If it doesn't (like SigLIP2), they apply Recursive Test-Time Registers (TTR). By appending an extra token during inference, they provide a "sink" for the encoder's artifacts before they ever reach the diffusion model.
Stage 2: Guarding the Denoiser
The authors add ~36 learnable diffusion registers into the DiT blocks (specifically starting from block 8). These registers stay with the sequence during denoising, absorbing the global context that would otherwise "explode" into an outlier patch.
Visual Evidence: Cleaning the Feature Map
The impact of DSR is most visible when looking at Norm Maps and PCA visualizations. Without registers, the feature maps are riddled with "bright spots" (outliers) that represent corrupted patch information. With DSR, the maps become smooth and semantically coherent.
Performance: Better, Faster, Stronger
The results across benchmarks are consistent:
- ImageNet-256: FID improved significantly (e.g., 5.89 → 4.58).
- Training Efficiency: DSR allows the model to reach the same quality milestones as the baseline while using 4x fewer training epochs.
- Text-to-Image: On benchmarks like GenEval, DSR showed a clear lead in following complex prompts.
Critical Insight & Conclusion
Why does this work? The authors hypothesize that Transformers naturally want to aggregate global information. Without registers, they are forced to "recruit" standard image patches to hold this global state, which destroys that patch's local meaning. By providing dedicated "register" slots, we let the image patches focus on being pixels and the registers focus on being the global "sink."
Takeaway: As we scale DiTs to larger resolutions and more complex text-to-image tasks, managing outlier tokens isn't just an optimization trick—it's a fundamental requirement for maintaining spatial fidelity.