기술 요약: 확산 트랜스포머의 이상치 토큰 제어

문제 정의

본 논문은 이미지 생성에 사용되는 확산 트랜스포머 (DiT) 내에서 발생하는 이상치 토큰 (outlier tokens) 현상을 조사합니다. 이전 연구 (예: Darcet 등) 는 인식 작업에서 비전 트랜스포머 (ViT) 가 비정상적으로 큰 특징 노름을 가진 소수의 토큰을 생성하여 제한된 지역 정보를 담고 있음에도 불구하고 불균형적으로 많은 주의를 흡수한다는 사실을 규명했으나, 이러한 토큰이 생성 모델에서 수행하는 역할은 아직 탐구되지 않았습니다.

저자들은 두 가지 중요한 격차를 확인합니다:

1. 생성 파이프라인의 보편성: 이상치 토큰은 현대의 표현 오토인코더 (RAE) 파이프라인의 사전 훈련된 ViT 인코더 (토크나이저) 에서뿐만 아니라 확산 트랜스포머 (디노이저) 내부에서도 나타납니다.
2. 결함의 본질: 이러한 토큰으로 인한 열화는 단순히 극단적인 손실 값 몇 개로 인한 결과가 아닙니다. 대신 저자들은 이상치 토큰이 손상된 지역 패치 의미론의 증상이라고 가정합니다. 훈련 중 고노름 토큰을 단순히 마스킹하는 것이 성능을 향상시키지 못한다는 사실은, 문제가 단순히 수치적 불안정성이 아니라 공간 구조와 지역 정보의 손실에 있음을 시사합니다.

방법론: 이중 단계 레지스터 (DSR)

이러한 문제를 해결하기 위해 저자들은 표현 인코더와 확산 생성기라는 RAE-DiT 파이프라인의 두 구성 요소 모두에 레지스터 토큰을 적용하는 통합 개입 전략인 이중 단계 레지스터 (Dual-Stage Registers, DSR) 를 제안합니다.

1. 비전 인코더의 레지스터

저자들은 사전 훈련된 레지스터의 가용성에 따라 두 가지 접근 방식을 사용하여 인코더 단계 (토크나이저) 의 이상치 토큰을 처리합니다:

• 훈련된 레지스터: DINOv2 와 같이 레지스터를 지원하는 인코더의 경우, 저자들은 사전 훈련된 레지스터 토큰을 활용하여 인코더 특징 맵의 고노름 아티팩트를 억제합니다.
• 재귀적 테스트 시간 레지스터 (TTR): 레지스터가 내장되지 않은 인코더 (예: SigLIP2) 의 경우, 저자들은 추론 시 입력 시퀀스에 추가적인 미훈련 토큰을 추가합니다. 결과 표현에서 2 단계 이상치가 감지되는 경우, 출력을 안정화하기 위해 TTR 을 재귀적으로 적용합니다.

2. 확산 트랜스포머의 레지스터

저자들은 안정화된 인코더를 사용하더라도 확산 트랜스포머 자체가 내부적으로 이상치 토큰을 발달시킨다는 것을 관찰합니다. 이는 최종 층에서 주로 나타나는 표준 ViT 와 달리, 주로 중간 층에서 발생합니다.

• 확산 레지스터: 그들은 학습 가능한 레지스터 토큰 소집합을 확산 트랜스포머에 직접 도입합니다. 이러한 토큰은 생성기와 함께 훈련되어 주의를 흡수하는 싱크 (attention sinks) 를 흡수하고 표현 공간을 안정화합니다.
• 추론 프로토콜: 추론 시에는 인코더와 확산 모델 모두에서 나온 모든 레지스터 토큰의 출력을 폐기하고 이미지 토큰 출력만 유지합니다.

주요 기여

• 내부 이상치 발견: 본 논문은 이상치 토큰이 인코더와 디노이저 모두에서 발견되며, 확산 모델의 중간 층에 독특한 집중 패턴을 보이며 DiT 에 만연함을 입증합니다.
• "극단적 값" 가설의 기각: 손실 마스킹 실험을 통해 저자들은 단순히 고노름 토큰을 훈련 손실에서 제외하는 것이 생성 품질을 향상시키지 못함을 보여줍니다. 이는 핵심 문제가 극단적 값의 존재가 아니라 지역 패치 의미론의 손상임을 결론짓게 합니다.
• 이중 단계 개입: DSR 제안은 인코더 (출처) 와 생성 과정 (디노이저) 에서 이상치를 해결함으로써 전체 생성 파이프라인을 안정화하는 간단하고 경량화된 메커니즘을 제공합니다.
• 일반화 가능성: 이 방법은 픽셀 공간, VAE 잠재 공간, 표현 공간 등 다양한 입력 공간과 SiT, JiT, RAE 기반 설계 등 다양한 아키텍처에서 효과적임이 입증되었습니다.

실험 결과

저자들은 ImageNet-1K 클래스 조건부 생성 (256×256) 과 대규모 텍스트 - 이미지 생성에서 DSR 을 평가합니다.

• ImageNet-1K 성능:
  • SigLIP2-B를 인코더로 사용할 때, DSR 은 800 에포크에서 FID 를 5.89 에서 4.58 로 낮추고 텍스트 - 이미지 작업의 GenEval 점수를 0.426 에서 0.466 으로 향상시킵니다.
  • DINOv2-B 인코더와 결합할 경우, FID 는 4.16 에서 3.95 로 개선됩니다.
  • DSR 은 수렴을 가속화하여 기준선 대비 4 배 적은 훈련 에포크로 동등한 품질을 달성합니다.
• 확장성: 개선 효과는 모델 크기 (DiT-S, B, XL) 전반에 걸쳐 일관되며, 계산 비용은 미미하게만 증가합니다 (약 10% GFLOPs 증가).
• 결손 연구:
  • 삽입 깊이: 레지스터는 초기부터 중간 층 (예: 블록 8 에서 시작) 에 도입될 때 가장 효과적입니다. 너무 일찍 또는 너무 늦게 도입하면 효과가 감소합니다.
  • 레지스터 수: 적절한 수의 레지스터 (예: 36 개) 가 최적의 성능을 내며, 너무 적거나 (1 개) 너무 많으면 (100 개) 성능이 저하됩니다.
• 대안과의 비교: DSR 은 입력 표현이 더 의미론적이 될수록 "인컨텍스트 컨디셔닝 (in-context conditioning, 입력 의존적 토큰 사용)"보다 우수한 성능을 보여주며, 레지스터가 전역 용량을 위한 더 강력한 메커니즘을 제공함을 시사합니다.

중요성과 주장

본 논문은 이상치 토큰 제어가 견고한 트랜스포머 기반 확산 파이프라인 구축을 위해 중요하지만 이전에는 간과되어 온 요소라고 주장합니다. 이러한 아티팩트가 지역 패치 의미론을 저하시키고 레지스터 기반 개입이 이를 효과적으로 완화할 수 있음을 입증함으로써, 저자들은 DSR 을 인코더와 생성기 안정성을 위한 기초적인 개선으로 위치시킵니다.

이 연구는 트랜스포머의 유연성이 전역 상호작용에 유익하지만, 생성 과정에서 공간 충실도를 방해하는 특정 아티팩트 (이상치 토큰) 를 도입한다는 점을 강조합니다. 제안된 솔루션은 복잡한 아키텍처 변경이나 전체 파이프라인의 처음부터 재훈련 없이 다양한 설정에서 생성 품질을 일관되게 향상시키는 간단하고 플러그 - 앤 - 플레이 방식의 개입으로 제시됩니다.