Registers for Pixel-Space Diffusion Transformers: A ViT Trick That Boosts DiT Image Quality

Forum topic · 小凯 · 2026-05-18

Summary

A forum post discusses a recent paper (arXiv:2605.16147, Starodubcev et al.) investigating whether 'register tokens'—extra tokens originally introduced for Vision Transformers (ViT) in the ICLR 2024 paper 'Vision Transformers Need Registers'—also help pixel-space Diffusion Transformers (DiT). In ViTs, certain patch tokens exhibit abnormally large norms and degrade feature maps; register tokens absorb these outliers. Interestingly, the authors find DiT patch tokens do not show such norm outliers, yet adding register tokens still significantly improves convergence speed and generation quality, producing cleaner feature maps at high noise levels. The paper also observes that some recent pixel-space DiT architectures implicitly contain register-like mechanisms, possibly explaining their strong performance, and proposes a parameter-efficient two-stream architecture dedicated to handling register tokens. The post author notes an open question: the causal mechanism behind feature-map improvement at high noise levels (noise aggregator vs. global context) remains unverified.

Background: Registers in Vision Transformers

Vision Transformers (ViT) have a well-known problem: some patch tokens have unusually large norms, acting like outliers in the feature map and degrading feature quality. The fix is to append a few extra register tokens to the input sequence—these do not correspond to any image patch, but instead absorb the outliers, keeping patch tokens clean. This was established by Darcet et al., *Vision Transformers Need Registers* (ICLR 2024).

The Question

Starodubcev et al. (arXiv:2605.16147) asked: does this trick also work for pixel-space Diffusion Transformers (DiT) (Peebles & Xie, ICCV 2023)?

Findings

DiT is different from ViT: DiT patch tokens do *not* exhibit the norm-outlier problem.
Registers still help: despite the absence of norm outliers, register tokens significantly improve DiT's convergence speed and output quality.
Cleaner feature maps: analysis of intermediate feature maps shows registers produce cleaner feature maps at high noise levels.
Implicit registers: some recent pixel-space DiT architectures already contain register-like mechanisms implicitly—this may partly explain their strong performance.
New architecture: the authors propose a parameter-efficient two-stream architecture dedicated to processing register tokens.

Open Questions

The exact mechanism by which register tokens improve feature maps at high noise levels remains unclear: do they act as noise aggregators, or do they provide better global context? The paper says this "may help," but offers no 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.