GEAR: Guided End-to-End AutoRegression for Image Synthesis

Visual generative models are typically trained in two stages. A tokenizer is first trained for reconstruction and then frozen, after which a generator is trained on its discrete indices or continuous latents. This decoupling leaves the tokenizer unaware of what the generator finds easy to model. We present GEAR (Guided End-to-end AutoRegression), which trains a vector-quantized (VQ) tokenizer and an autoregressive (AR) generator jointly and end-to-end, guided by representation alignment. The key obstacle is that the VQ index fed to the AR model is non-differentiable, so gradients cannot reach the tokenizer, and a straight-through estimator collapses. GEAR resolves this with a dual read-out of the codebook assignment. A hard, one-hot branch trains the AR with next-token prediction, while a differentiable soft branch carries a representation-alignment loss that flows back to guide only the tokenizer. The AR model thereby steers its tokenizer toward an index distribution it can predict more easily. This shifts the alignment burden from the tokenizer to the AR: the tokenizer's own features become less DINOv2-like while the AR's become more so, the opposite of diffusion-side recipes that make the latent itself semantic. GEAR speeds up ImageNet gFID convergence by up to 10x relative to the strong LlamaGen-REPA baseline, learns markedly better patch-level and spatially-coherent features, and generalizes across quantizers (VQVAE, LFQ, IBQ) and to text-to-image generation.

한 줄 요약

GEAR은 VQ 토크나이저와 AR 생성기를 엔드투엔드로 학습하여 이미지 생성 성능을 10배 가속화하는 모델이다.

핵심 기여도

VQ 토크나이저와 AR 생성기를 엔드투엔드로 학습하여, 기존 STE 기반 방법 대비 ImageNet gFID 수렴 속도를 최대 10× 향상.
하드-소프트 이중 할당 기법을 도입하여, AR의 다음 토큰 예측 손실(NTP)은 토크나이저에 전달되지 않도록 분리.
DINOv2 유사도는 토크나이저가 아닌 AR 생성기에서 발생하며, 토크나이저는 예측 가능한 인덱스 분포로 재구성됨.
VQVAE, LFQ, IBQ 등 다양한 퀀티파이어에 일반화 가능하며, 텍스트-이미지 생성에도 적용 가능.

핵심 아이디어

기존 시각 생성 모델은 토크나이저와 생성기를 분리하여 학습하는 두 단계 방식을 채택한다. 이는 토크나이저가 생성기가 예측하기 쉬운 인덱스를 생성하도록 고려하지 못하게 만든다. GEAR은 이 문제를 해결하기 위해 토크나이저와 AR 생성기를 엔드투엔드로 학습한다. 핵심은 VQ 인덱스가 미분 불가능한 argmax로 인해 기울기가 토크나이저로 전달되지 않는 문제를 해결하는 **하드-소프트 이중 할당**(dual read-out) 기법이다. 하드 브랜치는 AR의 NTP 손실을 통해 토크너 인덱스를 학습하고, 소프트 브랜치는 미분 가능한 대체 인덱스를 생성하여 토크나이저에만 손실을 전달한다. 이로 인해 토크나이저는 AR이 예측하기 쉬운 분포로 재구성되며, 생성기의 DINOv2 유사도는 토크나이저보다 AR에서 더 강하게 나타난다.

기술적 접근법

**VQ 토크나이저**: 이미지를 N×d 크기의 잠재 공간으로 인코딩하고, 학습 가능한 코드북에서 가장 가까운 인덱스를 선택.
**AR 생성기**: 토크나이저가 생성한 인덱스를 래스터 순서로 입력받아, 다음 토큰 예측(NTP) 손실로 학습.
**하드-소프트 이중 할당**: 하드 브랜치는 one-hot 인덱스로 AR을 학습하고, 소프트 브랜치는 temperature-weighted interpolation을 통해 미분 가능한 인덱스를 생성하여 토크나이저에만 손실을 전달.
**하이퍼파라미터**: temperature τ=0.1, alignment coefficient λ=0.5로 설정.
**데이터셋**: ImageNet 사용.
**모델**: LlamaGen-REPA 기반으로 B, L, XL 크기의 모델 실험.

주요 결과

**ImageNet gFID**: LlamaGen-REPA 대비 최대 10× 빠른 수렴.
**VQVAE**: gFID 14.72 → 10.63, rFID 1.724 → 1.640.
**LFQ**: gFID 18.68 → 14.78, rFID 2.421 → 2.129.
**IBQ**: gFID 20.25 → 12.97, rFID 1.973 → 1.716.
**모델 크기**: XL 모델에서 gFID 7.69, IS 115.5 달성.
**텍스트-이미지 생성**: ImageNet 학습 모델이 텍스트 조건에 따라 생성 가능.

의의 및 한계

GEAR은 토크나이저와 생성기를 엔드투엔드로 학습함으로써 기존 STE 기반 방법의 불안정성과 코드북 붕괴 문제를 해결한다. 이는 생성기의 예측 능력을 반영한 토크나이저 인덱스 분포를 학습하여, 더 빠른 수렴과 높은 생성 품질을 달성한다. 또한, 다양한 퀀티파이어와 텍스트-이미지 생성에 일반화 가능하다는 점에서 실용적 가치가 높다. 그러나, 하드-소프트 이중 할당은 복잡한 구조를 요구하며, τ와 λ와 같은 하이퍼파라미터 조정이 필요하다는 점에서 한계가 있다. 또한, STE 대신 사용하는 소프트 할당은 코드북 일관성을 약화시킬 수 있다.

실용적 활용

GEAR은 이미지 생성, 텍스트-이미지 생성, 다양한 퀀티파이어 기반 모델 개발 등 시각 생성 분야에서 널리 활용 가능하다. 특히, 생성 품질 향상과 수렴 속도 개선이 필요한 대규모 모델 학습에 적합하며, 코드북의 일관성과 재구성 능력을 유지하면서도 예측 가능성 높은 인덱스 분포를 학습할 수 있다.