Technical Summary: Registers Matter for Pixel-Space Diffusion Transformers

Problem Statement

Vision Transformers (ViTs) are known to exhibit high-norm patch-token outliers that degrade feature map quality and create artifacts in attention maps. This issue is effectively mitigated in ViTs by introducing "register tokens," which act as sinks for these outliers. As Diffusion Transformers (DiTs) increasingly adopt transformer architectures and shift toward training directly in pixel space (moving away from latent space), they structurally resemble ViTs more closely. This raises two critical questions:

1. Do DiTs inherit the high-norm patch-token outliers observed in ViTs?
2. Are register tokens beneficial for DiTs, particularly in pixel-space training regimes?

While attention sinks and special-token behaviors have been studied in Large Language Models (LLMs) and diffusion language models, their presence and functional role in image diffusion transformers remain underexplored.

Methodology

The authors conduct a systematic analysis of register tokens across different diffusion model spaces (pixel-space, VAE-space, and RAE-space) and compare them against standard ViTs (specifically DINOv2).

1. Empirical Analysis of Outliers: The study investigates token-wise feature norms and attention maps in DiTs (specifically pDiTs trained with flow matching on ImageNet) versus ViTs. They analyze whether patch tokens exhibit high-norm outliers in the absence of registers.
2. Register Token Integration: Register tokens are implemented as additional learnable tokens appended to the patch-token sequence, following the approach used in ViTs. These tokens are excluded from the training objective (loss function).
3. Internal Representation Analysis: The authors analyze intermediate feature maps using Total Variation (TV) to measure spatial smoothness and linear probing to evaluate the semantic content of register tokens. They examine how registers affect feature norms and the coherence of representations at different diffusion timesteps (noise levels).
4. Architectural Exploration: Motivated by the distinct roles of registers and patch tokens, the paper proposes a parameter-efficient dual-stream architecture. This design processes register and patch tokens through specialized pathways (specifically decoupling RMSNorm, MLP, and adaLN layers) while sharing attention mechanisms and most transformer parameters.
5. Ablation Studies: The study ablates the number of registers, the depth at which they are introduced, and the specific components to be dualized to determine optimal configurations.

Key Findings and Contributions

1. DiTs Differ Fundamentally from ViTs Regarding Outliers

Contrary to ViTs, standard DiTs (both latent and pixel-space) do not exhibit high-norm patch-token outliers. In ViTs, outliers are localized to specific tokens (often in low-information regions), whereas in DiTs, patch-token norms remain relatively uniform, and attention maps lack the low-information artifacts seen in ViTs.

2. Registers Improve Pixel-Space DiTs Despite Absence of Pre-existing Outliers

Surprisingly, adding register tokens to DiTs significantly improves convergence and generation quality, particularly in pixel-space models.

• Mechanism: Unlike in ViTs where registers absorb pre-existing outliers, registers in DiTs create high-norm tokens within themselves.
• Effect: The presence of registers consistently reduces the feature norms of all patch tokens. This leads to cleaner, smoother intermediate feature maps, especially at high noise levels (early diffusion timesteps).
• Dual Role: Register tokens specialize into two distinct roles:
  • Norm Sinks: Tokens with extremely high norms that absorb magnitude from patch tokens.
  • Global Information Carriers: Tokens with moderate norms that encode diverse semantic information about the image (e.g., foreground objects, background context), verified via linear probing and attention map visualization.

3. Space-Specific Efficacy

The benefits of register tokens are highly dependent on the training space:

• Pixel-Space: Registers yield the most significant improvements in generation quality (FID). Pixel-space models inherently produce noisier intermediate representations with higher feature norms, making the regularization effect of registers crucial.
• Latent-Space (VAE/RAE): Registers provide moderate gains or, in the case of RAE-based models, slightly degrade performance. Latent spaces are more structured and lower-dimensional, reducing the need for the strong regularization provided by registers.

4. Implicit Registers in Existing Architectures

The authors observe that recent pixel-space DiT architectures (e.g., JiT) that use in-context class conditioning (appending class embeddings to the token sequence) implicitly incorporate register-like mechanisms. These in-context tokens exhibit similar behavior to explicit registers: some act as norm sinks, while others encode global semantic information. This suggests that performance gains in such models may stem largely from this register-like behavior rather than the class information itself.

5. Parameter-Efficient Dual-Stream Architecture

To leverage these insights, the authors propose a compact dual-stream architecture that specializes the processing of register tokens.

• Design: Registers are introduced only in deeper layers (starting from block 4), where semantic structure is established. The architecture uses shared attention but decouples RMSNorm, MLP, and adaLN for registers and patches.
• Efficiency: This approach increases the model size by only ~14% (compared to ~77% for naive full duplication) while maintaining similar inference costs (GFLOPs).
• Performance: This architecture consistently improves generation quality across model scales (B, L, H) and resolutions (256x256, 512x512) compared to single-stream baselines.

Results

• Generation Quality: On ImageNet 256x256, pDiTs with registers consistently outperform those without. For example, a pDiT-B/16 model improved its FID from 7.39 to 5.30 at 200 epochs.
• Feature Smoothness: Models with registers show lower Total Variation (TV) ratios in intermediate features at high noise levels, indicating smoother and more coherent representations.
• Ablation: Registers are most effective when introduced after the first 4 layers and when using a larger number of tokens (e.g., 32 registers for pDiT-B, compared to 4 for ViTs).

Significance and Claims

The paper claims that register tokens are a critical, yet previously overlooked, component for high-quality pixel-space diffusion generation. The significance lies in:

1. Redefining the Role of Registers: Demonstrating that registers are not merely a fix for ViT-specific outliers but serve a broader regularization function in diffusion models by managing feature norms and enhancing spatial coherence.
2. Explaining Recent Successes: Providing a rationale for the strong empirical performance of recent pixel-space DiTs that utilize in-context conditioning, attributing part of their success to implicit register behavior.
3. Efficient Architecture Design: Offering a parameter-efficient dual-stream design that capitalizes on the distinct roles of registers and patches, achieving state-of-the-art generation quality with negligible runtime overhead.

The authors conclude that while their designs capture the promise of token specialization, the full potential of utilizing registers in pixel-space diffusion transformers remains an open question for future research.