Technical Summary: Mastering Outlier Tokens in Diffusion Transformers

Problem Statement

The work investigates the phenomenon of outlier tokens within Diffusion Transformers (DiTs) used for image generation. While prior work (e.g., Darcet et al.) demonstrated that Vision Transformers (ViTs) in recognition tasks frequently produce a small subset of tokens with abnormally large feature norms that absorb disproportionate attention while carrying limited local information, the role of these tokens in generative models remained unexplored.

The authors identify two critical gaps:

1. Ubiquity in Generative Pipelines: Outlier tokens occur not only in the pre-trained ViT encoders (tokenizers) of modern Representation Autoencoder (RAE) pipelines but also internally within the Diffusion Transformers (Denoisers) themselves.
2. Nature of the Defect: The degradation caused by these tokens is not merely the result of a few extreme loss values. Instead, the authors hypothesize that outlier tokens are symptoms of corrupted local patch semantics. Simply masking high-norm tokens during training yields no performance improvement, suggesting the problem lies in the loss of spatial structure and local information rather than just numerical instability.

Methodology: Dual-Stage Registers (DSR)

To address these issues, the authors propose Dual-Stage Registers (DSR), a unified intervention strategy that applies register tokens to both components of the RAE-DiT pipeline: the Representation Encoder and the Diffusion Generator.

1. Registers in Vision Encoders

The authors address outlier tokens in the encoder stage (the tokenizer) using two approaches, depending on the availability of pre-trained registers:

• Trained Registers: For encoders like DINOv2 that support registers, the authors utilize pre-trained register tokens to suppress high-norm artifacts in the encoder's feature maps.
• Recursive Test-Time Registers (TTR): For encoders without integrated registers (e.g., SigLIP2), the authors add additional untrained tokens to the input sequence at inference time. In cases where a second layer of outliers is detected in the resulting representation, they apply TTR recursively to stabilize the output.

2. Registers in Diffusion Transformers

The authors observe that even with a stabilized encoder, the Diffusion Transformer itself develops internal outlier tokens, which occur predominantly in middle layers (in contrast to standard ViTs, where outliers often appear in the final layers).

• Diffusion Registers: They introduce a small set of learnable register tokens directly into the Diffusion Transformer. These tokens are trained jointly with the generator to absorb attention sinks and stabilize the representation space.
• Inference Protocol: During inference, the outputs of all register tokens (from both the encoder and the diffusion model) are discarded; only the image token outputs are retained.

Main Contributions

• Discovery of Internal Outliers: The work demonstrates that outlier tokens are ubiquitous in DiTs, appearing in both the encoder and the denoiser, with a distinct pattern of concentration in the middle layers of the diffusion model.
• Rejection of the "Extreme Value" Hypothesis: Through loss-masking experiments, the authors show that simply excluding high-norm tokens from the training loss does not improve generation quality. This leads to the conclusion that the core problem is the corruption of local patch semantics, not merely the presence of extreme values.
• Dual-Stage Intervention: The proposal of DSR offers a simple, lightweight mechanism to stabilize the entire generation pipeline by addressing outliers at the source (encoder) and during the generation process (denoiser).
• Generalizability: The method proves effective across various input spaces (pixel space, VAE latencies, and representation spaces) and architectures (SiT, JiT, RAE-based designs).

Experimental Results

The authors evaluate DSR on class-conditional generation on ImageNet-1K (256×256) and large-scale text-to-image generation.

• ImageNet-1K Performance:
  • Using SigLIP2-B as the encoder, DSR reduces the FID from 5.89 to 4.58 (with 800 epochs) and improves GenEval scores on text-to-image tasks from 0.426 to 0.466.
  • Combined with DINOv2-B encoders, the FID improves from 4.16 to 3.95.
  • DSR accelerates convergence, achieving comparable quality with 4× fewer training epochs compared to the baseline.
• Scalability: Improvements are consistent across various model sizes (DiT-S, B, XL), with only a marginal increase in computational cost (approximately 10% GFLOPs increase).
• Ablation Studies:
  • Insertion Depth: Registers are most effective when introduced in early to middle layers (e.g., starting at Block 8). Introducing them too early or too late yields diminishing returns.
  • Number of Registers: A moderate number of registers (e.g., 36) delivers optimal performance; too few (1) or too many (100) degrade performance.
• Comparison with Alternatives: DSR outperforms "In-Context Conditioning" (which uses input-dependent tokens) because the input representation becomes more semantic, suggesting that registers provide a more robust mechanism for global capacity.

Significance and Claims

The work claims that controlling outlier tokens is a critical, yet previously overlooked, component in building robust, Transformer-based diffusion pipelines. By demonstrating that these artifacts degrade local patch semantics and that register-based interventions can effectively mitigate them, the authors position DSR as a fundamental improvement for the stability of both encoder and generator.

The work highlights that while the flexibility of Transformers is advantageous for global interactions, it introduces specific artifacts (outlier tokens) that impair spatial accuracy during generation. The proposed solution is presented as a simple plug-and-play intervention that consistently improves generation quality across diverse settings without requiring complex architectural changes or a full retraining of the entire pipeline from scratch.