워테르슈타인 그래디언트 플로우를 이용한 이산 자기회귀 사전 확률 학습

Learning Discrete Autoregressive Priors with Wasserstein Gradient Flow

이산 이미지 토크나이저는 일반적으로 두 단계로 훈련됩니다. 첫 번째 단계는 재구성을 위한 것이고, 두 번째 단계는 고정된 토큰 시퀀스에 적합된 사전 모델을 사용하는 것입니다. 이러한 분리는 토크나이저가 나중에 토큰을 생성할 모델을 인식하지 못하게 합니다. 결과적으로, 학습된 토큰은 이미지 정보를 잘 보존할 수 있지만, 자기회귀(AR) 사전 모델이 왼쪽에서 오른쪽으로 예측하기 어렵게 만드는 문제가 있습니다. 우리는 삼각 변분 일관성(TVC)을 사용하여 이러한 불일치를 분석했습니다. TVC는 잠재 변수 학습을 세 가지 일관성 조건으로 분해하며, 즉 조건부-우도 일관성, 사전 일관성 및 사후 일관성입니다. TVC는 두 단계 훈련이 재구성 측면은 유지하지만, 토크나이저 목표 내에서 사전 일관성이 유지되지 않는다는 것을 보여줍니다. 즉, AR 사전 모델이 훈련에 참여하기 전에 전체 토큰 분포가 고정됩니다. 이러한 관점에서 영감을 받아, 우리는 토크나이저 훈련 중에 분포 수준의 사전 일치 신호를 추가하면서 재구성 목표는 변경하지 않았습니다. 이 신호를 워테르슈타인 그래디언트 플로우 업데이트를 사용하여 최적화합니다. 이산 범주형 토큰의 경우, 업데이트는 보조 AR 모델이 토크나이저의 현재 토큰 분포를 추적하고 대상 AR 사전 모델과의 토큰 수준의 대비를 이루는 방식으로 작동합니다. 이 방법은 두 개의 AR 모델을 통과하는 순방향 연산만 필요하며, 어느 모델에도 역전파를 수행하지 않습니다. 결과적으로, wAR-Tok이라는 토크나이저는 AR 손실을 줄이고, CIFAR-10 및 ImageNet 데이터셋에서 유사한 재구성 품질을 유지하면서 생성 FID를 개선합니다.

Discrete image tokenizers are commonly trained in two stages: first for reconstruction, and then with a prior model fitted to the frozen token sequences. This decoupling leaves the tokenizer unaware of the model that will later generate its tokens. As a result, the learned tokens may preserve image information well but still be difficult for an autoregressive (AR) prior to predict from left to right. We analyze this mismatch using Tripartite Variational Consistency (TVC), which decomposes latent-variable learning into three consistency conditions: conditional-likelihood consistency, prior consistency, and posterior consistency. TVC shows that two-stage training preserves the reconstruction side but leaves prior consistency outside the tokenizer objective: the overall token distribution is fixed before the AR prior participates in training. Motivated by this view, we add a distribution-level prior-matching signal during tokenizer training, while keeping the reconstruction objective unchanged. We optimize this signal with a Wasserstein-gradient-flow update. For hard categorical tokens, the update reduces to a token-level contrast between an auxiliary AR model that tracks the tokenizer's current token distribution and the target AR prior. It requires only forward passes through the two AR models and does not backpropagate through either of them. The resulting tokenizer, wAR-Tok, reduces AR loss and improves generation FID on CIFAR-10 and ImageNet at comparable reconstruction quality.