Beyond Stationarity: Rethinking Codebook Collapse in Vector Quantization

이 논문은 벡터 양자화에서 발생하는 코드북 붕괴 현상의 근본 원인이 인코더의 비정상적 변화에 있음을 규명하고, 이를 해결하기 위해 제안한 NSVQ 와 TransVQ 두 가지 새로운 방법을 통해 CelebA-HQ 데이터셋에서 코드북 활용도와 재구성 품질을 획기적으로 개선했음을 보여줍니다.

핵심

🎨 1. 배경: "완벽한 팔레트"와 "죽은 화가"

인공지능이 그림을 그릴 때, 수백만 가지 색상을 모두 기억할 수는 없습니다. 대신 **약 1,000~10,000 개의 색상 (코드북, Codebook)**을 미리 준비해 두고, 그림을 그릴 때 이 중에서 가장 비슷한 색을 골라 사용합니다. 이를 '벡터 양자화'라고 합니다.

하지만 여기서 문제가 생깁니다.
훈련을 시작하면 인공지능은 그림을 더 잘 그리기 위해 계속 학습합니다. 그런데 이상한 일이 발생합니다. 수천 개의 색상 중 일부만 계속 쓰이고, 나머지 수천 개는 아예 쓰이지 않게 되는 것입니다. 마치 화가에게 100 개의 물감을 주었는데, 오직 빨강과 파랑만 계속 쓰고 나머지 98 개는 말라버린 채 창고 구석에 방치된 것과 같습니다.

이를 **'코드북 붕괴 (Codebook Collapse)'**라고 부릅니다. 인공지능이 더 많은 색상 (코드) 을 쓸수록 오히려 효율이 떨어지는 역설적인 상황입니다.

🔍 2. 원인 발견: "움직이는 카메라"와 "고정된 표적"

저자들은 이 문제의 원인을 이론적으로 찾아냈습니다. 바로 "움직이는 카메라 (Encoder)" 때문입니다.

• 비유: imagine you are taking photos of a moving target (움직이는 표적) with a camera that keeps changing its focus (초점이 계속 변하는 카메라).
  • 처음에는 카메라가 표적을 잘 찍습니다.
  • 하지만 카메라가 학습되면서 초점과 각도가 조금씩 변합니다 (Encoder Drift).
  • 문제는 표적 (데이터) 이 변하는 속도에 맞춰서, 미리 준비해 둔 색상 (코드북) 들이 따라가지 못한다는 것입니다.
  • 특정 색상이 한 번 쓰이지 않게 되면, 카메라가 계속 변하기 때문에 그 색상은 영원히 다시 쓰이지 않게 됩니다. 결국 그 색상은 '죽은 코드 (Dead Code)'가 되어버립니다.

기존 방법들은 "색을 다시 섞어보자"거나 "무작위로 교체하자"는 식의 임시방편 (휴리스틱) 이었습니다. 하지만 저자들은 **"카메라가 움직이니까, 색상도 함께 움직여야 한다"**는 핵심 통찰을 얻었습니다.

💡 3. 해결책: 두 가지 새로운 방법

저자는 이 문제를 해결하기 위해 두 가지 창의적인 방법을 제안했습니다.

🚀 방법 1: NS-VQ (움직임을 공유하는 방법)

• 비유: 카메라가 움직일 때, 사용하지 않은 색상들도 그 움직임을 미리 감지하고 따라가게 하는 것입니다.
• 어떻게?: 지금 그림을 그릴 때 '빨강' 색을 썼다면, '파랑'이나 '초록' 같은 다른 색상들도 "아, 카메라가 이쪽으로 움직였구나!"라고 알아채고 살짝 움직이게 합니다.
• 효과: 어떤 색이든 한 번도 쓰이지 않고 방치되는 일이 사라집니다. 모든 색이 살아있게 됩니다.

🧠 방법 2: TransVQ (지능적인 변신)

