변분 오토인코딩 이산 확산 (VADD): 차원 간 상관관계 모델링 강화

이 논문은 **변분 오토인코딩 이산 확산 (Variational Autoencoding Discrete Diffusion, VADD)**이라는 새로운 프레임워크를 제안합니다. 이는 기존 마스킹 확산 모델 (Masked Diffusion Models, MDM) 의 한계를 극복하고, 특히 적은 수의 디노이징 (denoising) 단계에서도 고품질 샘플을 생성할 수 있도록 차원 간 상관관계를 효과적으로 모델링하는 것을 목표로 합니다.

1. 문제 정의 (Problem)

최근 이산 상태 공간 (discrete state space) 을 다루는 확산 모델은 이미지, 텍스트 생성 등에서 autoregressive 모델과 경쟁력 있는 성능을 보여주고 있습니다. 특히 MDM 은 데이터를 점진적으로 마스킹 (forward) 하고, 역과정에서 여러 차원을 동시에 예측 (unmasking) 하는 병렬 생성 방식을 통해 autoregressive 모델보다 빠른 추론 속도를 제공합니다.

그러나 MDM 의 주요 한계점은 다음과 같습니다:

• 차원 간 독립성 가정: MDM 의 역과정 (backward process) 에서 각 차원의 분포는 서로 독립적인 카테고리 분포의 곱으로 모델링됩니다. 이는 실제 데이터에 존재하는 복잡한 차원 간 의존성 (inter-dimensional dependencies) 을 포착하지 못하게 합니다.
• 적은 단계에서의 성능 저하: 디노이징 단계 (sampling steps) 가 적을수록 한 번에 마스킹이 해제되는 차원의 수가 많아지는데, 이때 독립성 가정이 강하게 작용하여 샘플의 품질이 급격히 떨어집니다.
• 기존 해결책의 비효율성: 기존 연구들은 사전 훈련된 autoregressive 모델이나 상관관계 모델을 사용하여 내부 루프 (inner-loop) 샘플링을 수행함으로써 이를 보완하려 했으나, 이는 추가적인 계산 비용을 발생시킵니다.

2. 방법론 (Methodology)

저자들은 MDM 의 디노이징 분포에 잠재 변수 (latent variable) 구조를 도입하여 VADD를 제안했습니다.

2.1 잠재 변수 기반 디노이징 모델

기존 MDM 이 $p_\theta(x_s|x_t)$를 명시적인 분포로 정의하는 대신, VADD 는 이를 잠재 변수 $z$를 가진 잠재 변수 모델로 정의합니다:
$$p_\theta(x_s|x_t) = \int p_\theta(x_s|x_t, z)p(z)dz$$
여기서 $z$는 고수준의 의미 (high-level semantics) 를 제어하는 잠재 변수로 작용하여, 마스킹된 상태 $x_t$에서 깨끗한 데이터 $x_0$를 복원하는 다양한 모드 (multimodality) 를 포착하고 차원 간의 상관관계를 암시적으로 학습하도록 돕습니다.

2.2 변분 오토인코딩 (VAE) 메커니즘

잠재 변수 $z$의 주변 분포 (marginal distribution) 를 직접 계산하는 것은 불가능 (intractable) 하므로, VAE 프레임워크를 차용합니다:

• 인식 모델 (Recognition Model): $r_\phi(z|x_0, x_t)$를 도입하여 사후 분포 $p_\theta(z|x_0, x_t)$를 근사합니다.
• Double Evidence Lower Bound (DELBO): 기존 ELBO 를 확장하여, 디노이징 모델 ($\theta$) 과 인식 모델 ($\phi$) 을 공동으로 최적화하는 새로운 하한 (lower bound) 을 유도했습니다.
• KL Annealing: 사후 붕괴 (posterior collapse, 인식 모델이 데이터를 무시하고 사전 분포로 수렴하는 현상) 를 방지하기 위해 KL 발산 항에 가중치 $\lambda$를 동적으로 적용합니다.

2.3 텍스트 생성을 위한 아키텍처 설계

