채널별 벡터 양자화(CVQ)

왜 중요한가

전통적인 patch-wise 벡터 양자화는 코드북 활용 저하와 AR에 불리한 2D 토큰 순서를 야기한다. CVQ는 채널 차원에서 벡터 양자화를 수행해 1D 토큰 시퀀스를 형성하고 100% 코드북 활용을 달성하며, 채널 단위의 AR인 CAR로 이미지 생성을 글로벌 구조에서 세부 텍스처까지 점진적으로 향상시킨다. 이를 통해 더 큰 코드북의 효과를 활용하고 AR 학습의 순서를 보다 자연스럽게 정렬할 수 있다.

핵심 기여

CVQ: 채널별 벡터 양자화의 제안

이미지를 2D patch 대신 h×w×c 피처의 각 채널(z(:,:,k))에 대해 양자화를 적용하는 CVQ를 제안한다. 이로써 1D 채널 시퀀스(X)를 구성하고, 100% 코드북 활용과 재구성 품질 향상을 달성한다.

CAR: next-channel prediction

채널 단위 토큰으로 구성된 1D 시퀀스에 대해 decoder-only 트랜스포머를 이용한 AR 학습을 수행한다. p(X) = ∏k p(x^(k) | x^(<k)) 형태의 확률 분포를 학습한다.

Nested channel dropout으로 AR-friendly ordering 확보

채널의 초기 부분만 활성화되는 c_keep 구성과 확률적 샘플링으로 coarse-to-fine 채널 순서를 학습한다. 이를 통해 patch 기반의 어려움 없이 AR 학습의 안정성과 성능을 향상시킨다.

원-패스 재구성 및 AR 생성의 성능 향상

256 토큰에서 rFID 2.60, 1024 토큰에서 rFID 0.88 등 재구성 품질이 개선되며, 텍스트-이미지 생성에서도 CAR가 기존 AR baselines를 상회하는 성능을 보인다.

다중 해상 확장 및 1D 토크나이저의 경쟁력

CVQ는 1D 토크나이저 중에서도 비교적 간단한 구성으로도 높은 코드북 활용과 재구성 품질을 달성하며, 65K 코드북까지 안정적으로 활용 가능성을 확인한다.

핵심 아이디어 이해하기

단계1: patch-wise VQ는 패치 간 중복으로 코드북이 자주 업데이트되지 않아 코드북 collapse가 발생하기 쉽다. 단계2: CVQ는 Z를 채널 단위로 분해하고 z(:,:,k)를 각각 코드북에 매칭하는 channel-wise quantization으로 바꿔 각 채널에 고유한 정보를 분리시키고 코드북 활용을 높인다. 단계3: 채널(tokens)들을 순차적으로 예측하는 CAR는 1D 토큰 시퀀스로 AR 학습이 가능하게 하며, nested dropout으로 coarse-to-fine 채널 순서를 학습해 AR 친화성을 높인다. 단계4: 재구성 품질은 16K 이상의 코드북에서도 100% 활용과 AR 시퀀스의 효율성으로 개선되며, 텍스트-이미지 생성에서도 강한 AR 성능을 보인다.

방법론

Z ∈ R^{h×w×c}에서 Z의 각 채널 z(:,:,k) ∈ R^{h×w×1}에 대해 코드북 C_channel = {e^(n)}를 두고 양자화를 수행한다. 각 k에 대해 z^(k)q = argmin_en∈Cchannel || z^(k) − en ||^2 이고, z^(k)q는 z^(k) + sg[en − z^(k)]로 갱신한다. 역전파는 STE를 사용한다. 채널은 순서를 갖지 않으므로 tokenizer 학습 시 nested dropout으로 상위 채널이 먼저 채워지도록 coarse-to-fine 순서를 만든다. X = {x_i}{i=1}^c의 1D 토큰 시퀀스로 간주되며 p(X) = ∏{k=1}^c p(x^(k) | x^(<k))의 AR likelihood를 decoder-only transformer로 학습한다. 입력 채널들을 h×w 차원으로 feed하기 위해 2-layer MLP 프로젝터를 사용해 차원을 정렬하고 LLM 백본과 연동한다. CVQ tokenizer 학습은 재구성loss, commitment loss, LPIPS loss, PatchGAN discriminator를 포함한 표준 VQGAN 스타일의 손실로 수행한다. 또, Stage I에서는 MLP 프로젝터와 LLM head만 학습하며 Stage II에서 전체 파라미터를 미세조정한다.

주요 결과

