기술 요약: 이미지 합성을 위한 가이드형 엔드투엔드 자기회귀(GEAR)

1. 문제 정의

VQ-VAE와 같은 벡터 양자화(VQ) 토크나이저를 학습시켜 이미지를 재구성한 후 이를 고정(frozen)하고, 그 결과로 생성된 이산적 인덱스(discrete indices) 상에서 AR 생성기를 학습시키는 두 단계 학습 패러다임은 현대의 시각적 생성 모델(자기회귀 트랜스포머 및 확산 모델 포함)이 따르는 일반적인 방식이다.

이러한 디커플링(decoupling)은 근본적인 불일치를 초래한다: 토크나이저는 순수하게 재구성 충실도(reconstruction fidelity)를 위해 최적화되지만, 결과로 나오는 이산적 시퀀스가 다운스트림 생성기가 인과적으로 모델링하기 용이한 형태인지에 대해서는 고려하지 않는다. 최근 확산 모델 연구(예: REPA-E)에서는 외부 시맨틱 인코더(DINOv2 등)와 잠재 표현을 정렬함으로써 엔드투엔드 공동 학습을 성공적으로 탐구했으나, 이를 이산적 AR 생성에 적용하는 것은 훨씬 더 어렵다.

핵심적인 장애물은 **비미분성(non-differentiability)**이다. 연속적인 잠재 표현에서 이산적인 VQ 인덱스로 매핑하는 과정에는 arg max 연산이 포함된다. 결과적으로 AR 생성기로부터 오는 그래디언트가 토크나이저로 흐를 수 없다. 이 간극을 메우기 위해 **Straight-Through Estimator (STE)**를 사용하는 시도들은 불안정하여, 종종 코드북 붕괴(codebook collapse, 재구성 충실도를 희생하면서 예측 가능성을 극대화하기 위해 토크나이저가 몇 개의 지배적인 코드만을 출력하는 현상)를 유발하며, 이는 치명적인 실패(gFID $\approx$ 105)로 이어진다.

2. 방법론: GEAR

저자들은 코드북 붕괴 없이 VQ 토크나이저와 AR 생성기를 엔드투엔드로 공동 학습시키는 프레임워크인 **GEAR (Guided End-to-end AutoRegression)**를 제안한다.

핵심 메커니즘: 이중 판독(Dual Read-Out)

GEAR는 모든 위치에 대해 코드북 할당의 이중 판독을 도입하여 비미분성 문제를 해결한다:

1. 하드 브랜치 (One-Hot):
   • 표준 이산 인덱스($q_i = \text{arg max}_k A_{ik}$)를 사용한다.
   • 이산 토큰을 AR 모델에 입력하여 다음 토큰 예측(NTP) 및 하드 표현 정렬 손실(hard representation alignment loss)을 수행한다.
   • 결정적 특징: 이 브랜치는 토크나이저에 대해 미분 불가능하다. NTP 및 하드 정렬 손실로부터 오는 그래디언트는 오직 AR 모델만 업데이트하며, 토크나이저를 저엔트로피 코드로 붕괴시키지 않는다.
2. 소프트 브랜치 (Differentiable):
   • 코드북 항목들에 대한 온도 가중 소프트맥스 보간($\pi_i = \text{softmax}(A_i / \tau)$)을 사용한다.
   • 부드럽고 연속적인 임베딩을 AR 모델에 입력한다.
   • 결정적 특징: 이 브랜치는 완전히 미분 가능하다. 이는 표현 정렬 손실(예: DINOv2 피처와의 정렬)을 전달하며, 이 손실은 AR 백본을 통해 흘러가 오직 토크나이저만을 업데이트한다.

최적화 전략

두 모듈은 서로 분리된 목적 함수로 업데이트된다:

• 토크나이저 ($\theta_{tok}$): 표준 VQ 재구성 손실($L_{VQ}$)과 소프트 정렬 손실($L^s_{align}$)에 의해 업데이트된다. 토크나이저는 예측 가능하고 시맨틱하게 정렬된 인덱스 분포를 형성하도록 유도되지만, 예측 손실 자체에 의해서는 유도되지 않는다.
• AR 생성기 ($\theta_{AR}$): NTP 손실($L_{NTP}$)과 하드 정렬 손실($L^h_{align}$)에 의해 업데이트된다.

이 설계는 토크나이저가 AR 모델이 예측하기 쉽고 시맨틱 피처와 잘 정렬될 수 있는 토큰을 생성하도록 학습시키는 동시에, 높은 재구성 충실도를 유지하도록 보장한다.

3. 주요 기여

