기술 요약: 픽셀 공간 확산 트랜스포머에서 레지스터의 중요성

문제 제기

비전 트랜스포머 (ViT) 는 특징 맵의 품질을 저하시키고 어텐션 맵에 아티팩트를 생성하는 고노름 (high-norm) 패치 토큰 이상치 (outliers) 를 나타내는 것으로 알려져 있습니다. 이 문제는 ViT 에 '레지스터 토큰 (register tokens)'을 도입함으로써 효과적으로 완화되는데, 이는 이러한 이상치를 흡수하는 싱크 (sinks) 역할을 합니다. 확산 트랜스포머 (DiT) 가 트랜스포머 아키텍처를 점점 더 채택하고 잠재 공간 (latent space) 에서 픽셀 공간 (pixel space) 으로 직접 훈련하는 방향으로 전환함에 따라, 구조적으로 ViT 와 더 유사해지고 있습니다. 이는 두 가지 중요한 질문을 제기합니다:

1. DiT 는 ViT 에서 관찰된 고노름 패치 토큰 이상치를 계승하는가?
2. 레지스터 토큰은 특히 픽셀 공간 훈련 체계에서 DiT 에 유익한가?

어텐션 싱크와 특수 토큰의 동작은 대규모 언어 모델 (LLM) 과 확산 언어 모델에서 연구되었으나, 이미지 확산 트랜스포머에서의 존재와 기능적 역할은 아직 충분히 탐구되지 않았습니다.

방법론

저자들은 픽셀 공간, VAE 공간, RAE 공간 등 다양한 확산 모델 공간에서 레지스터 토큰에 대한 체계적인 분석을 수행하고, 이를 표준 ViT(특히 DINOv2) 와 비교합니다.

1. 이상치에 대한 실증적 분석: 이 연구는 ImageNet 에서 플로우 매칭 (flow matching) 으로 훈련된 pDiT 와 ViT 간의 토큰별 특징 노름과 어텐션 맵을 조사합니다. 레지스터가 없는 상태에서 패치 토큰이 고노름 이상치를 나타내는지 분석합니다.
2. 레지스터 토큰 통합: 레지스터 토큰은 ViT 에서 사용된 접근 방식을 따르며 패치 토큰 시퀀스에 추가되는 추가 학습 가능한 토큰으로 구현됩니다. 이러한 토큰은 훈련 목적 함수 (손실 함수) 에서 제외됩니다.
3. 내부 표현 분석: 저자들은 총 변동 (Total Variation, TV) 을 사용하여 중간 특징 맵의 공간적 부드러움을 측정하고, 선형 프로빙 (linear probing) 을 통해 레지스터 토큰의 의미론적 내용을 평가합니다. 또한 레지스터가 다양한 확산 시간 단계 (노이즈 수준) 에서 특징 노름과 표현의 일관성에 어떻게 영향을 미치는지 검토합니다.
4. 아키텍처 탐색: 레지스터와 패치 토큰의 명확한 역할에 영감을 받아, 논문은 매개변수 효율적인 듀얼 스트림 아키텍처를 제안합니다. 이 설계는 어텐션 메커니즘과 대부분의 트랜스포머 매개변수를 공유하면서 레지스터와 패치 토큰을 전문화된 경로 (특히 RMSNorm, MLP, adaLN 계층을 분리) 를 통해 처리합니다.
5. 애블레이션 연구: 이 연구는 최적의 구성을 결정하기 위해 레지스터의 수, 도입되는 깊이, 그리고 듀얼화할 특정 구성 요소를 제거하여 분석합니다.

주요 발견 및 기여

1. 이상치에 관한 DiT 와 ViT 의 근본적 차이

ViT 와는 대조적으로, 표준 DiT(잠재 공간 및 픽셀 공간 모두) 는 고노름 패치 토큰 이상치를 나타내지 않습니다. ViT 에서 이상치는 특정 토큰 (종종 정보량이 적은 영역에 위치) 에 국한되는 반면, DiT 에서는 패치 토큰 노름이 상대적으로 균일하게 유지되며 ViT 에서 관찰되는 정보량이 적은 아티팩트가 없는 어텐션 맵을 보입니다.

2. 기존 이상치가 없음에도 레지스터가 픽셀 공간 DiT 를 개선함

놀랍게도, DiT 에 레지스터 토큰을 추가하면 특히 픽셀 공간 모델에서 수렴과 생성 품질이 크게 향상됩니다.

• 메커니즘: ViT 에서 레지스터가 기존 이상치를 흡수하는 것과 달리, DiT 의 레지스터는 스스로 고노름 토큰을 생성합니다.
• 효과: 레지스터의 존재는 모든 패치 토큰의 특징 노름을 일관되게 감소시킵니다. 이는 특히 높은 노이즈 수준(초기 확산 시간 단계) 에서 더 깨끗하고 부드러운 중간 특징 맵으로 이어집니다.
• 이중 역할: 레지스터 토큰은 두 가지 명확한 역할로 전문화됩니다:
  • 노름 싱크 (Norm Sinks): 패치 토큰으로부터 크기를 흡수하는 극도로 높은 노름을 가진 토큰.
  • 전역 정보 운반자 (Global Information Carriers): 선형 프로빙과 어텐션 맵 시각화를 통해 검증된 바와 같이, 이미지 (전경 객체, 배경 컨텍스트 등) 에 대한 다양한 의미론적 정보를 인코딩하는 중간 정도의 노름을 가진 토큰.

