Registers Matter for Pixel-Space Diffusion Transformers

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.

📋 한 줄 요약

**[Diffusion Transformer / Register]** Pixel-space DiT가 ViT와 달리 token outlier가 없으나 register token이 수렴·생성 품질 크게 향상, 듀얼 스트림 아키텍처 제안.

🎯 핵심 기여도

ViT가 feature map 품질을 저하시키는 high-norm patch-token outlier를 보이며, register token이 이를 효과적으로 완화한다는 사전 연구 배경.
확산 모델이 점점 트랜스포머 아키텍처 채택·pixel-space 학습으로 이동하며 ViT와 형태가 가까워짐 — register token이 DiT에도 유용한가라는 질문 제기.
DiT가 ViT와 핵심 차이 — patch-token outlier를 보이지 않음을 발견.
그럼에도 register token이 pixel-space DiT의 수렴·생성 품질을 크게 향상시킨다는 흥미로운 발견.
Register token이 high noise level에서 더 깨끗한 feature map 생성 — pixel-space 생성에서 효과의 원인일 수 있음.
최근 pixel-space DiT 아키텍처가 register-like 메커니즘을 암묵적으로 통합하고 있음 관찰 — 강한 경험적 성능의 부분적 설명일 수 있음.

💡 핵심 아이디어

DiT는 ViT의 token outlier 문제를 보이지 않음에도 register token이 효과적인데, 이는 register가 outlier 완화가 아닌 high noise level에서 깨끗한 feature map 형성에 기여하기 때문이며, 이 통찰이 register token에 특화된 듀얼 스트림 아키텍처 설계로 이어진다.

🔬 기술적 접근법

**모델/방법론**: Pixel-space DiT의 register token 분석 + parameter-efficient dual-stream 아키텍처.
**핵심 기법**: (1) DiT의 patch-token outlier 부재 분석, (2) Register token 추가의 수렴·품질 영향 측정, (3) High noise level에서 register가 만드는 feature map 분석, (4) 최근 pixel-space DiT의 implicit register-like 메커니즘 식별, (5) Register token 처리를 특화한 parameter-efficient dual-stream 아키텍처.

📊 주요 결과

DiT가 ViT와 달리 patch-token outlier를 보이지 않음.
Register token이 pixel-space DiT의 수렴·생성 품질 크게 향상.
Register token이 high noise level에서 더 깨끗한 feature map 생성.
듀얼 스트림 아키텍처가 negligible runtime overhead로 pixel-space 생성 품질 개선.

💭 의의 및 한계

**의의**: pixel-space DiT의 내부 메커니즘에 새 통찰 제공, register token의 작동 원리에 대한 이해 갱신, pixel-space 생성 아키텍처 설계 가이드. **한계**: pixel-space DiT 중심 분석으로 latent-space DiT나 다른 확산 아키텍처로의 일반화는 추가 검증 필요.

🚀 실용적 활용

Pixel-space 확산 모델 학습 가속·품질 향상.
DiT 아키텍처 설계 지침.
비전 트랜스포머 일반의 register 활용 이해.