Simplified and Generalized Masked Diffusion for Discrete Data

Masked (or absorbing) diffusion is actively explored as an alternative to autoregressive models for generative modeling of discrete data. However, existing work in this area has been hindered by unnecessarily complex model formulations and unclear relationships between different perspectives, leading to suboptimal parameterization, training objectives, and ad hoc adjustments to counteract these issues. In this work, we aim to provide a simple and general framework that unlocks the full potential of masked diffusion models. We show that the continuous-time variational objective of masked diffusion models is a simple weighted integral of cross-entropy losses. Our framework also enables training generalized masked diffusion models with state-dependent masking schedules. When evaluated by perplexity, our models trained on OpenWebText surpass prior diffusion language models at GPT-2 scale and demonstrate superior performance on 4 out of 5 zero-shot language modeling tasks. Furthermore, our models vastly outperform previous discrete diffusion models on pixel-level image modeling, achieving 2.75 (CIFAR-10) and 3.40 (ImageNet 64x64) bits per dimension that are better than autoregressive models of similar sizes. Our code is available at https://github.com/google-deepmind/md4.

한 줄 요약

이 연구는 마스킹 확산 모델을 단순화하고 일반화하여 이산 데이터 생성 성능을 크게 향상시킨다.

핵심 기여도

마스킹 확산 모델의 ELBO가 가중치가 적용된 교차 엔트로피 손실의 적분임을 증명.
상태에 따라 마스킹 스케줄을 조정하는 GenMD4 모델 제안.
OpenWebText에서 GPT-2 규모의 확산 언어 모델을 초과, 5개 중 4개의 zero-shot 작업에서 우수한 성능.
CIFAR-10에서 2.75, ImageNet 64x64에서 3.40 bits/dim 달성, 이는 유사한 크기의 자동회귀 모델보다 우수.

핵심 아이디어

기존 마스킹 확산 모델은 복잡한 수식과 불명확한 관계로 인해 최적화가 어려웠다. 본 연구는 마스킹 확산 모델의 ELBO가 단순한 가중치가 적용된 교차 엔트로피 손실의 적분임을 보여주어, 모델 파라미터화와 학습 목적을 명확히 정의했다. 이는 모델의 이해와 최적화를 크게 단순화한다. 또한, 상태에 따라 마스킹 스케줄을 조정하는 GenMD4를 제안하여, 기존 모델에서의 제한을 극복하고 성능을 향상시켰다. 이는 기존의 ELBO 표현식을 기반으로 하되, 마스킹 스케줄이 모델의 상태에 따라 달라지도록 확장한 것이다.

기술적 접근법

**ELBO 표현식**: 마스킹 확산 모델의 ELBO는 시간에 따른 교차 엔트로피 손실의 가중치 적분으로 표현된다.
**GenMD4**: 상태에 따라 마스킹 스케줄이 조정되는 일반화된 마스킹 확산 모델.
**학습 목적**: ELBO는 Riemann 합으로 표현되며, T → ∞ 시 더 정밀한 적분이 가능하다.
**마스킹 스케줄**: α_t = σ(λ_t)로 정의하며, λ_t는 log-SNR(t)로 표현된다.
**모델 파라미터화**: μ_θ(x_t, t) ∈ Δ^{m+1}로 정의되며, softmax를 적용한 신경망 출력으로 파라미터화된다.

주요 결과

OpenWebText에서 GPT-2 규모의 확산 언어 모델을 초과, 5개 중 4개의 zero-shot 작업에서 우수한 성능.
CIFAR-10에서 2.75 bits/dim, ImageNet 64x64에서 3.40 bits/dim 달성.
이는 유사한 크기의 자동회귀 모델보다 우수한 성능을 나타낸다.
GenMD4는 상태 독립적인 모델 대비 테스트 로그 확률에서 추가적인 개선을 보인다.

의의 및 한계

본 연구는 마스킹 확산 모델을 단순화하고 일반화하여, 이산 데이터 생성 분야에서의 활용 가능성을 크게 확장했다. 특히, ELBO의 단순화는 모델 이해와 최적화를 용이하게 하며, GenMD4는 상태에 따른 마스킹 스케줄을 도입하여 성능을 향상시켰다. 그러나 일부 작업(예: text8)에서는 자동회귀 모델에 비해 여전히 경쟁력이 부족하며, GenMD4는 과적합에 취약해 zero-shot 전이 성능이 저하될 수 있다. 또한, 상태에 따른 스케줄을 사용하는 추론 과정은 더 복잡하다는 점도 한계로 지적된다.

실용적 활용

이 연구는 텍스트 생성, 이미지 생성, 그래프 생성 등 이산 데이터 생성 분야에서 활용 가능하다. 특히, GenMD4는 상태에 따라 마스킹 스케줄을 조정할 수 있어, 다양한 데이터 유형에 맞춘 유연한 모델링이 가능하다. 이는 NLP, 컴퓨터 비전, 생물정보학 등 여러 분야에서 실용적 활용이 기대된다.