VAE 계열 생성 모델에서의 자주 언급되는 문제점은 다음과 같다. 디코더를 PixelCNN이나 WaveNet처럼 매우 강력한 autoregressive 모델로 만들수록, 인코더가 열심히 정보를 latent에 담지 않아도 디코더가 알아서 복원을 잘 해버리는 경우가 생긴다는 점이다. 흔히 posterior collapse라고 부르는 현상이다. 요약하면 “latent가 무시되고, 의미 있는 표현을 못 배운다”는 맥락이다.
VQ-VAE는 이 문제를 구조적으로 회피하고자 한다. 핵심 아이디어는 간단하다.
latent를 연속 가우시안이 아니라 이산(discrete) 심볼로 만든다.
그 이산화는 “가장 가까운 코드북 벡터로 치환”하는 Vector Quantization(VQ) 방식으로 한다.
이후 디코더는 그 이산 심볼(코드 인덱스)에 대응하는 코드북 임베딩을 입력으로 받아 복원한다.
즉, “이미지/오디오를 일종의 토큰 시퀀스로 압축한 뒤, 그 토큰을 기반으로 복원/생성하는 2단계 구조”가 자연스럽게 나온다.
2. 전체 구조 요약: Encoder – Codebook – Quantization – Decoder
VQ-VAE는 구조 자체는 오토인코더와 유사하되, 가운데 병목(bottleneck)이 특별하다.
Encoder: 입력 x를 받아 연속 벡터 ze(x)를 만든다.
Codebook: K개의 코드 벡터 e1,…,eK를 가진다(embedding table).
Quantization: ze(x)와 가장 가까운 코드 ek를 찾아 그 코드로 스냅(snap) 한다.
Decoder: 양자화된 벡터 zq(x)=ek로부터 x를 복원한다.
여기서 “latent가 이산”이라는 말은, 실제로는 인덱스 k 가 이산이라는 뜻이고, 디코더는 그 인덱스에 해당하는 코드북 벡터 ek를 입력으로 받는다고 이해하면 된다.
이 구조가 갖는 직관은 명확하다. 인코더는 연속 공간에서 적당한 위치를 찍고, 그 위치를 가장 가까운 코드북 센터로 강제로 매핑한다. 이 과정에서 “연속값의 미세한 흔들림”이 제거되면서 표현이 안정적으로 정리되는 효과도 있다.
3. 문제점: argmin은 미분이 안 된다 (ST의 등장배경)
VQ에서 핵심 연산은 “가장 가까운 코드 선택”이다. 보통 다음처럼 쓴다.
문제는 argmin 미분 불가능하다는 점이다. 여기서 VQ-VAE가 선택한 트릭이 Straight-Through Estimator(ST) 이다.
forward에서는 진짜로 hard quantization을 한다.
backward에서는 “디코더 입력 zq에 대한 그래디언트”를 그대로 인코더 출력 ze에 흘려보낸다.
즉, 앞으로는 끊어지는데 뒤로는 그냥 통과시키는 셈이다. 엄밀한 미분은 아니지만, 학습이 잘 굴러가게 만드는 실용적 선택이라고 보면 된다.
4. VQ-VAE의 핵심 수식: loss(식 3)를 항별로 외우는 게 제일 중요하다
VQ-VAE를 공부할 때 가장 중요한 부분은, 목적함수(Eq. 3)의 각 항이 누구를 업데이트하는지를 분리해서 이해하는 것이다. 논문은 다음 형태의 loss를 둔다.
여기서 sg[⋅]는 stop-gradient로, forward는 통과하지만 backward에서 gradient를 0으로 만든다.
4.1 재구성 항 (reconstruction)
디코더가 복원을 잘 하도록 만드는 항이다. ST 덕분에 인코더도 이 항의 영향을 받는다.
4.2 코드북 업데이트 항 (VQ loss)
여기서는 sg[ze]때문에 인코더는 업데이트되지 않는다. 오직 코드북 ek만 ze 쪽으로 이동한다. 직관적으로는 “각 데이터 포인트가 선택한 코드북 벡터를 자기 쪽으로 끌어온다”는 느낌이며, k-means에서 centroid를 업데이트하는 것과 유사하다.
4.3 커밋먼트 항 (commitment loss)
이번에는 코드북에 stop-gradient가 걸려 코드북은 고정되고, 인코더 출력 ze만 코드북 벡터 근처에 머물도록 압박한다. 이 항이 없으면 인코더 출력이 불안정하게 커지거나(codebook이 따라가기 어려움) 코드 선택이 흔들릴 수 있어, “내가 선택한 코드에 붙어서 쓰겠다”는 커밋을 강제한다고 이해하면 된다.
정리하면 다음 문장으로 이해하면 된다.
재구성: “복원을 잘하자(디코더 중심, 인코더도 간접 영향)”
VQ loss: “코드북이 데이터(인코더 출력) 쪽으로 오게 하자”
commitment: “데이터(인코더 출력)가 코드북에 달라붙게 하자”
5. “VAE인데 KL은 어디?”: 학습 단계에서는 사실상 상수 취급
일반 VAE는 ELBO에서 KL 항이 중요하지만, VQ-VAE는 학습에서 prior를 일단 uniform categorical로 두고, posterior q(z∣x)는 결정론적(one-hot) 선택이므로 KL이 파라미터에 크게 관여하지 않는 형태로 정리된다(논문 흐름상 “상수처럼 처리”되는 느낌이다).
대신 VQ-VAE는 다음 단계에서 prior를 별도로 학습한다. 이 부분이 중요하다.
6. “2-stage 학습”이 VQ-VAE의 운영 방식이다
VQ-VAE는 보통 다음 순서로 쓴다.
VQ-VAE(Encoder+Codebook+Decoder) 를 학습해, 입력을 이산 토큰(코드 인덱스 맵) 으로 잘 압축하도록 만든다.
그 다음, 얻어진 토큰들 z에 대해 p(z)를 autoregressive prior(PixelCNN/WaveNet 등)로 따로 학습한다.
생성할 때는
prior에서 토큰을 샘플링한 뒤
decoder로 실제 이미지/오디오를 복원한다.
이 방식의 장점은 분명하다. 픽셀 공간에서 바로 autoregressive를 돌리면 차원이 너무 커서 비효율적이지만, 토큰 공간은 훨씬 작고 구조적이어서 “전역 패턴”을 모델링하기 좋다.
7. VQ-VAE가 실무/연구에서 중요했던 이유
개인적으로 VQ-VAE의 의의는 다음 두 가지로 정리할 수 있다.
이산 토큰화를 통해 강한 디코더 환경에서도 표현 학습이 잘 유지되는 기반을 만들었다.
“이미지/오디오를 토큰으로 바꿔서 다루는” 관점이 이후 모델들(VQGAN, dVAE, Jukebox, DALLE 계열, 토큰 기반 이미지 생성)에 넓게 이어졌다.
즉, VQ-VAE는 단순히 한 편의 모델이라기보다 “연속 신호를 토큰화해 다루는 표준 파이프라인”을 확립한 논문에 가깝다.
+ (RQ-VAE 곁다리)
VQ-VAE를 이해한 뒤 RQ-VAE를 보면, 핵심 차이는 한 문장으로 정리할 수 있따.
VQ-VAE는 “코드 1개”로 양자화하지만, RQ-VAE는 “잔차를 여러 번” 양자화해 코드 여러 개의 합으로 표현력을 키운다.
1) Residual Quantization(RQ)의 직관
어떤 벡터 z를 코드 하나로 근사하면 오차가 남는다. RQ는 그 오차(잔차)를 다시 양자화한다.
1단계: z≈e(k1)
잔차: r1=z−e(k1)
2단계: r1≈e(k2)
… 반복
최종: z^=∑i=1De(ki)
이렇게 하면 같은 코드북 크기 K에서도 조합 가능한 표현 수가 대략 K^D로 늘어나는 효과가 있다. 즉, “K를 키우지 않고도 표현력을 키우는 방식”이다.
2) RQ-VAE가 노리는 실용적 목표
RQ-VAE 논문은 특히 autoregressive 이미지 생성에서 “시퀀스 길이” 문제를 강조한다.
토큰 맵 해상도가 클수록(예: 32×32) AR 길이가 커져 비용이 급증한다.
길이를 줄이려면(예: 16×16로 다운샘플) 압축이 강해져 복원이 나빠지는데,
RQ는 “깊이 D”로 표현력을 보완해서 더 작은 토큰 맵을 쓰면서도 품질을 지키려는 전략이다.
3) RQ-Transformer까지 이어지는 설계
RQ-VAE가 토큰을 “위치마다 D개(스택)”로 내놓으니, Stage2에서는 그 스택을 예측해야 한다. 논문은 포지션 축(T)과 깊이 축(D)을 분리한 구조로 AR 비용을 줄이는 설계를 제안한다(공간 먼저, 깊이 다음).
정리: VQ-VAE의 핵심
정리하면, VQ-VAE는 다음 세 가지의 이해가 중요한 듯 하다.
코드북 + 최근접 양자화라는 bottleneck이 무엇인지
미분 불가능 문제를 ST + stop-gradient로 우회하는 방식
Eq.3의 세 항이 각각 decoder / codebook / encoder를 어떻게 업데이트하는지
그리고 RQ-VAE는 VQ-VAE를 확장하여, “코드 1개” 대신 “코드 여러 개의 합(잔차 양자화)”으로 표현력을 올려 더 강한 압축(짧은 시퀀스) 을 가능하게 만든 접근으로 이해할 수 있다.
+RQ-VAE논문은 한국 연구진분들이 출판하신 논문이다. 새삼 대단함을 다시 느낀다. 구글이 내 논문을 인용한다면 어떤 기분일까?