• 비유: 색상 팔레트 전체를 스마트하게 변신시키는 마법 거울을 거치게 하는 것입니다.
• 어떻게?: 카메라가 변하면, 색상들끼리 서로 대화하며 (트랜스포머 구조) "우리가 함께 움직여야 해!"라고 결정합니다. 마치 한 무리의 춤꾼들이 리더의 동작에 맞춰 전체 군무의 방향을 바꾸는 것처럼요.
• 효과: 색상들이 개별적으로 움직이는 게 아니라, 전체 팔레트가 하나의 유기체처럼 데이터의 변화에 맞춰 유연하게 적응합니다.

🏆 4. 결과: 더 좋은 그림, 더 많은 색상 활용

이 두 방법을 실험해 보니 놀라운 결과가 나왔습니다.

1. 색상 활용도 100%: 준비한 수천 개의 색상 중 하나도 버려지지 않고 모두 사용되었습니다.
2. 더 선명한 그림: 색을 더 잘 골라 쓸 수 있게 되니, 인공지능이 그린 얼굴이나 사물의 질이 기존 방법보다 훨씬 좋아졌습니다.
3. 이론적 증명: 단순히 "잘 되네"가 아니라, "왜 잘 되는지"에 대한 수학적인 이유도 증명했습니다.

📝 요약: 한 줄로 정리하면?

"인공지능이 그림을 그릴 때, 카메라가 변하면 색상 팔레트도 함께 변해야 한다. 그래야 모든 색이 살아나고 더 멋진 그림을 그릴 수 있다!"

이 연구는 인공지능이 더 효율적이고 강력한 모델을 만들 수 있는 새로운 이론적 토대를 마련했다는 점에서 매우 중요합니다.

기술 요약

1. 문제 정의 (Problem Definition)

벡터 양자화 (Vector Quantization, VQ) 는 VQ-VAE, VQ-GAN, 잠재 확산 모델 (Latent Diffusion Models) 등 현대 생성 모델의 핵심 구성 요소입니다. 그러나 VQ 는 훈련 과정에서 **코드북 붕괴 (Codebook Collapse)**라는 심각한 문제에 직면해 있습니다. 이는 코드북의 크기가 커질수록 많은 수의 코드 벡터가 훈련 중 사용되지 않고 비활성화되는 현상입니다.

• 기존 접근법의 한계: 기존 연구들은 확률적 양자화, 코드북 재설정, 분포 정규화 등 다양한 휴리스틱 (heuristic) 방법을 제시했으나, 이러한 방법들은 이론적 근거가 부족하고 실제 성능이 일관되지 않습니다.
• 핵심 원인: 본 논문은 VQ-VAE 에서 **인코더 업데이트의 비정상성 (Non-stationarity)**이 코드북 붕괴의 근본 원인임을 이론적으로 규명했습니다. 인코더 파라미터가 훈련 중에 변화 (drift) 하면, 선택되지 않은 코드 벡터들은 업데이트를 받지 못하고 점차 비활성화되어 '죽은 코드 (dead codes)'가 됩니다.

2. 방법론 (Methodology)

저자는 인코더의 비정상적 움직임을 해결하기 위해 두 가지 새로운 VQ 방법을 제안합니다.

가. 비정상성 벡터 양자화 (NS-VQ, Non-Stationary Vector Quantization)

• 개념: 인코더의 드리프트 (drift) 를 선택되지 않은 코드 벡터에도 전파하여 모든 코드가 업데이트되도록 합니다.
• 구현 방식:
  • 커널 기반 규칙 (Kernel-based rule) 을 사용하여 현재 배치의 입력 $x_i$에 대한 인코더 업데이트 ($\Delta E$) 를 다른 코드 $c_{qj}$에도 적용합니다.
  • 신경망 탄젠트 커널 (NTK) 을 가우스 RBF 커널로 근사화하여 계산 효율성을 높였습니다.
  • 새로운 임베딩 손실 함수: 선택된 코드뿐만 아니라 비선택된 코드에도 가중치 (거리 기반) 를 부여하여 업데이트를 유도합니다.
  • 수정된 STE (Straight-Through Estimator): 인코더 드리프트 하에서 안정성을 높이기 위해 기존 STE 업데이트 규칙을 재정의했습니다.

나. 트랜스포머 기반 벡터 양자화 (TransVQ, Transformer-based Vector Quantization)

