기술 요약: 와asserstein 기울기 흐름을 이용한 이산 자기회귀 사전 학습

1. 문제 제기

이산 이미지 토크나이저는 일반적으로 이단계 레시피를 사용하여 훈련됩니다. 먼저, 토크나이저는 이미지를 재구성하도록 훈련됩니다 (조건부 가능도에 초점). 그 다음, 생성된 고정된 토큰 시퀀스에 대해 자기회귀 (AR) 사전이 훈련됩니다. 이 분리는 근본적인 불일치를 초래합니다. 토크나이저는 이미지 정보를 보존하도록 학습되지만, AR 사전의 귀납적 편향을 인지하지 못하게 됩니다. 결과적으로 학습된 토큰은 재구성에는 탁월할 수 있지만, AR 모델이 왼쪽에서 오른쪽으로 순차적으로 예측하기에는 어려울 수 있습니다.

기존 방법들은 순차적 구조를 장려하기 위해 디코더를 수정 (예: Tail Dropout, Diffusion Decoders) 함으로써 이를 완화하려 시도합니다. 그러나 이러한 접근 방식들은 여전히 디코더 중심의 감독에 의존하며, AR 사전의 모델링 난이도를 직접 활용하여 토크나이저의 학습을 안내하지는 못합니다. furthermore, 연속 잠재 공간용으로 설계된 표준 분포 수준 매칭 기법들 (예: WAE-MMD, InfoVAE) 은 양자화를 통한 미분 가능한 경로가 부재하기 때문에 경직된 이산 토크나이저에 직접 적용되지 못합니다.

2. 방법론

이론적 프레임워크: 삼분산 변분 일관성 (TVC)

저자들은 잠재 변수 학습을 분석하기 위해 **삼분산 변분 일관성 (Tripartite Variational Consistency, TVC)**을 도입합니다. TVC 는 생성 목적 함수를 베이즈 정리에서 유도된 세 가지 일관성 조건으로 분해합니다:

1. 조건부 가능도 일관성: 모델의 디코더와 데이터에 고정된 사후 분포 간의 정렬.
2. 사전 일관성: 모델의 사전과 인코더의 집계 사후 분포 간의 정렬.
3. 사후 일관성: 모델의 사후 분포와 데이터에 고정된 사후 분포 간의 정렬.

저자들은 표준 이단계 훈련이 1 단계에서는 가능도 일관성을 만족하고 2 단계에서는 사전 일관성을 만족하지만, 집계 사후 분포가 본질적으로 "AR 친화적"이 되도록 토크나이저에 신호를 제공하지는 못한다고 주장합니다. 반면, ELBO 스타일 훈련은 인스턴스별 사전 일관성을 강제하여 종종 사후 분포 붕괴를 초래합니다.

제안된 솔루션: WGF-DPD 및 wAR-Tok

이 격차를 해소하기 위해, 본 논문은 개별 인스턴스를 제약하지 않고 토큰의 집계 분포를 목표 AR 사전과 매칭하는 **분포적 사전 발산 (Distributional Prior Divergence, DPD)**을 제안합니다. 이산 토큰에 대해 이를 최적화하기 위해 **와asserstein 기울기 흐름 (Wasserstein Gradient Flow, WGF)**을 사용합니다.

핵심 메커니즘인 WGF-DPD는 다음과 같이 작동합니다:

1. 프록시 모델: 토크나이저 ($q_\phi$) 가 생성한 현재 집계 토큰 분포를 추적하도록 경량 보조 AR 모델 ($q_\psi$) 이 훈련됩니다.
2. 점수 기반 업데이트: 집계 사후 분포와 목표 사전 간의 KL 발산에 대한 와asserstein 기울기에서 토크나이저의 업데이트 방향이 유도됩니다. 이는 입자 수준의 기울기 업데이트로 이어집니다:
   $$\nabla_z \mathcal{L}_{DPD} \approx -\nabla_z \log p_\theta(x|z) + (\log q_\psi(z) - \log p_\theta(z))$$
   • 첫 번째 항은 표준 재구성 기울기입니다.
   • 두 번째 항은 사전 매칭 점수로, 프록시 모델과 목표 AR 사전의 로그 확률 차이를 나타냅니다.
