StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training
Encoder-Decoder와 Codebook 간의 상호 의존성 문제를 해결하여 VQ 학습의 안정성과 코드북 활용도를 높이는 프레임워크
논문이 다루는 내용
최근 공유 투영(shared-projection) 방식의 코드북 기법이 발전했음에도 불구하고, VQ 학습의 불안정성은 여전히 해결해야 할 과제로 남아 있습니다. 기존 방식은 인코더-디코더와 코드북 학습이 서로의 성능에 의존하는 얽힘(entanglement) 문제로 인해 학습 과정에서 시스템이 쉽게 붕괴되는 경향이 있습니다. 본 논문은 각 모듈이 독립적인 역할을 수행할 수 있도록 학습 목적 함수를 재정의하는 StableVQ를 제안합니다. 구체적으로 Dynamic STE, Region VQ Loss, Decoupled Schedule이라는 세 가지 기법을 통해 모듈 간의 간섭을 최소화합니다. 실험 결과, ImageNet 데이터셋에서 코드북 활용도와 재구성 품질을 동시에 개선하며 안정적인 학습 성능을 입증했습니다.
핵심 결과
-
Dynamic STE: 이산 정규화 하에서도 인코더가 재구성 공간을 안정적으로 최적화하도록 지원
-
Region VQ Loss: 인코더의 진동에 의존하지 않고 코드북이 인코더 출력 분포를 추적할 수 있도록 설계
-
Decoupled Schedule: 인코더-디코더와 코드북의 서로 다른 최적화 역학을 고려한 독립적 학습률 스케줄링 적용
실무에서 볼 만한 점
추가적인 파라미터 없이 기존 VQ 구조에 즉시 적용 가능하며, 학습 중 발생하는 코드북 붕괴 문제를 방지할 수 있습니다.
읽을 때 확인할 점
-
기존 VQ-VAE/VQGAN 학습 시 코드북 활용도가 낮을 때 Dynamic STE 적용 테스트
-
인코더와 코드북의 학습 속도 차이를 조절하는 Decouled Schedule 실험
-
공유 투영(shared-projection) 기반 코드북 구조에 Region VQ Loss 통합