텍스트 생성에 Transformer 를 적용할 때의 효율성을 위해 다음과 같은 설계를 채택했습니다:

• 디노이징 모델: 잠재 변수 $z$의 임베딩을 AdaLN (Adaptive Layer Normalization) 레이어를 통해 토큰 임베딩에 적응적으로 주입합니다.
• 인식 모델: $x_0$와 $x_t$를 입력으로 받지만, $x_t$가 $x_0$의 부분 마스킹 버전이라는 점을 활용하여 계산 비용을 줄입니다. 마스킹된 토큰에만 AdaLN 을 적용하는 'Siamese' 구조를 사용하여 효율성을 높였습니다.

3. 주요 기여 (Key Contributions)

1. 새로운 프레임워크 제안: 이산 확산 모델에 잠재 변수 구조를 통합하여 차원 간 상관관계를 모델링하는 VADD 를 최초로 제안했습니다.
2. 효율적인 학습 및 추론: 인식 모델을 통한 변분 추론을 통해 안정적인 학습을 가능하게 했으며, 추론 시에는 인식 모델이 필요하지 않아 기존 MDM 과 동일한 빠른 추론 속도를 유지합니다.
3. 적은 단계에서의 우수한 성능: 디노이징 단계가 적을 때 발생하는 독립성 가정의 한계를 잠재 변수를 통해 해결하여, 적은 단계로도 고품질 샘플을 생성할 수 있음을 입증했습니다.

4. 실험 결과 (Results)

VADD 는 2D toy 데이터, 픽셀 단위 이미지 생성, 텍스트 생성 등 다양한 작업에서 기존 MDM (MDLM) 베이스라인과 비교하여 우수한 성능을 보였습니다.

• 2D Toy Data: 체커보드, 스위스롤, 서클 등 다중 모드 분포에서 VADD 는 1 단계 (T=1) 샘플링만으로도 MDLM 이 실패하는 복잡한 상관관계를 정확히 복원했습니다.
• 이미지 생성 (MNIST, CIFAR-10):
  • FID 점수: CIFAR-10 에서 10 단계 샘플링 시 MDLM 의 FID 334.3 대비 VADD 는 170.3 으로 압도적인 성능 향상을 보였습니다.
  • BPD (Bits Per Dimension): 비이진화 MNIST 및 CIFAR-10 에서 더 낮은 BPD 를 기록하여 생성 모델링 용량이 향상되었음을 확인했습니다.
• 텍스트 생성 (LM1B, OpenWebText):
  • Generative Perplexity: 적은 샘플링 단계 (예: 16~64 단계) 에서 MDLM 대비 현저히 낮은 생성 퍼플렉시티를 기록했습니다.
  • Zero-shot Perplexity: 6 개 벤치마크 데이터셋에서 MDLM 과 유사하거나 더 나은 성능을 보였습니다.
  • 학습 효율성: VADD 의 학습 속도는 MDLM 의 약 0.66 배 수준으로, 두 모델을 동시에 학습함에도 불구하고 2 배 이상의 오버헤드가 발생하지 않았습니다.

5. 의의 및 결론 (Significance)

이 연구는 이산 확산 모델의 핵심 약점인 '차원 간 독립성 가정'을 잠재 변수 모델링을 통해 해결했다는 점에서 의의가 큽니다. 특히, 실제 배포 환경에서 추론 속도가 중요한 경우 (적은 디노이징 단계가 필요한 상황) 에 VADD 는 기존 MDM 보다 월등히 높은 샘플 품질을 보장합니다.

또한, VADD 는 사전 훈련된 강력한 Teacher 모델에 의존하는 증류 (distillation) 기반 방법론과 달리, 훈련 데이터로부터 잠재 변수 디노이징 모델을 처음부터 학습한다는 점에서 차별화됩니다. 이는 이산 확산 모델의 표현력을 높이고, 텍스트 및 이미지 생성 분야에서 더 빠르고 고품질의 생성 모델을 구축할 수 있는 새로운 방향을 제시합니다.