Vision Transformer (ViT) has a well-known problem: some patch tokens exhibit abnormally large norms, acting as outliers within the image feature map and degrading its quality. The established fix is to add a few extra "register tokens" to the input sequence—tokens that do not correspond to any image patch and exist solely to absorb these outliers, keeping the patch tokens clean (Darcet et al., *Vision Transformers Need Registers*, ICLR 2024).
Starodubcev et al., in a paper (arXiv:2605.16147), asked whether this trick also works for pixel-space Diffusion Transformers (DiT).
Interestingly, DiT differs from ViT here: DiT's patch tokens do not show norm outliers. However, register tokens still significantly improve both convergence speed and output quality. By analyzing intermediate feature maps, they found that register tokens produce cleaner feature maps at high noise levels. They also observed that several recent pixel-space DiT architectures implicitly contain register-token-like mechanisms—which may partly explain their strong performance. Finally, the paper proposes a parameter-efficient two-stream architecture specifically designed to handle register tokens.
Unclear points
The precise mechanism behind "register tokens improving feature maps at high noise levels" remains unexplained: do they act as noise aggregators, or do they provide better global context? The paper says it "may help" but does not offer causal verification.
References
1. Starodubcev, N., et al. (2026). *Registers Matter for Pixel-Space Diffusion Transformers*. arXiv:2605.16147 [cs.CV]. 2. Darcet, T., et al. (2024). *Vision Transformers Need Registers*. ICLR 2024. 3. Peebles, W., & Xie, S. (2023). *Scalable Diffusion Models with Transformers*. ICCV 2023.