주요 재구성 벤치마크에서 256 tokens의 CVQ는 rFID 2.60, SSIM 0.565, PSNR 20.94를 달성하며, 1024 tokens의 CVQ는 rFID 0.88, SSIM 0.723, PSNR 25.02를 달성한다(코드북 사용률 100%). patch-based VQ 대비 100% 코드북 활용과 재구성 품질 향상을 보인다. 텍스트-이미지 생성에서 CAR(4B)는 GenEval 0.88, CAR(8B) 0.92로 상위 AR baseline인 NextStep-1/Emu3를 경쟁하거나 우수한 성능을 보인다. MJHQ-30K에서 CVQ의 FID은 6.42로 1D 토큰 기반의 강력한 성능을 기록한다. nested dropout 도입 시 GenEval +0.12, DPG +9.38의 향상을 보이며 재구성과 생성 간의 균형이 개선된다.

기술 상세

CVQ는 기존 VQ의 patch-wise 양자화를 채널(channel) 축으로 바꾼 아이디어이다. Z ∈ R^{h×w×c}에서 Z의 각 채널 z(:,:,k)를 코드북 C_channel의 엔트리로 양자화한다. z^(k)q = argmin{e∈Cchannel} || z^(k) − e ||^2 이며, forward는 z^(k)q = z^(k) + sg[e_n − z^(k)]. 역전파는 STE를 사용한다. 채널별 코드북은 서로 다른 코드워드 벡터를 가지며, 채널 간의 정보가 서로 다른 코드북으로 매핑되므로 코드북의 재활용이 감소한다. CVQ의 토큰 시퀀스 X = {x_i}{i=1}^c에서 각 x_i는 채널의 코드북 인덱스를 나타낸다. CAR은 이 X에 대해 p(X) = ∏_{k=1}^c p(x^(k) | x^(<k))를 학습한다. 입력 채널은 h×w 차원으로 변환하기 위해 두-layer MLP를 사용해 LLM 백본에 맞춘 차원을 맞춘 뒤 transformer에 전치한다. 채널 간 순서가 없으므로 nested dropout을 통해 coarse-to-fine 순서를 학습한다. CVQ tokenizer의 학습 손실은 recon/L1, commitment, LPIPS, GAN 손실로 구성되며 PatchGAN discriminator를 사용한다. Nested dropout의 비율 α는 0~50% 범위에서 실험되며, α가 커질수록 AR 생성 성능이 향상되지만 재구성 품질은 크게 손실되지 않는다. 실험 데이터는 ImageNet-1K, LAION- aesth., CC12M 등 다수의 텍스트-이미지 데이터셋으로 Stage I과 Stage II에서 학습이 수행된다.

실무 활용

CVQ는 토큰 수를 고정한 채 해상도 확장을 가능하게 하며, 더 큰 코드북으로의 확장을 통해 재구성 및 생성 성능을 개선할 수 있다. 기존 patch 기반 VQ 대비 AR 학습을 보다 자연스럽게 만들고, unified 비전-언어 모델 등에 쉽게 적용 가능하다.

고해상도 텍스트-투-이미지 생성 파이프라인에 CVQ 토크나이저를 도입
LLM-연계 multimodal 모델에 CVQ를 결합해 학습 효율성 및 생성 품질 향상
동영상 확장에 채널 차원을 시간 축으로 확장하는 CVQ-Temporal 확장 연구의 기초
저용량 토큰으로도 큰 해상도 표현이 필요한 모바일 비주얼 생성 애플리케이션

코드 공개 여부: 공개

코드 저장소 보기

키워드

Channel-wise Vector Quantizationimage tokenizationvector quantizationpatch-wise tokenschannel-wise tokensvisual autoregressive frameworknext-channel predictioncodebook utilization

용어 해설

코드북 활용도(Codebook Utilization): — 코드북의 벡터가 학습 과정에서 골고루 업데이트되어 전체 코드북이 활발히 사용되는 정도를 뜻한다. CVQ는 채널 단위 양자화를 통해 코드북 활용도를 거의 100%에 가깝게 유지시키는 경향이 있다.
네스티드 드롭아웃(Nested Dropout): — 다음에 올 토큰의 순서를 coarse-to-fine으로 정렬하도록 내부 레이어를 점진적으로 비활성화하는 학습 기법. CVQ에서 AR 학습의 채널 순서를 형성하는 데 사용된다.
다음 채널 예측(Next-Channel Prediction): — 이미지 생성을 채널 단위로 순차 예측하는 AR 방식. patch-by-patch 대신 채널 간의 의존성을 순차적으로 모델링한다.
채널별 양자화(Channel-wise Quantization): — Z를 채널 단위로 분해하여 각 채널(z(:,:,k))을 코드북에 양자화하는 새로운 양자화 축. 2D patch 대신 1D 채널 시퀀스를 생성한다.
VQ-VAE: — 연속 표현을 이산 코드로 매핑하는 초기 연구 흐름. CVQ의 아이디어는 이를 채널 축으로 확장하는 것에 있다.
자귀적 예측(Autoregressive): — 다음 토큰을 이전 토큰들로 조건화해 순차적으로 예측하는 모델링 패러다임. CVQ는 이를 2D 패치가 아닌 채널 단위로 적용한다.