이미지 생성 모델의 벡터 양자화 학습 안정성 개선 방법
Key Point
이미지 생성 모델의 핵심 컴포넌트인 벡터 양자화의 학습 불안정성을 해결하는 실무적 방법론을 제시하므로, 생성형 AI 모델 개발자와 연구자에게 즉시 적용 가능한 지침을 제공한다.
핵심 요약
- 자동회귀 및 마스킹 이미지 생성 모델의 기초가 되는 벡터 양자화(VQ) 토크나이저는 학습 과정이 불안정하다는 문제를 지적했다.
- 인코더-디코더와 코드북 학습이 얽혀 있어 두 모듈이 협력할 때만 작동하므로, 학습이 가중될수록 붕괴되는 취약함이 있다.
- StableVQ는 동적 STE를 통해 인코더의 학습 목표를 안정화하고, 리전 VQ 손실로 코드북이 독립적으로 전체 추적을 보장하며, 분리된 스케줄로 각 모듈이 별도의 학습률을 갖도록 한다.
- 새로운 학습 가능 파라미터를 추가하지 않으면서 공유 투영 코드북 위에 적용 가능하다.
- ImageNet 실험 결과 다양한 코드북 크기와 초기화 설정에서 학습 안정성, 코드북 활용도, 재구성 품질이 일관되게 개선되었다.