엔코더와 코드북의 얽힌 최적화를 분리해 VQ 토크나이저를 안정화하는 StableVQ

무엇인가

이 논문은 VQ-VAE 계열의 이산 시각 토크나이저에서 코드북 붕괴와 학습 불안정을 다룬다. VQ 토크나이저는 이미지 특징을 학습된 코드북의 이산 토큰으로 바꿔 자기회귀 생성 모델이나 마스크 생성 모델이 쓸 수 있게 하므로, 코드북 활용률과 재구성 품질이 하류 생성 품질의 상한을 좌우한다. 최근 shared-projection 코드북은 코드를 f_theta(e_k)처럼 공유 함수로 재매개변수화해 gradient sparsity를 완화하고 활용률을 크게 높였지만, 저자들은 학습 안정성이 여전히 해결되지 않은 핵심 문제라고 주장한다. 그 원인을 엔코더-디코더와 코드북 학습이 서로 얽혀 있어 각 모듈이 독립적으로 자기 역할을 다하지 못하는 데서 찾는다. 논문이 제시하는 실패 모드는 세 가지다. 코드 스케일이 토큰 스케일보다 작으면 장기 저활용과 commitment loss 급등, 심하면 NaN으로 간다. 코드 스케일이 토큰 스케일보다 크면 활용률이 빠르게 오르다가 손실 신호가 포화되면서 대부분의 코드가 영구 dead code가 된다. 코드북이 잘 쓰이던 상황에서도 엔코더 출력 분포가 급변할 때 코드북이 추적을 놓치면 잘못된 STE gradient가 발산을 증폭해 활용률이 순간적으로 붕괴한다.

어떻게 동작하나

StableVQ는 이 문제를 모듈별 책임 분리로 풀며, 세 가지 구성 요소를 제안한다. 첫째 Dynamic STE는 엔코더의 학습 목표를 바로잡는다. 같은 코드에 여러 토큰이 배정됐을 때 상대적으로 먼 토큰의 STE gradient가 가장 신뢰하기 어렵다는 관찰에서 출발해, 각 토큰에 w_ij = sg[d*_{k*_ij} / ||z_ij - f_theta(e_{k*_ij})||^2]라는 가중치를 준다. 배정된 코드에 가장 잘 맞는 토큰은 w_ij = 1로 온전한 gradient를 받고, 상대 오차가 커질수록 가중치가 줄어든다. 이를 z_hat_ij = w_ij·z_ij + sg[z_hat_ij - w_ij·z_ij]로 STE에 반영해 먼 토큰의 gradient를 약화한다. 모든 토큰이 잘 맞으면 w_ij가 1에 가까워져 표준 STE로 되돌아가므로 별도 임계값 하이퍼파라미터가 필요 없다. 둘째 Region VQ Loss는 코드북의 학습 목표를 다시 설계한다. 표준 VQ loss는 현재 선택된 코드에만 직접 목표를 주고, shared projection을 통한 나머지 코드 신호는 간접적이고 약해진다. 이를 해결하려고 길이 W의 FIFO 큐로 최근 활성 코드 집합 A_t와 지속적 비활성 집합 N_t를 나눈다. 각 활성 소스 코드 k는 배정된 토큰 수 n_k에 비례하는 quota q_k를 받아 N_t 안의 가장 가까운 q_k개 코드로 목표를 전파한다. 활성 코드는 자기 토큰 평균을, 비활성 코드는 전파받은 소스 목표의 평균을 t_j로 삼고, L_code = (1/|M_t|) Σ_{k∈M_t} ||f_theta(e_k) - sg[t_k]||^2로 코드북을 갱신한다. 셋째 Decoupled Schedule은 엔코더-디코더와 코드북이 서로 다른 최적화 동역학을 가진다고 본다. 엔코더-디코더는 복잡한 다목적 재구성 문제라 warmup-plus-annealing이 유리하고, 코드북은 엔코더 출력 분포를 계속 추적해야 하므로 초기부터 충분히 큰 상수 학습률이 유리하다는 것이다. StableVQ는 shared-projection 코드북 위에 올라가며 추가 학습 파라미터가 없다고 주장한다.

무엇과 다른가

실험은 ImageNet 256×256에서 VQGAN-style 엔코더-디코더, 다운샘플링 계수 f=16, 이미지당 16×16=256개 토큰으로 수행됐다. 평가 지표는 rFID, LPIPS, 그리고 ImageNet 검증셋에서 활성화된 코드 비율로 계산한 코드북 활용률이다. 주요 베이스라인은 shared-projection 계열의 SimVQ와 FVQ이며, 각 원래 설정을 재구현하고 save-then-reload 대신 on-the-fly 재구성으로 공정하게 측정했다고 밝힌다. StableVQ는 기본적으로 가장 단순한 단일 선형 shared projector를 쓴다. 논문은 SimVQ와 FVQ가 각자의 표준 설정에서 100% 활용률을 달성하지만, SimVQ는 상수 학습률 때문에 annealing이 없어 재구성 품질이 크게 떨어지고, FVQ는 ViT block 깊이와 patch size 같은 프로젝터 설계에 의존하며 patch embedding 때문에 코드북 크기가 완전제곱수로 제한된다고 지적한다. 반면 StableVQ는 단일 선형 투영만으로 더 복잡한 프로젝터를 쓰는 방법과 대등하거나 더 나은 재구성 품질을 보이고, 프로젝터별 설계나 구조적 제약 없이 더 다양한 난이도에서 완전 활용률을 유지한다고 주장한다. 또한 코드북-토큰 분포 불일치 상황에서 활용률 회복을 측정하는 UR-AUC를 도입했고, SimVQ와 FVQ는 제한된 불일치 조건에서만 느리게 회복하지만 StableVQ는 다양한 분포 관계에서 빠르게 완전 활용률을 되찾는다고 보고한다. 생성 실험은 IBQ 설정을 따라 StableVQ 토큰 위에 클래스 조건부 자기회귀 트랜스포머를 학습했고, ImageNet 256×256에서 경쟁력 있는 결과를 Table 2에 제시한다고만 밝힌다. 제공된 원문 발췌에는 Table 1과 Table 2의 구체적인 rFID, LPIPS, UR-AUC 수치가 실려 있지 않아 정확한 값은 제시할 수 없다.

