이산 데이터를 위한 단순화 및 일반화된 마스크 확산 모델

Simplified and Generalized Masked Diffusion for Discrete Data

June 6, 2024

저자: Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, Michalis K. Titsias

cs.AI

초록

마스킹(또는 흡수) 확산 모델은 이산 데이터의 생성 모델링을 위한 자기회귀 모델의 대안으로 활발히 연구되고 있습니다. 그러나 이 분야의 기존 연구는 불필요하게 복잡한 모델 구성과 다양한 관점 간의 명확하지 않은 관계로 인해 최적의 매개변수화, 훈련 목표, 그리고 이러한 문제를 해결하기 위한 임시 조정에 제약을 받아왔습니다. 본 연구에서는 마스킹 확산 모델의 전체 잠재력을 발휘할 수 있는 간단하고 일반적인 프레임워크를 제공하는 것을 목표로 합니다. 우리는 마스킹 확산 모델의 연속 시간 변분 목적 함수가 교차 엔트로피 손실의 단순한 가중치 적분임을 보여줍니다. 또한, 우리의 프레임워크는 상태 의존적 마스킹 스케줄을 사용하여 일반화된 마스킹 확산 모델을 훈련할 수 있도록 합니다. OpenWebText 데이터셋에서 훈련된 우리의 모델은 GPT-2 규모의 이전 확산 언어 모델을 perplexity 기준으로 능가하며, 5개의 제로샷 언어 모델링 작업 중 4개에서 우수한 성능을 보였습니다. 더 나아가, 우리의 모델은 픽셀 수준의 이미지 모델링에서 이전의 이산 확산 모델을 크게 능가하며, CIFAR-10에서 2.78, ImageNet 64x64에서 3.42 비트/차원을 달성하여 유사한 크기의 자기회귀 모델과 비교할 만하거나 더 나은 성능을 보였습니다.

English

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.78~(CIFAR-10) and 3.42 (ImageNet 64times64) bits per dimension that are comparable or better than autoregressive models of similar sizes.