3. 공간별 효과성

레지스터 토큰의 이점은 훈련 공간에 크게 의존합니다:

• 픽셀 공간: 레지스터는 생성 품질 (FID) 에서 가장 큰 개선을 가져옵니다. 픽셀 공간 모델은 본질적으로 더 높은 특징 노름을 가진 더 노이즈가 많은 중간 표현을 생성하므로, 레지스터의 정규화 효과가 중요합니다.
• 잠재 공간 (VAE/RAE): 레지스터는 중간 정도의 개선을 제공하거나, RAE 기반 모델의 경우 성능을 약간 저하시킵니다. 잠재 공간은 더 구조화되어 있고 차원이 낮아 레지스터가 제공하는 강력한 정규화가 덜 필요합니다.

4. 기존 아키텍처의 암시적 레지스터

저자들은 컨텍스트 내 클래스 조건부 (in-context class conditioning, 토큰 시퀀스에 클래스 임베딩을 추가) 를 사용하는 최근 픽셀 공간 DiT 아키텍처 (예: JiT) 가 암시적으로 레지스터와 유사한 메커니즘을 통합하고 있음을 관찰합니다. 이러한 컨텍스트 내 토큰은 명시적 레지스터와 유사한 행동을 보입니다: 일부는 노름 싱크 역할을 하고, 다른 일부는 전역 의미론적 정보를 인코딩합니다. 이는 이러한 모델의 성능 향상이 클래스 정보 자체보다는 이러한 레지스터와 유사한 행동에서 비롯되었을 가능성을 시사합니다.

5. 매개변수 효율적인 듀얼 스트림 아키텍처

이러한 통찰력을 활용하기 위해 저자들은 레지스터 토큰 처리를 전문화하는 컴팩트한 듀얼 스트림 아키텍처를 제안합니다.

• 설계: 레지스터는 의미론적 구조가 확립되는 더 깊은 레이어 (블록 4 부터 시작) 에서만 도입됩니다. 아키텍처는 공유 어텐션을 사용하지만 레지스터와 패치를 위해 RMSNorm, MLP, adaLN 을 분리합니다.
• 효율성: 이 접근 방식은 모델 크기를 단순한 전체 복제 (~77%) 에 비해 불과 **약 14%**만 증가시키면서 유사한 추론 비용 (GFLOPs) 을 유지합니다.
• 성능: 이 아키텍처는 단일 스트림 베이스라인에 비해 모델 규모 (B, L, H) 와 해상도 (256x256, 512x512) 전반에 걸쳐 생성 품질을 일관되게 향상시킵니다.

결과

• 생성 품질: ImageNet 256x256 에서 레지스터가 있는 pDiT 는 없는 경우보다 일관되게 우수한 성능을 보입니다. 예를 들어, pDiT-B/16 모델은 200 에포크에서 FID 를 7.39 에서 5.30 으로 개선했습니다.
• 특징 부드러움: 레지스터가 있는 모델은 높은 노이즈 수준에서 중간 특징의 총 변동 (TV) 비율이 낮아 더 부드럽고 일관된 표현을 나타냅니다.
• 애블레이션: 레지스터는 처음 4 레이어 이후에 도입되고 더 많은 수의 토큰 (예: ViT 의 4 개 대비 pDiT-B 의 32 개 레지스터) 을 사용할 때 가장 효과적입니다.

중요성 및 주장

이 논문은 레지스터 토큰이 고품질 픽셀 공간 확산 생성을 위한 중요하지만 이전에는 간과되었던 구성 요소라고 주장합니다. 그 중요성은 다음과 같습니다:

1. 레지스터 역할의 재정의: 레지스터가 ViT 고유의 이상치에 대한 단순한 해결책이 아니라, 특징 노름을 관리하고 공간적 일관성을 향상시킴으로써 확산 모델에서 더 넓은 정규화 기능을 수행함을 입증합니다.
2. 최근 성공의 설명: 컨텍스트 내 조건부를 활용하는 최근 픽셀 공간 DiT 의 강력한 실증적 성능에 대한 근거를 제공하며, 그 성공의 일부를 암시적 레지스터 행동에 귀속시킵니다.
3. 효율적인 아키텍처 설계: 레지스터와 패치의 명확한 역할을 활용하는 매개변수 효율적인 듀얼 스트림 설계를 제시하여, 거의 없는 런타임 오버헤드로 최첨단 생성 품질을 달성합니다.

저자들은 그들의 설계가 토큰 전문화의 가능성을 포착하지만, 픽셀 공간 확산 트랜스포머에서 레지스터를 활용하는 완전한 잠재력은 향후 연구를 위한 열린 질문으로 남아 있다고 결론지었습니다.