• 가이드형 엔드투엔드 학습: GEAR는 미분 가능한 소프트 할당 브릿지를 도입하여, AR 모델의 표현 정렬 목적 함수가 토크나이저를 가이드할 수 있게 함으로써 표준 STE 방식이 극복하지 못한 비미분적 인덱스 장벽을 넘어선다.
• 정렬 부담의 전이 (Shifted Alignment Burden): 잠재 표현 자체를 시맨틱하게 만드는 확산 모델 측 방법론(예: REPA-E, VA-VAE)과 달리, GEAR는 정렬의 부담을 AR 생성기로 옮긴다.
  • 토크나이저의 피처는 덜 DINOv2스러워지며, 인덱스 분포를 더 예측 가능한 저엔트로피 구조로 재편한다.
  • AR 생성기의 은닉 상태는 특히 **패치 레벨(patch level)**에서 더 DINOv2스러워지며, 풍부한 국소적 공간 구조를 포착한다.
• 범용성: 이 메커니즘은 하위 양자화기(quantizer)에 무관하며, VQVAE, LFQ, IBQ 전반에서 성능 향상을 입증했다. 또한 클래스 조건부 생성에서 텍스트-이미지 합성으로 효과적으로 전이된다.

4. 실험 결과

실험은 ImageNet-1K (256x256) 및 GPIC 텍스트-이미지 벤치마크에서 수행되었다.

• 수렴 속도: GEAR는 강력한 LlamaGen-REPA 베이스라인과 비교하여 ImageNet gFID 수렴 속도를 최대 10배까지 가속화한다. STE를 사용하는 나이브한 엔드투엔드 변형은 완전히 발산한다.
• 생성 품질:
  • ImageNet: 300 에포크에서 GEAR (775M 파라미터)는 CFG 적용 시 2.52의 gFID를 달성하여, LlamaGen-REPA (2.68)를 능가하고 강력한 잠재 확산 모델(예: REPA-E의 1.12)에 근접하는 성능을 보인다.
  • Text-to-Image: GPIC에서 GEAR는 모든 학습 단계(50k ~ 390k)에서 LlamaGen-REPA를 일관되게 앞서며, DINOv2 피처 공간에서의 Fréchet 거리(FDD)를 더 낮게 달성한다.
• 재구성: 재구성을 위해 시맨틱을 희생하는 방법들과 달리, GEAR는 엔드투엔드로 학습하면서도 재구성 품질(rFID)을 개선하거나 유지한다.
• Ablation Study:
  • STE 실패: 소프트 브랜치를 STE로 교체하면 코드북 붕괴(gFID $\approx$ 105)가 발생한다.
  • 온도(Temperature): $\tau=0.1$의 온도가 생성 품질과 재구성 사이의 균형을 맞춘다.
  • 모델 크기: GEAR는 Base, Large, XLarge 스케일 모두에서 성능을 향れて시키며, 공동 학습된 토크나이저는 AR 모델의 크기가 커짐에 따라 개선된 재구성을 보여준다.

5. 의의 및 주장

본 논문은 GEAR가 이산 토큰 파이프라인을 위한 통합된 엔드투엔드 시각적 생성을 향한 중요한 진전임을 주장한다.

• 이산적 병목 현상 극복: GEAR는 예측 손실을 토크나이저 업데이트로부터 분리함으로써, STE의 불안정성 없이 이산 AR 모델을 위한 엔드투엔드 학습이 가능하다는 것을 입증했다.
• 새로운 정렬 패러다임: 저자들은 이산 AR 모델을 위한 최적의 전략이 토크나이저를 시맨틱하게 만드는 것이 아님을 강조한다. 대신, 토크나이저는 예측 가능성(저엔트로피, 구조화된 인덱스)을 위해 최적화되어야 하며, 시맨틱 정렬은 AR 모델의 은닉 상태 내에서 나타난다. 이러한 "국소적 패치 레벨 구조"가 다음 토큰 예측을 효율적으로 만든다.
• 통합 시스템의 기반: 저자들은 이산적 차기 토큰(next-token) 형식이 긴 컨텍스트와 멀티모달 입력을 처리하는 데 적합하여, 통합 이해 및 생성 시스템(예: LLM의 RLHF, DPO와 같은 성숙한 정렬 기술 계승)에 독보적으로 유리하다고 본다. GEAR는 이러한 시스템을 확장하기 위한 견고한 토대를 제공한다.

저자들은 GEAR가 AR과 확산 모델 사이의 격차를 좁히고 있지만, 최상위 확산 모델에서 사용되는 연속 VAE와 비교했을 때 VQ 토크나이저의 이산적 특성으로 인한 **재구성 천장(reconstruction ceiling)**이 여전히 제한 사항임을 인정한다. 향-후 연구로는 압축과 시퀀스 길이를 분리하는 방식(예: 패치화된 AR)을 탐구하여 충실도를 더욱 높이는 것이 제안되었다.