Упрощенная и обобщенная маскированная диффузия для дискретных данных

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 и демонстрируют превосходную производительность на 4 из 5 нулевых задач языкового моделирования. Более того, наши модели значительно превосходят предыдущие дискретные модели диффузии в моделировании изображений на уровне пикселей, достигая 2,78 (CIFAR-10) и 3,42 (ImageNet 64 на 64) бита на размерность, которые сравнимы или лучше, чем авторегрессионные модели схожего размера.

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.