VQ-Diffusion: 이산 공간에서의 조건부 잠재 확산

VQ-Diffusion은 중국과학기술대학과 Microsoft가 공동 개발한 조건부 잠재 확산 모델입니다. 대부분의 일반적인 확산 모델이 연속 공간에서 작동하는 것과 달리, VQ-Diffusion은 이산적인 벡터 집합으로 구성된 양자화된 잠재 공간에서 노이징 및 디노이징 과정을 수행합니다.

아키텍처 및 기술 구현

VQ-Diffusion은 차원 축소를 위해 VQ-VAE를 사용하고, 확산 과정을 위해 인코더‑디코더 트랜스포머를 결합합니다.

VQ-VAE 잠재 공간

이미지는 먼저 VQ-VAE 인코더를 사용해 이산 토큰 또는 임베딩 벡터로 인코딩됩니다. 이 과정은 이미지를 패치로 나눈 뒤 각 패치를 고정 크기 어휘 코드북의 가장 가까운 항목으로 교체하는 방식으로 진행됩니다. VQ-Diffusion은 Taming Transformers의 VQGAN 변형을 사용하며, 확산 학습 과정 동안 고정된 사전 학습된 VQ-VAE를 활용합니다.

순방향 확산 과정

순방향 과정에서 잠재 토큰은 깨끗한 상태에서 노이즈가 섞인 상태로 전이합니다. 각 토큰은 다음 중 하나를 수행할 수 있습니다:

  1. 그대로 유지한다.
  2. 동일한 확률로 다른 잠재 벡터로 재샘플링된다.
  3. 마스킹된다.

토큰이 마스킹되면 이후에도 계속 마스킹된 상태를 유지합니다. 이 과정은 하이퍼파라미터 $\alpha_{t}$, $\beta_{t}$, $\gamma_{t}$에 의해 제어되며, $\gamma_{t}$는 토큰이 마스킹될 확률, $\alpha_{t} + \beta_{t}$는 토큰이 그대로 유지될 확률을 의미합니다. 마스킹되지 않은 개별 잠재 벡터로 전이될 확률은 $\beta_{t}$입니다.

역과정 근사

노이즈를 역전시켜 이미지를 생성하기 위해 인코더‑디코더 트랜스포머가 사용되어, 프롬프트 ($y$)에 조건화된 비노이즈 잠재 ($x_{0}$)의 클래스를 근사합니다.

  • Encoder: 가중치가 고정된 CLIP 텍스트 인코더.
  • Decoder: 모든 잠재 픽셀에 대해 마스크되지 않은 전역 어텐션을 제공하고, 벡터 임베딩에 대한 범주형 분포의 로그 확률을 출력하는 트랜스포머.

디코더가 한 번의 순전파로 비노이즈 잠재 전체 분포를 예측하기 때문에 현재 상태 $x_{t}$에 대한 전역 자기‑어텐션을 제공합니다.

VQ-Diffusion과 다른 생성 모델 비교

VQ-Diffusion은 연속 확산 모델과 자동 회귀(AR) 모델 모두에 대한 기술적 대안을 제시합니다.

연속 확산 모델과의 비교

대부분의 최신 확산 모델은 연속적이며, 가우시안 노이즈를 반복적으로 추가하고 U‑Net을 통해 그 노이즈를 예측함으로써 역과정을 근사합니다. 반면 VQ-Diffusion은 이산 값을 사용하며, $x_{0}$에 대한 분포를 직접 예측하는 것이 노이즈를 예측하는 것보다 더 명확한 목표가 됩니다.

자동 회귀(AR) 모델과의 비교

VQ-Diffusion은 AR 트랜스포머 기반 이미지 모델에서 흔히 발생하는 세 가지 주요 문제를 해결합니다: 해상도가 증가함에 따라 추론 속도가 선형적으로 감소하는 현상, 오류 누적, 그리고 방향성 편향.

  • 추론 속도: AR 모델은 이미지 확률을 픽셀별로 순차적으로 조건화합니다(래스터 스캔 순서). VQ-Diffusion의 추론 복잡도는 확산 단계 수가 잠재 픽셀 수보다 적을 때 우수합니다. ITHQ 데이터셋(32×32 잠재 해상도)에서는 100 단계까지 학습되어 약 10배의 빅‑O 개선을 보이며, 실제로는 AR 방식보다 최대 15배 빠르면서도 이미지 품질이 더 높습니다.
  • 오류 교정: AR 모델이 교사 강제(teacher‑forcing)를 필요로 하는 반면, VQ-Diffusion은 마스킹과 무작위 토큰 교체를 모두 포함한 학습을 통해 잘못 예측된 토큰을 스스로 교정하도록 학습됩니다.
  • 전역 컨텍스트: VQ-Diffusion은 $x_{t}$를 예측하면서 $x_{t-1}$을 전역적으로 고려하지만, AR 모델은 미래 픽셀을 보지 않도록 마스크된 어텐션에 제한됩니다.

구현 및 추가 최적화

VQ-Diffusion은 diffusers 라이브러리의 VQDiffusionPipeline을 통해 Hugging Face에 통합되었습니다. 모델은 ITHQ, CUB-200, Oxford-102, MSCOCO, Conceptual Captions, LAION-400M, ImageNet 등 다양한 데이터셋으로 학습되었습니다.

추가 최적화로는 역확산 단계를 건너뛰는 시간 스트라이드 $\Delta t$를 이용한 빠른 추론 전략이 포함됩니다. 또한 "Improved Vector Quantized Diffusion Models"에서는 샘플 품질을 향상시키는 이산 클래스‑프리 가이던스와 공동 분포 문제를 해결하기 위한 대체 추론 전략을 제시합니다.

Sources