[논문리뷰] Registers Matter for Pixel-Space Diffusion Transformers

링크: 논문 PDF로 바로 열기

메타데이터

저자: Nikita Starodubcev, Ilia Sudakov, Ilya Drobyshevskiy, Artem Babenko, Dmitry Baranchuk

---

1. Key Terms & Definitions (핵심 용어 및 정의)

• Register Tokens: 고정된 수의 학습 가능한(learnable) 토큰으로, 모델이 핵심 정보를 보존하거나 특정 계산 역할을 수행하도록 설계된 추가 토큰입니다.
• Pixel-Space Diffusion Transformers (pDiTs): latent space를 거치지 않고 직접 픽셀 수준에서 데이터 분포를 학습하는 확산 모델 기반의 트랜스포머입니다.
• Norm Sinks: 입력 토큰들의 feature norm이 비정상적으로 높아지는 것을 방지하기 위해, 큰 값을 흡수하는 역할을 수행하는 특정 토큰을 의미합니다.
• Dual-Stream Architecture: 서로 다른 토큰 유형(예: patch vs register)을 개별적으로 처리하면서 필요한 정보만 교환하도록 설계된 파라미터 효율적인 아키텍처입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 Register Tokens가 기존 ViTs에서의 고질적인 문제인 '고 norm 아웃라이어(high-norm patch-token outliers)'를 해결하는 것과 달리, DiTs에서의 구체적인 역할과 효과는 미비하게 탐구되었다는 점에 주목합니다. 기존 연구들은 ViTs에서 발생하는 어텐션 맵 아티팩트를 줄이기 위해 Register Tokens를 도입했으나, 저자들의 분석 결과 DiTs는 이러한 patch-token 아웃라이어를 본래 갖지 않음에도 불구하고 Register Tokens 도입 시 생성 품질이 크게 향상됨을 확인했습니다 [Figure 2, 3]. 이에 따라 저자들은 DiTs에서 Register Tokens가 왜 성능 향상을 이끄는지 그 메커니즘을 규명하고, 이를 효과적으로 활용하는 새로운 아키텍처를 제안하고자 합니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

저자들은 Pixel-Space DiTs의 성능 향상을 위해 register 토큰의 특성을 활용한 parameter-efficient dual-stream architecture를 제안합니다. 이 모델은 초기 층에서는 공통적으로 입력을 처리하고, 깊은 층(layer 4 이후)에서 patch token과 register token의 경로를 분리하여 각 토큰이 고유한 역할을 수행하도록 최적화합니다 [Table 3]. 핵심 연구 결과에 따르면, Register Tokens는 크게 두 가지 역할을 수행합니다: 일부 토큰은 전역 의미 정보(global semantic information)를 인코딩하고, 나머지 토큰은 patch tokens의 feature norm을 흡수하는 norm sinks 역할을 하여 내부 피처 맵을 평탄화(smooth)합니다 [Figure 4, 6]. 실험적으로, 제안된 구조는 JiT 베이스라인 대비 FID 수치를 유의미하게 개선하였으며, 특히 ImageNet 256x256 설정에서 3.71에서 3.41로 성능을 향상시켰습니다 [Table 5, 6]. 또한, 이러한 성능 향상은 모델 사이즈 증가율 대비 약 14%의 파라미터 증가만으로 달성되어 높은 파라미터 효율성을 입증했습니다.

4. Conclusion & Impact (결론 및 시사점)

본 논문은 Register Tokens가 단순한 아웃라이어 제거 장치를 넘어, pixel-space 확산 모델의 내부 표현을 최적화하고 공간적 일관성(spatial coherence)을 높이는 핵심적인 역할을 수행함을 증명합니다. 특히 Register Tokens와 In-context class-conditioning 메커니즘 간의 기능적 유사성을 밝혀내어, 최근 고성능 DiT 모델들의 설계 원리를 이론적으로 정립하였습니다. 본 연구에서 제안된 dual-stream 설계 전략은 향후 확산 모델 아키텍처를 설계함에 있어 정보 흐름을 토큰 유형별로 특화하는 새로운 방법론을 제시하며, 고해상도 생성 모델의 학습 효율성을 극대화하는 데 크게 기여할 것으로 기대됩니다.

Figure 3 — DiT에서의 토큰별 feature norm

Figure 5 — Register 유무에 따른 피처 스무스니스 비교

Figure 6 — Register 토큰의 이중 역할

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.