VQ-VAE

Vector Quantised-Variational AutoEncoder, VQ-VAE; 벡터 양자화 변분 오토인코더
인코더 출력을 학습되는 코드북의 가장 가까운 벡터로 바꿔, 이미지·음성을 이산 토큰열로 표현하는 오토인코더

VQ-VAE(Vector Quantised-Variational AutoEncoder, 벡터 양자화 변분 오토인코더)는 오토인코더의 잠재 표현을 연속 벡터가 아니라 유한한 코드북(codebook) 항목의 번호로 만드는 모델이다. 2017년 딥마인드의 판덴오르트(Aäron van den Oord), 비냘스(Oriol Vinyals), 카북추오글루(Koray Kavukcuoglu)가 발표했다.^[1] 이미지는 32×32개의 정수 격자, 음성은 정수열이 되므로 그 위에 자기 회귀 모델이나 트랜스포머를 얹어 언어 모델처럼 생성할 수 있다.

이름에 VAE가 들어가지만 변분 오토인코더처럼 잠재 변수를 표본 추출하지 않는다. 코드 선택은 결정적이고, 사전 분포는 학습 중에는 균등 분포로 두었다가 나중에 따로 학습한다. VQ-VAE 자체보다 "이미지·소리를 토큰으로 바꾸는 토크나이저"라는 쓰임이 더 중요해져, DALL-E(1세대)의 이산 VAE, VQGAN, 신경 오디오 코덱 SoundStream·EnCodec, 영상 토크나이저 MAGVIT-v2로 이어졌다.

구조

(1) 인코더 $E$ 가 입력 $x$ 를 $H \times W \times D$ 특성맵 $z_{e} (x)$ 로 줄인다.
(2) 코드북 ${e_{1}, \dots, e_{K}} \subset ℝ^{D}$ 에서 위치마다 가장 가까운 벡터를 고른다. $k = \arg \min_{j} ‖ z_{e} (x) - e_{j} ‖_{2}$ , $z_{q} (x) = e_{k}$ .
(3) 디코더 $G$ 가 $z_{q} (x)$ 로 $x$ 를 복원한다.
(4) 학습이 끝나면 코드 번호 격자 위에 사전 분포 모델(이미지는 PixelCNN, 음성은 WaveNet)을 따로 학습하고, 생성할 때는 이 모델로 번호를 뽑아 디코더에 넣는다.

원 논문은 ImageNet 128×128×3 이미지를 $K = 512$ 짜리 코드 32×32개로 표현했다. 8비트 픽셀 기준으로 비트 수가 약 42.6분의 1로 줄어드는데도 복원 이미지는 원본보다 약간 흐린 정도였다.^[1]

학습

손실 함수

$ℒ = \underset{복원}{\underset{⏟}{- \log p (x ∣ z_{q} (x))}} + \underset{코드북}{\underset{⏟}{‖ sg [z_{e} (x)] - e ‖_{2}^{2}}} + β \underset{약속(commitment)}{\underset{⏟}{‖ z_{e} (x) - sg [e] ‖_{2}^{2}}}$

$sg$ 는 정지 기울기(stop-gradient) 연산자로, 순전파에서는 항등이고 기울기는 0이다.

복원 항은 디코더와 (아래 직통 추정기를 거쳐) 인코더를 학습시킨다.
코드북 항은 고른 코드 벡터를 인코더 출력 쪽으로 끌어당긴다. K-평균 군집화의 중심 갱신과 같은 역할이다.
약속 항은 인코더 출력이 코드에서 멀리 떠다니며 커지지 않도록 붙잡는다. 논문은 β를 0.1~2.0으로 바꿔도 결과가 거의 같았고 실험에는 0.25를 썼다.^[1]

사전 분포가 균등하다고 가정하므로 ELBO의 KL 항은 상수( $\log K$ )가 되어 학습에서 빠진다. 디코더가 강력한 자기회귀 모델이면 잠재 변수를 무시해 버리는 VAE의 사후 분포 붕괴(posterior collapse)도 이산 병목 덕분에 피할 수 있다고 논문은 설명했다.