어떻게 쓰나

파일럿 연구와 절제 실험은 각 구성 요소의 필요성을 뒷받침한다. Dynamic STE 파일럿에서는 코드북을 고정하고 엔코더-디코더만 학습시켰을 때 표준 STE가 commitment loss 급등과 재구성 손실 진동을 일으켜 결국 NaN으로 발산하는 반면, Dynamic STE는 손실을 안정적으로 유지했다. Region VQ Loss 파일럿에서는 엔코더를 고정하고 코드북만 최적화했을 때 표준 VQ loss가 5000스텝 후에도 활용률 약 12.5%에 머문 반면, Region VQ Loss는 500스텝 만에 완전 활용률에 도달하고 유지했다. Decoupled Schedule 파일럿에서는 FVQ 구조에서 코드북을 상수 학습률로 두어도 성능 저하가 없었지만 엔코더-디코더를 상수 학습률로 두면 품질이 떨어졌고, SimVQ 구조에서는 코드북이 복잡한 스케줄보다 안정적이고 충분히 높은 상수 학습률에서 이득을 봤다. 코드북 확장 절제 실험은 코드북을 매우 작은 범위로 초기화한 흔한 설정에서 단일 선형 프로젝터와 warmup-plus-annealing, 최대 학습률 1e-4를 사용했고, Decoupled Schedule을 쓸 때 코드북 학습률은 상수 1e-3으로 두었다. 결과적으로 Region VQ Loss만 쓰면 코드북이 너무 느리게 갱신돼 NaN 붕괴가 났고, Dynamic STE만 쓰면 엔코더는 안정되지만 활용률은 매우 낮았으며, Decoupled Schedule만 쓰면 분포 추적은 좋아지지만 표준 VQ loss 아래에서는 완전 활용률에 못 미쳤다. Region VQ Loss를 Dynamic STE나 Decoupled Schedule 중 하나와 결합하면 완전 활용률과 강한 재구성 품질이 회복됐고, 세 요소를 모두 쓰면 가장 완전한 해법이 됐다. 코드북 축소 절제 실험은 코드북 차원 1024와 FVQ 설정의 더 강한 ViTBlock shared projector를 사용했으며, FVQ는 두 초기화 모두에서 완전 활용률을 달성하지 못하고 초기화에 민감했지만 Region VQ를 추가하면 초기화와 무관하게 완전 활용률과 개선된 재구성 품질을 얻었다.

전제와 한계

개발자 관점에서 이 논문은 shared-projection 코드북을 쓰는 VQ 토크나이저 학습에서 코드북 활용률 저하, commitment loss 스파이크, NaN, 초기화 민감성이 나타날 때 참고할 만하다. 적용 지점은 명확하다. 엔코더 쪽에는 토큰별 상대 거리 기반 STE 가중치를 넣고, 코드북 쪽에는 FIFO 큐와 비활성 코드로의 목표 전파를 넣고, 옵티마이저 쪽에는 엔코더-디코더와 코드북의 학습률 스케줄을 분리한다. 특히 단일 선형 프로젝터만으로도 복잡한 프로젝터에 필적하는 결과를 노린다는 점은 구현 부담을 낮춘다. 다만 실무에서는 코드북 활용률만 보지 말고 commitment loss 안정성, 재구성 지표, 초기화 변화에 따른 UR-AUC 성격의 회복 곡선을 함께 확인해야 한다. Region VQ Loss는 FIFO 큐 길이 W와 quota 전파 규칙이라는 설정값을 가지며, Decoupled Schedule도 코드북 상수 학습률 값을 정해야 하므로 완전히 하이퍼파라미터가 없는 방법은 아니다.

저자들이 명시한 별도의 한계 절은 제공된 원문에 없다. 대신 전제는 분명하다. StableVQ는 shared-projection 코드북 위에서 동작하고, 실험은 ImageNet 256×256, VQGAN-style f=16, 256토큰, 특정 코드북 크기와 초기화 설정에 한정된다. 생성 실험도 IBQ 설정의 클래스 조건부 자기회귀 트랜스포머에 의존하며, StableVQ 학습에는 이전 연구를 따른 판별자 기반 적대적 감독이 쓰인다. 따라서 다른 해상도, 다른 도메인, 더 큰 코드북, 다른 토크나이저 구조에서도 같은 안정성과 활용률이 보장되는지는 추가 검증이 필요하다. 그럼에도 논문의 핵심 주장은 분명하다. VQ 학습의 오랜 불안정은 양자화 자체의 한계가 아니라 엔코더-디코더와 코드북의 얽힌 최적화 목표에서 비롯되며, 모듈별 책임을 분리하면 완전 코드북 활용률을 취약한 우연이 아니라 원리적으로 보장 가능한 속성으로 만들 수 있다는 것이다.