Vereenvoudigde en Gegeneraliseerde Gemaskeerde Diffusie voor Discrete Gegevens

Simplified and Generalized Masked Diffusion for Discrete Data

June 6, 2024

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

cs.AI

Samenvatting

Gemaskerde (of absorberende) diffusie wordt actief onderzocht als alternatief voor autoregressieve modellen voor generatieve modellering van discrete data. Echter, bestaand werk op dit gebied is belemmerd door onnodig complexe modelformuleringen en onduidelijke relaties tussen verschillende perspectieven, wat heeft geleid tot suboptimale parameterisatie, trainingsdoelen en ad hoc aanpassingen om deze problemen tegen te gaan. In dit werk streven we ernaar om een eenvoudig en algemeen raamwerk te bieden dat het volledige potentieel van gemaskerde diffusiemodellen ontsluit. We laten zien dat het continue-tijd variatie-objectief van gemaskerde diffusiemodellen een eenvoudige gewogen integraal van kruisentropieverliezen is. Ons raamwerk maakt het ook mogelijk om gegeneraliseerde gemaskerde diffusiemodellen te trainen met toestandsafhankelijke maskeringsschema's. Wanneer geëvalueerd op perplexiteit, overtreffen onze op OpenWebText getrainde modellen eerdere diffusietaalmodellen op GPT-2-schaal en tonen ze superieure prestaties op 4 van de 5 zero-shot taalmodelleertaken. Bovendien presteren onze modellen aanzienlijk beter dan eerdere discrete diffusiemodellen op pixel-niveau beeldmodellering, met 2,78 (CIFAR-10) en 3,42 (ImageNet 64x64) bits per dimensie, wat vergelijkbaar of beter is dan autoregressieve modellen van vergelijkbare grootte.

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.