직통 추정기

$\arg \min$ 은 미분할 수 없다. VQ-VAE는 역전파 때 디코더 입력 $z_{q}$ 에 대한 기울기를 인코더 출력 $z_{e}$ 로 그대로 복사한다. 직통 추정기(straight-through estimator, STE)라 부르는 방법이다. 인코더 출력과 디코더 입력이 같은 $D$ 차원 공간에 있으므로, 복사된 기울기는 "복원 손실을 줄이려면 인코더 출력을 어느 쪽으로 옮겨야 하는가"를 알려 주고, 출력이 옮겨 가면 다음 순전파에서 다른 코드가 선택된다. 코드에서는 다음 한 줄로 구현한다.

$z_{q} \leftarrow z_{e} + sg [z_{q} - z_{e}]$

EMA 코드북 갱신

코드북 항 대신 지수 이동 평균(EMA)으로 코드 벡터를 갱신할 수 있다. 코드 $e_{i}$ 에 배정된 인코더 출력의 개수 $N_{i}$ 와 합 $m_{i}$ 를 감쇠율 γ로 누적하고 $e_{i} = m_{i} / N_{i}$ 로 둔다. 미니배치용 온라인 K-평균인 셈이다. 원 논문 부록은 γ = 0.99가 잘 작동했다고 적었고,^[1] 신경 오디오 코덱 SoundStream과 EnCodec도 감쇠율 0.99의 EMA 갱신을 썼다. EnCodec은 쓰이지 않는 코드를 현재 배치의 표본으로 갈아 끼우는 재초기화도 함께 썼다.^[2]

코드북 붕괴

가장 흔한 실패는 일부 코드만 계속 선택되고 나머지는 한 번도 쓰이지 않는 코드북 붕괴(codebook collapse, 죽은 코드)다. 유효 어휘가 줄어 복원 품질과 생성 다양성이 모두 떨어진다. 대응책은 다음과 같다.

EMA 갱신과 죽은 코드 재초기화(배치의 인코더 출력이나 K-평균 결과로 다시 채움)
인수분해 코드와 L2 정규화: ViT-VQGAN은 코드 검색을 768차원 대신 32차원이나 8차원의 저차원 공간에서 하고, 인코더 출력과 코드를 L2 정규화해 코사인 유사도로 찾게 해서 코드북을 8192개로 키워도 사용률을 높였다.^[3]
코드북 자체를 없애기: FSQ와 LFQ(아래)

구현 예

양자화층만 떼어 낸 최소 구현이다. 코드북 손실과 약속 손실, 직통 추정기를 그대로 옮겼다. PyTorch 2.5에서 실행해 출력 기울기가 인코더 쪽으로 그대로(모두 1) 전달되는 것을 확인했다.

import torch, torch.nn as nn, torch.nn.functional as F

class VectorQuantizer(nn.Module):
    def __init__(self, num_codes=512, dim=64, beta=0.25):
        super().__init__()
        self.codebook = nn.Embedding(num_codes, dim)
        self.codebook.weight.data.uniform_(-1 / num_codes, 1 / num_codes)
        self.beta = beta

    def forward(self, z_e):                          # z_e: (B, D, H, W) 인코더 출력
        B, D, H, W = z_e.shape
        flat = z_e.permute(0, 2, 3, 1).reshape(-1, D)
        dist = (flat.pow(2).sum(1, keepdim=True)    # ||z||^2 - 2 z.e + ||e||^2
                - 2 * flat @ self.codebook.weight.t()
                + self.codebook.weight.pow(2).sum(1))
        idx = dist.argmin(1)                         # 가장 가까운 코드 번호
        z_q = self.codebook(idx).view(B, H, W, D).permute(0, 3, 1, 2)
        codebook_loss = F.mse_loss(z_q, z_e.detach())       # ||sg[z_e] - e||^2
        commit_loss = F.mse_loss(z_e, z_q.detach())         # ||z_e - sg[e]||^2
        z_q = z_e + (z_q - z_e).detach()             # 직통 추정기
        return z_q, codebook_loss + self.beta * commit_loss, idx.view(B, H, W)

