이미지 합성을 위한 안내된 종단간 자기회귀 학습(GEAR)

대부분의 비전 생성 파이프라인은 토크나이저를 재구성 목적으로 학습하고 고정하여 생성기가 예측하기 쉬운 토큰 분포를 얻지 못한다. GEAR는 autoregressive(AR) 생성기의 표현 정렬 신호를 미분 가능한 소프트 어사인먼트로 토크나이저에 전달해 토큰의 예측 가능성을 높였다. 이 접근은 학습 수렴을 크게 가속하고 패치 수준의 지역 구조를 강화하여 최종 샘플 품질을 개선했다.

전통적인 두 단계 파이프라인에서는 토크나이저가 재구성 목표에만 맞춰져 생성기가 예측하기 쉬운 이산 인덱스 분포를 반영하지 못한다. 이로 인해 생성기는 고분산의 세부 정보를 가진 라텐트를 모델링해야 하고 토크나이저와 생성기 목표 사이에 충돌이 발생한다. GEAR는 이 충돌을 완화하기 위해 생성기의 표현 정렬 목표를 토크나이저로 안내하는 방법을 채택했다. GEAR의 핵심은 각 위치의 코드북 할당을 하드(one-hot)와 소프트(weighted mixture)로 동시에 읽어들여 두 경로를 분리하는 점이다. 하드 경로는 inference와 동일한 불연속 토큰 시퀀스를 AR에 전달하여 next-token prediction(NTP) 손실과 하드 정렬 손실이 오직 AR만 업데이트하도록 만든다. 소프트 경로는 온도 τ로 제어되는 softmax를 이용해 AR 임베딩 테이블의 가중합을 구성하고, 이 미분 가능한 경로를 통해 정렬 손실이 토크나이저(인코더·코드북·디코더)만 업데이트하도록 한다. 이 설계는 토크나이저가 DINOv2와 같은 외부 인코더를 직접 모방하는 대신 토큰 분포를 더 예측 가능하게 재조직하게 만든다. 결과적으로 토크나이저의 패치 수준 DINOv2 유사도는 감소하지만 AR의 각 레이어는 패치 단위에서 더 DINOv2-like한 특징을 획득했다. 이 변화는 NTP 학습 난이도를 낮추어 수렴 속도를 크게 향상시키고 샘플의 공간적 일관성을 높였다.

전체 접근은 VQ 토크나이저와 causal AR 트랜스포머를 동일한 학습 루프에서 최적화하되 업데이트 경로를 분리하는 것이다. 토크나이저는 재구성 손실(ℒ_VQ)과 소프트 정렬 손실(λℒ_align^s)로만 갱신되고 AR은 NTP 손실(ℒ_NTP)과 하드 정렬 손실(λℒ_align^h)로만 갱신된다. 이렇게 하면 NTP 손실이 토크나이저로 흘러 코드북 붕괴를 일으키는 상황을 원천 차단한다. 핵심 메커니즘은 할당 행렬 A의 두 가지 읽기이다. 각 위치 i에서 A_i,k = -||z_i - c_k||^2로 점수를 만들고 하드 읽기는 argmax로 one-hot 인덱스를 선택해 AR 임베딩 테이블 E에서 해당 행을 뽑는다. 소프트 읽기는 π_i = softmax(A_i/τ)로 가중합 u_i^s = π_i E를 만들며 이 항은 인코더와 코드북에 미분 경로를 제공한다. 학습 절차는 이중-브랜치 전방 전달 후 두 종류의 정렬 손실과 NTP 손실을 계산하는 것이다. 소프트 정렬은 AR 백본을 동결한 채 토크나이저만 업데이트하도록 그래디언트를 제한하고, 하드 정렬과 NTP는 AR 파라미터만 업데이트한다. 기본 하이퍼파라미터로 τ=0.1과 λ=0.5가 실험 전반에서 균형 잡힌 성능을 보였고 STE 대체 실험에서 STE가 코드북 붕괴를 유발함이 관찰되었다.