3. 효율적 최적화: 중요한 점은 이 업데이트가 프록시 및 목표 AR 모델을 통한 **순방향 전달 (forward passes)**만 필요하다는 것입니다. AR 모델을 통한 역전파는 필요하지 않습니다. 기울기는 비미분 가능한 양자화 단계를 거쳐 Straight-Through Estimator (STE) 를 통해 인코더 매개변수로 전파됩니다.

결과적으로 생성된 토크나이저인 wAR-Tok은 재구성 목적 함수와 WGF-DPD 신호와 함께 공동으로 훈련됩니다.

3. 주요 기여

• TVC 프레임워크: 잠재 변수 학습의 이론적 분해로, "사전 일관성"을 별도의 목적 함수로 명시적으로 식별하여 이단계 훈련이 토크나이저가 사전의 모델링 난이도를 인지하지 못하게 하는 이유를 설명합니다.
• WGF-DPD 알고리즘: 와asserstein 기울기 흐름을 사용하여 집계 토큰 분포를 AR 사전과 정렬하는 이산 토크나이저를 위한 새로운 최적화 방법입니다. 이는 이산 양자화기를 통한 역전파의 비실현성이나 적대적 매칭의 불안정성을 피합니다.
• wAR-Tok: WGF-DPD 신호를 토크나이저 훈련에 통합하는 실용적 구현으로, 인코더 측 정렬 신호를 생성하기 위해 보조 AR 모델을 통한 순방향 전달만 필요합니다.

4. 실험 결과

이 방법은 1D 토큰화, Tail Dropout, Diffusion Decoders, Tied Timesteps 등을 포함한 베이스라인과 비교하여 CIFAR-10 및 ImageNet 256×256에서 평가되었습니다.

• CIFAR-10:
  • 일치된 재구성 품질 (rFID $\approx$ 5.0) 하에서, wAR-Tok 은 생성 FID (gFID) 14.50 및 AR 생성 손실 (gLoss) 0.42를 달성하여 가장 강력한 베이스라인 (ImageFolder/One-D-Piece: gFID 17.60, gLoss 0.78) 을 크게 앞섰습니다.
  • 일치된 훈련 반복 횟수 하에서, wAR-Tok 은 gFID 11.30 및 gLoss 0.34에 도달하여 더 빠른 수렴과 더 높은 품질을 입증했습니다.
• ImageNet 256×256:
  • 베이스라인과 동일한 인코더 - 디코더 및 AR 사전 아키텍처를 사용하여, wAR-Tok 은 gFID 를 6.32 에서 5.42로, 평가 AR 손실을 7.89 에서 7.07로 감소시켰으며, 재구성 품질에서 약간의 트레이드오프 (rFID 2.05 $\to$ 2.35) 가 있었습니다.
• 훈련 역학: AR 훈련 곡선은 베이스라인에 비해 wAR-Tok 이 훈련 내내 더 낮은 교차 엔트로피 손실과 더 높은 토큰 예측 정확도를 유지함을 보여주며, 학습된 토큰이 본질적으로 AR 사전이 모델링하기 쉽다는 것을 나타냅니다.

5. 의의 및 주장

본 논문은 wAR-Tok이 토크나이저 훈련 중에 분포 수준의 사전 매칭 신호를 도입함으로써 이산 토크나이저와 자기회귀 사전 간의 불일치를 효과적으로 해결한다고 주장합니다.

• 효율성: 이 방법은 계산적으로 효율적이며, 훈련 단계당 프록시 및 목표 두 번의 순방향 전달만 추가하고 AR 모델을 통한 값비싼 역전파를 피합니다.
• 일반성: AR 사전에 초점을 맞추고 있지만, TVC 프레임워크와 WGF-DPD 접근 방식은 사전 무관하게 제시되며, 계산 가능한 로그 확률을 가진 모든 이산 사전에 적용 가능합니다.
• 겸손함: 저자들은 의도적으로 최근 토크나이저에서 일반적인 보조 기법들 (예: DINO 정렬, 증류, Diffusion-디코더 후 훈련) 을 배제하여 WGF-DPD 의 효과를 분리했으므로 절대적인 최첨단 (SOTA) 성능을 주장하지 않는다고 명시적으로 밝힙니다. 결과는 통제된 환경에서 제안된 정렬 목적 함수의 유효성을 보여줍니다.

이 연구는 단순한 재구성 가능성뿐만 아니라 하류 생성 모델에 의한 토큰의 "예측 가능성"을 명시적으로 최적화하는 것이 더 효율적이고 고품질의 이산 생성 시스템으로 이어진다는 것을 시사합니다.