* (5) 학습 루프: loss = recon_loss(decoder(z_q), x) + vq_loss

idx가 토큰이다. 2단계에서는 이 정수 격자를 래스터 순서로 펴서 자기회귀 트랜스포머를 학습하거나, MaskGIT처럼 가린 토큰을 한꺼번에 맞히는 양방향 트랜스포머를 학습한다.

발전

모델	연도	바뀐 점	쓰임
VQ-VAE	2017	코드북, 직통 추정기, 약속 손실	이미지·영상·음성 생성. 음성 잠재 코드를 41개 음소에 대응시키면 49.3% 정확도(무작위 7.2%)로 비지도 음소 학습 가능
VQ-VAE-2^[4]	2019	상·하위 두 단계 계층 코드, 강한 자기회귀 사전 분포	ImageNet에서 당시 GAN에 견줄 품질, 모드 붕괴 없음
DALL-E 이산 VAE^[5]	2021	직통 추정기 대신 검벨-소프트맥스 완화, 어휘 8192, 8배 축소	256×256 이미지를 32×32 토큰으로 만들어 텍스트 토큰과 한 줄로 이어 자기회귀 생성
VQGAN^[6]	2020	지각 손실과 패치 판별자 적대 손실을 더해 선명한 복원, 트랜스포머 사전 분포	메가픽셀 이미지 합성, 잠재 확산 모델의 전신
ViT-VQGAN	2021	ViT 인코더·디코더, 인수분해·L2 정규화 코드	ImageNet 256×256 FID 4.17(기존 VQGAN 17.04)^[3]
SoundStream RVQ^[7]	2021	잔차 벡터 양자화, 양자화기 드롭아웃	3~18kbps 가변 비트율 신경 오디오 코덱
FSQ^[8]	2023	학습 코드북 없이 차원마다 정해진 몇 단계로 반올림	보조 손실 없이 코드 전부 사용
MAGVIT-v2 LFQ^[9]	2023	코드 임베딩 차원을 0으로 줄인 검색 없는 양자화	큰 어휘의 이미지·영상 공용 토크나이저

잔차 벡터 양자화

잔차 벡터 양자화(residual vector quantization, RVQ)는 첫 코드북으로 양자화하고 남은 잔차를 둘째 코드북으로, 그 잔차를 셋째 코드북으로 양자화하는 식으로 $N_{q}$ 단을 쌓는다. 양자화된 벡터는 각 단 코드의 합이다. 코드북 하나로 같은 정밀도를 내려면 크기가 지수적으로 커져야 하는 문제를 피한다. SoundStream은 학습 때 예제마다 단 수 $n_{q}$ 를 1과 $N_{q}$ 사이에서 무작위로 골라 일부만 쓰는 양자화기 드롭아웃으로, 한 모델이 3kbps부터 18kbps까지 여러 비트율을 지원하게 했다. 24kHz 음성에서 3kbps SoundStream이 12kbps Opus보다 좋은 주관 평가를 받았다.^[7] 메타의 EnCodec도 RVQ를 써서 24kHz에서 1.5, 3, 6, 12, 24kbps를 지원하고, 코드북당 1024개 항목(10비트)을 최대 32개 썼다.^[2]

이렇게 만든 오디오 토큰은 오디오 언어 모델의 입력이 된다. 구글의 AudioLM은 신경 오디오 코덱의 이산 코드를 음질용 토큰으로 써서 오디오 생성을 언어 모델링 문제로 바꿨고,^[10] 메타의 MusicGen은 여러 줄의 압축 토큰 열을 한 트랜스포머로 생성한다.^[11] → 음성 합성

코드북 없는 양자화