클래스 조건부 ImageNet 256×256 실험에서 GEAR는 동일한 AR 및 학습 예산 조건에서 LlamaGen-REPA 대비 일관된 품질 향상을 보였다. 예컨대 300-epoch 설정에서 B/L/XL 스케일에서 gFID와 IS가 개선되었고, CFG 적용 시 gFID가 여러 스케일에서 감소하는 추세가 관찰되었다. 장기 스케줄과 크기 확장 실험에서도 모델 크기가 커질수록 GEAR의 이득이 유지되었다. 텍스트-투-이미지 GPIC 단일 에폭 실험에서는 동일한 Qwen3-1.7B 텍스트 인코더와 데이터에서 GEAR가 LlamaGen-REPA보다 FDD와 기타 품질 지표를 꾸준히 개선했다. AR을 GEAR로 학습한 토크나이저와 고정 토크나이저를 비교한 추가 실험에서 GEAR 토크나이저는 NTP 손실에서 약 2.5배 빠른 수렴과 REPA 정렬 손실에서 약 11.1배 빠른 수렴을 달성했다. 이 가속은 토크나이저가 더 예측 가능한 인덱스 분포를 형성했기 때문으로 해석되며 고정 토크나이저 파이프라인에 드롭인했을 때에도 이득이 전이되었다. 다수의 소스-별 실험과 ablation이 수행되었고 주요 발견으로는 소프트 온도 τ의 영향, 정렬 계수 λ의 최적화값(기본 λ=0.5), STE 사용 시의 붕괴 위험, 그리고 VQVAE/LFQ/IBQ 전반에서의 이득이 확인되었다. 온도 τ는 너무 작으면 소프트 경로가 하드와 지나치게 유사해져 가이던스가 약해지고 너무 크면 일관성은 떨어지나 gFID는 개선되는 경향이 있었다. 전체적으로 소프트 어사인먼트를 통한 안내가 안정성과 성능을 동시에 개선하는 핵심 요소로 드러났다.

전체 아키텍처는 학습 가능한 VQ 토크나이저(인코더 ℰ, 코드북 𝒞, 디코더 𝒟)와 L-layer의 causal transformer AR로 구성된다. AR의 입력은 토크나이저에서 생성된 인덱스 q에 대해 AR 전용 임베딩 행렬 E에서 뽑은 행을 사용하며 소프트 브랜치에서는 π_i E의 가중합을 이용한다. 토크나이저 손실은 재구성 ℒ_VQ(재구성, LPIPS, GAN, 엔트로피, 커밋먼트 항 포함)이며 AR은 ℒ_NTP를 최적화한다. 핵심 수학은 할당 행렬 A와 온도 τ를 통한 softmax π=softmax(A/τ)이다. 이 π는 u^s=πE를 통해 AR 임베딩 공간으로 매핑되고 정렬 손실 ℒ_align^s가 계산되면 그 그래디언트는 AR 백본의 상위 파라미터를 동결한 상태로 오직 토크나이저 파라미터로 흐른다. 하드 브랜치의 u^h는 argmax(A)로 결정되는 불연속 임베딩을 사용해 ℒ_NTP와 ℒ_align^h가 AR 파라미터만 갱신하도록 보장한다. 표현 분석은 CKNNA와 CKA로 레이어별 유사도를 평가하고 Moran’s I 및 FFT 고주파 비율로 공간적 자기상관과 고주파 에너지 분포를 측정했다. 이 분석에서 GEAR의 AR은 패치 수준 CKA/CKNNA와 Moran’s I에서 DINOv2 기준선에 더 근접하는 것으로 관찰되었다. 학습 세부사항으로는 ImageNet-1K 256×256 학습, GPIC 단일 에폭 텍스트-투-이미지 비교, 기본 τ=0.1과 λ=0.5 사용, 일부 장기 실험에서는 400k 스텝 공동 학습을 수행했다. 실험 구현은 AR 임베딩 테이블과 코드북을 분리하고 소프트 브랜치의 전방 전달을 정렬 레이어 ℓ까지만 실행해 계산 비용을 낮추는 방식으로 최적화되었다. STE 대체 실험에서는 STE 사용 시 코드북 붕괴(gFID≈105, rFID 대폭 악화)가 재현되어 soft-assignment의 필요성이 실험적으로 확증되었다. 논문은 또한 CFG(결과 샘플링에서 가이드 스케일)의 효과를 탐색하여 샘플 품질과 다양성 트레이드오프를 보고했다.

논문이 명시한 주요 한계는 이산 VQ 토크나이저로 인한 재구성 상한이다. VQ-AR 파이프라인에서는 토크나이저의 다운샘플률이 시퀀스 길이와 직접 연결되어 있어 재구성 fidelity를 개선하려면 시퀀스 길이와 계산량이 증가한다. 따라서 연속 VAE 기반의 REPA-E처럼 더 높은 재구성 성능을 얻기 위해서는 토크나이저-AR 간의 압축·계산 분리 또는 다중 토큰 예측과 같은 구조적 변화가 필요하다.