Registers Matter for Pixel-Space Diffusion Transformers

Abstract: Vision Transformers (ViTs) are known to exhibit high-norm patch-token outliers that degrade feature map quality, a problem effectively mitigated by \textit{register tokens}. As diffusion models increasingly adopt transformer architectures and move toward pixel-space training, they become closer in form to ViTs, raising the question of whether register tokens are also useful for Diffusion Transformers (DiTs). In this work, we show that DiTs differ from ViTs in a key respect: they do not exhibit patch-token outliers. Interestingly, register tokens significantly improve convergence and generation quality of pixel-space DiTs. By analyzing intermediate representations, we find that register tokens produce cleaner feature maps at high noise levels, which may contribute to their effectiveness in pixel-space generation. We further observe that recent pixel-space DiT architectures implicitly incorporate register-like mechanisms, which may partially account for their strong empirical performance. Motivated by these insights, we investigate a parameter-efficient dual-stream architecture that specializes processing for register tokens and improves pixel-space generation quality with negligible runtime overhead.

• The paper demonstrates that register tokens significantly lower feature norm variability and improve FID scores by acting as norm sinks and aggregating global information.
• The study reveals that registers are most effective in pixel-space settings when introduced later and in higher quantities, with a dual-stream architecture yielding notable performance gains at minimal computational cost.
• The work emphasizes that specialized, register-aware designs enhance intermediate feature interpretability and robustness across early and mid-stage denoising processes.

Register Tokens in Pixel-Space Diffusion Transformers: Role, Mechanisms, and Architectural Implications

Overview

"Registers Matter for Pixel-Space Diffusion Transformers" (2605.16147) delivers a comprehensive analysis of the function of register tokens in Diffusion Transformers (DiTs), particularly in the pixel space setting. While register tokens were originally introduced in Vision Transformers (ViTs) to mitigate high-norm outlier artifacts—in particular, to regularize attention distributions and improve feature map interpretability—the motivation, efficacy, and internal dynamics of register usage in DiTs had not been systematically characterized prior to this work. This paper establishes that, despite notable differences from the ViT setting (notably, the absence of patch-token outliers in DiTs), register tokens significantly enhance both convergence and final generation quality for pixel-space DiTs. Mechanistic analysis shows that these benefits arise through distinct phenomena from the ViT case—including norm sinking and global information aggregation—ultimately producing cleaner and more structured intermediate features, especially in the early and mid denoising stages. These insights are further leveraged to propose an efficient dual-stream DiT architecture with register-aware specialization, yielding systematic FID improvements at minimal computational cost.

Register Tokens in DiTs vs. ViTs

A primary empirical result is that, unlike ViTs, both latent-space and pixel-space DiTs do not intrinsically generate high-norm patch-token outliers. In ViTs, high-norm outliers (typically in low-information/background regions) motivate the deployment of registers to 'sink' global information and prevent these artifacts from contaminating the patch representations. However, in DiTs, attention maps and token feature distributions are near-uniform in the absence of registers, lacking the low-information artifacts found in ViT models.

Despite this, adding registers to DiTs induces the emergence of high-norm tokens—now contained within the registers themselves. Surprisingly, this addition consistently and substantially reduces FID (for example, FID improves from 7.39 to 5.30 in base pDiT models on ImageNet 256×256), with effects strictly dominant in pixel-space settings and minimal or even negative effects in some latent-space models.

Mechanistic Insights: Smoothing and Information Specialization

Analysis of intermediate representations reveals that register tokens serve two principal roles in pixel-space DiTs:

1. Norm Sinking: A subset of register tokens absorbs magnitude (i.e., become high-norm outliers), thereby reducing the feature norms of all patch tokens. This norm reduction regularizes local variability in the spatial features, resulting in smoother and more spatially coherent intermediate feature maps.
2. Global Information Aggregation: Other register tokens, as probed via class linear separability, encode diverse semantic and global properties of the input image, attending to distinctive semantic regions. The effect is most pronounced at high-noise diffusion steps ($t \in [0, 0.2]$), which are especially critical for flow-matched models in pixel space.

Total Variation (TV) and correlation decay metrics further quantify smoother spatial organization induced by registers—a phenomenon not observed in self-supervised ViTs with registers. This suggests that the benefit of registers in DiTs is not limited to simply absorbing outliers but more fundamentally to regularizing high-dimensional and noisy intermediate signal propagation, especially in pixel spaces with higher intrinsic noise.

Registers in the Broader Context of DiT Architectures

The analysis extends to contemporary architectural trends in DiTs, such as models that exploit in-context conditioning by appending auxiliary class or text tokens. The paper demonstrates that these in-context tokens functionally act as implicit registers, exhibiting similar norm sinking and global semantics behaviors. Empirically, most of the generation quality gains previously attributed to explicit class conditioning can be traced to this implicit register-like effect, rather than the mere inclusion of class information per se.

Furthermore, ablation studies clarify two critical design insights for register deployment:

• Registers are only beneficial if introduced after several initial transformer layers (after layer 4 in pDiT-B), unlike in ViTs where early introduction is effective. Early-layer registers lack semantic structure and act as ineffective norm sinks, degrading overall performance.
• Pixel-space DiTs require substantially more register tokens (e.g., 32) to reach optimality compared to ViTs, reflecting the higher complexity and variability in pixel-level generative modeling.

Dual-Stream Architecture: Efficient Specialization for Registers

Capitalizing on the distinct roles of registers and patch tokens, the authors develop and validate a parameter-efficient dual-stream DiT architecture. In this design, key transformer components (MLP, adaLN, RMSNorm) are decoupled for register vs. patch tokens, with parameter sharing or LoRA-based specialization employed for efficiency. Registers are only inserted in deeper blocks (as motivated empirically), and the architecture increases parameters by ~14% with no significant runtime penalty (in GFLOPs).

This dual-stream approach further improves FID across multiple model and resolution scales (e.g., FID drops from 3.71 to 3.41 on ImageNet 256×256 for base models) and remains effective when combined with recent representational alignment techniques (such as PixelREPA).

Practical and Theoretical Implications

The results carry several implications:

• For DiT Training: Registers should be included as standard practice in pixel-space and potentially high-noise or high-dimensional DiT setups, introduced only after sufficient early representation learning has occurred, and in sufficient quantity to absorb and specialize global information.
• For DiT Architecture Design: Specialized register-aware parameterization is beneficial and scales efficiently. Dual-stream or otherwise specialized handling of heterotypic tokens (registers, text/class, etc.) may become foundational in future state-of-the-art generative models.
• For Broader Transformer Research: The phenomenon of global information sinking and specialization of tokens appears repeatedly across modalities—vision, text, video, and multimodal generative models—inviting further unified theoretical treatment.
• For Interpretability and Robustness: Register tokens improve not only numerical metrics (e.g., FID) but also internal feature structure and interpretability, reinforcing the utility of architectural interventions for model robustness and transparency.

Conclusion

This study rigorously establishes that register tokens play a vital, mechanistically unique role in pixel-space Diffusion Transformers, despite the absence of the outlier artifacts that originally motivated them in ViTs. Their introduction consistently improves pixel-space DiT representations by regularizing spatial features and aggregating global semantic information. The dual-stream architectures proposed to support register specialization demonstrate that further, parameter-efficient advances are feasible. These findings set a new standard for DiT architectural design and open several avenues for deeper theoretical understanding and practical innovation in transformer-based generative models.