FSQ(finite scalar quantization): 인코더 출력을 10차원 미만으로 줄이고 각 차원을 $⌊ L / 2 ⌋ \tanh (z)$ 같은 함수로 묶은 뒤 정수로 반올림한다. 차원마다 단계 수 $L_{i}$ 를 정하면 암묵적 코드북 크기는 $\prod_{i} L_{i}$ 다. 기울기는 똑같이 직통 추정기로 흘린다. 논문은 약속 손실, 코드북 재초기화, 엔트로피 벌점 같은 장치 없이도 코드를 모두 쓰고 MaskGIT 이미지 생성 등에서 VQ와 비슷한 성능을 냈다고 보고했다.^[8]
LFQ(lookup-free quantization): 코드 임베딩 차원을 0으로 줄여 각 차원의 부호만 남기는 방식이다. 거리 계산으로 코드를 찾을 필요가 없어 어휘를 크게 키울 수 있고, MAGVIT-v2는 이 토크나이저로 언어 모델이 ImageNet과 Kinetics 생성에서 확산 모델을 앞섰다고 주장했다.^[9]

이미지 토크나이저로서의 쓰임

VQ-VAE 계열은 "2단계 생성"의 1단계를 맡는다. 1단계 토크나이저가 지각적으로 중요하지 않은 세부를 압축하고, 2단계 모델은 짧아진 토큰열에서 전체 구조만 배운다.

자기회귀 생성: DALL-E(1세대), VQGAN + 트랜스포머처럼 토큰을 한 줄로 펴서 다음 토큰을 예측한다.
마스크 생성: MaskGIT은 양방향 트랜스포머가 가린 토큰을 여러 개씩 동시에 채워, 래스터 순서 자기회귀 디코딩보다 최대 64배 빨랐다고 보고했다.^[12]
잠재 확산: 잠재 확산 모델 논문은 VQ 정규화 오토인코더(VQGAN의 양자화층을 디코더 쪽으로 흡수한 형태)와 KL 정규화 오토인코더를 모두 실험했다.^[13] 확산 모델은 양자화 전의 연속 잠재에서 학습한다. → 확산 모델
멀티모달 언어 모델: 이미지를 이산 토큰으로 바꾸면 텍스트와 같은 어휘 공간에서 하나의 자기회귀 모델로 다룰 수 있다. 이것이 이미지·음성을 "말하는" 멀티모달 모델에서 토크나이저 연구가 중요한 이유다. → 멀티모달 인공지능

연속 잠재(VAE)와 비교

항목	VQ-VAE 계열(이산)	KL 정규화 VAE(연속)
잠재 표현	정수 토큰 격자	실수 벡터 격자
2단계 모델	자기회귀·마스크 트랜스포머(교차 엔트로피)	확산·플로 매칭(회귀)
복원 품질	어휘 크기와 단 수로 비트율이 정해짐, 세부 손실 가능	같은 압축률에서 대체로 유리
학습 난점	코드북 붕괴, 직통 추정기의 편향	KL 가중치 조절, 사후 분포 붕괴
장점	언어 모델 기술을 그대로 적용, 압축·전송에 바로 쓰는 비트열	연속 생성 모델과 궁합이 좋음

같이 보기

각주