• 개념: 인코더 업데이트가 특징 표현에 드리프트 항을 추가한다면, 코드북 자체도 학습 가능한 변환을 통해 적응적으로 조정해야 한다는 아이디어입니다.
• 구현 방식:
  • 코드북 $C$를 고정된 상태로 두고, 가중치 $\phi$를 학습하는 가벼운 매핑 함수 $P_\phi(\cdot)$를 도입합니다.
  • 각 코드 벡터를 토큰으로 간주하고, 단일 헤드의 선형 어텐션 (Linear Attention) 레이어와 작은 MLP 를 통과시켜 변환된 코드북 $C'$을 생성합니다.
  • 이론적 보장: SimVQ 와 같은 기존 선형 매핑 방법과 달리, TransVQ 는 k-means 해로의 수렴 조건을 유지하면서도 전체 코드북의 부드러운 적응을 가능하게 합니다.

3. 주요 기여 (Key Contributions)

1. 이론적 분석: VQ-VAE 의 인코더 업데이트 비정상성이 코드북 붕괴의 근본 원인임을 최초로 이론적으로 증명했습니다.
2. NS-VQ 제안: 커널 기반 업데이트 규칙을 통해 인코더 드리프트를 비선택 코드로 전파하여 코드북 활용도를 극대화하는 새로운 변형을 제시했습니다.
3. TransVQ 제안: k-means 수렴 조건을 해치지 않으면서 학습 가능한 매핑을 통해 전체 코드북을 적응적으로 변환하는 트랜스포머 기반 아키텍처를 개발했습니다.
4. 실험적 검증: CelebA-HQ 데이터셋을 통한 광범위한 실험을 통해 제안된 방법들이 기존 VQ 변형들보다 우수한 재구성 품질과 거의 완벽한 코드북 활용도를 달성함을 입증했습니다.

4. 실험 결과 (Results)

실험은 CelebA-HQ (256x256) 데이터셋에서 수행되었으며, rFID, LPIPS, SSIM 지표를 사용하여 평가되었습니다.

• 코드북 활용도 (Utilization): 기존 VQGAN-FC 나 VQVAE2 는 코드북 크기가 커질수록 활용도가 급격히 떨어지는 반면, NS-VQ 와 TransVQ 는 모든 코드북 크기 (64~8912) 에서 100% 에 가까운 활용도를 유지했습니다.
• 재구성 품질 (Reconstruction Quality):
  • rFID (낮을수록 좋음): TransVQ 가 64 차원 코드에서 13.70, NS-VQ 가 14.01 로, 기존 최첨단 방법 (SimVQ: 14.37, VQGAN-FC: 17.57) 보다 우수한 성능을 보였습니다.
  • SSIM/LPIPS: TransVQ 와 NS-VQ 는 더 높은 SSIM 과 더 낮은 LPIPS 값을 기록하여 이미지 재구성 품질이 뛰어남을 확인했습니다.
• 배치 크기 영향: 이론적 분석대로 배치 크기가 커질수록 코드북 붕괴가 완화되고 재구성 오류 (rFID) 가 감소하는 경향을 실험적으로 확인했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 이론과 실전의 연결: 코드북 붕괴가 단순한 최적화 실패가 아니라 비정상성 (non-stationarity) 에 기인한 구조적 문제임을 규명함으로써, VQ 기반 생성 모델에 대한 이론적 기반을 강화했습니다.
• 확장성: 제안된 방법들은 이미지 재구성뿐만 아니라 대규모 시각 - 언어 모델 (VLM) 및 확산 모델 등 다양한 생성 작업에 적용 가능한 확장 가능한 기반을 제공합니다.
• 한계 및 향후 과제: NS-VQ 는 하이퍼파라미터 ($\sigma^2$) 조정이 필요하고, TransVQ 는 추가적인 계산 비용이 발생한다는 한계가 있습니다. 향후 연구에서는 이러한 매개변수를 자동으로 제어하는 적응형 메커니즘 개발 및 확산/자기회귀 아키텍처와의 통합이 기대됩니다.

이 논문은 VQ 기술의 한계를 극복하고 더 안정적이고 효율적인 생성 모델을 구축하기 위한 중요한 이론적, 실용적 통찰을 제공합니다.