↑ ^1.0 ^1.1 ^1.2 ^1.3 Aaron van den Oord, Oriol Vinyals, Koray Kavukcuoglu, 〈Neural Discrete Representation Learning〉, NeurIPS 2017. arXiv:1711.00937
↑ ^2.0 ^2.1 Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi, 〈High Fidelity Neural Audio Compression〉, 2022. arXiv:2210.13438
↑ ^3.0 ^3.1 Jiahui Yu 외, 〈Vector-quantized Image Modeling with Improved VQGAN〉, ICLR 2022. arXiv:2110.04627
↑ Ali Razavi, Aaron van den Oord, Oriol Vinyals, 〈Generating Diverse High-Fidelity Images with VQ-VAE-2〉, NeurIPS 2019. arXiv:1906.00446
↑ Aditya Ramesh 외, 〈Zero-Shot Text-to-Image Generation〉, ICML 2021. arXiv:2102.12092
↑ Patrick Esser, Robin Rombach, Björn Ommer, 〈Taming Transformers for High-Resolution Image Synthesis〉, CVPR 2021. arXiv:2012.09841
↑ ^7.0 ^7.1 Neil Zeghidour 외, 〈SoundStream: An End-to-End Neural Audio Codec〉, 2021. arXiv:2107.03312
↑ ^8.0 ^8.1 Fabian Mentzer 외, 〈Finite Scalar Quantization: VQ-VAE Made Simple〉, 2023. arXiv:2309.15505
↑ ^9.0 ^9.1 Lijun Yu 외, 〈Language Model Beats Diffusion: Tokenizer is Key to Visual Generation〉, ICLR 2024. arXiv:2310.05737
↑ Zalán Borsos 외, 〈AudioLM: a Language Modeling Approach to Audio Generation〉, 2022. arXiv:2209.03143
↑ Jade Copet 외, 〈Simple and Controllable Music Generation〉, NeurIPS 2023. arXiv:2306.05284
↑ Huiwen Chang 외, 〈MaskGIT: Masked Generative Image Transformer〉, CVPR 2022. arXiv:2202.04200
↑ Robin Rombach 외, 〈High-Resolution Image Synthesis with Latent Diffusion Models〉, CVPR 2022. arXiv:2112.10752

[vqvae-1] 1.0 ^1.1 ^1.2 ^1.3 Aaron van den Oord, Oriol Vinyals, Koray Kavukcuoglu, 〈Neural Discrete Representation Learning〉, NeurIPS 2017. arXiv:1711.00937

[encodec-2] 2.0 ^2.1 Alexandre Défossez, Jade Copet, Gabriel Synnaeve, Yossi Adi, 〈High Fidelity Neural Audio Compression〉, 2022. arXiv:2210.13438

[vitvqgan-3] 3.0 ^3.1 Jiahui Yu 외, 〈Vector-quantized Image Modeling with Improved VQGAN〉, ICLR 2022. arXiv:2110.04627

[vqvae2-4] Ali Razavi, Aaron van den Oord, Oriol Vinyals, 〈Generating Diverse High-Fidelity Images with VQ-VAE-2〉, NeurIPS 2019. arXiv:1906.00446

[dalle-5] Aditya Ramesh 외, 〈Zero-Shot Text-to-Image Generation〉, ICML 2021. arXiv:2102.12092

[vqgan-6] Patrick Esser, Robin Rombach, Björn Ommer, 〈Taming Transformers for High-Resolution Image Synthesis〉, CVPR 2021. arXiv:2012.09841

[soundstream-7] 7.0 ^7.1 Neil Zeghidour 외, 〈SoundStream: An End-to-End Neural Audio Codec〉, 2021. arXiv:2107.03312

[fsq-8] 8.0 ^8.1 Fabian Mentzer 외, 〈Finite Scalar Quantization: VQ-VAE Made Simple〉, 2023. arXiv:2309.15505

[magvit2-9] 9.0 ^9.1 Lijun Yu 외, 〈Language Model Beats Diffusion: Tokenizer is Key to Visual Generation〉, ICLR 2024. arXiv:2310.05737

[audiolm-10] Zalán Borsos 외, 〈AudioLM: a Language Modeling Approach to Audio Generation〉, 2022. arXiv:2209.03143

[musicgen-11] Jade Copet 외, 〈Simple and Controllable Music Generation〉, NeurIPS 2023. arXiv:2306.05284

[maskgit-12] Huiwen Chang 외, 〈MaskGIT: Masked Generative Image Transformer〉, CVPR 2022. arXiv:2202.04200

[ldm-13] Robin Rombach 외, 〈High-Resolution Image Synthesis with Latent Diffusion Models〉, CVPR 2022. arXiv:2112.10752

[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]

[9]

[10]

[11]

[12]

[13]