MDM-Prime: Революция в Masked Diffusion Models с частичным размаскированием токенов

Masked Diffusion Models и их ограничения

Masked Diffusion Models (MDMs) эффективно генерируют дискретные данные, такие как текст и символические последовательности, постепенно размаскируя токены на нескольких шагах. На каждом шаге токены либо маскируются, либо размаскируются. Однако существует значительная неэффективность: до 37% шагов обратного процесса не изменяют последовательность, что приводит к избыточной обработке и потере вычислительных ресурсов.

Развитие дискретных диффузионных моделей

Дискретные диффузионные модели эволюционировали от простых бинарных данных до сложных задач, включая генерацию текста и изображений с использованием различных стратегий шума. Недавние инновации включают упрощение целей обучения, комбинирование MDM с авторегрессивными методами, управление с помощью energy-based моделей, выборочное повторное маскирование и методы дистилляции для сокращения количества шагов сэмплинга. Подходы с использованием непрерывного шума для дискретных данных, такие как Bit Diffusion, сталкиваются с проблемами вычислимости правдоподобия из-за квантизации.

Введение MDM-Prime и частичного маскирования

MDM-Prime, разработанная исследователями из Vector Institute, NVIDIA и Национального университета Тайваня, представляет новую схему частичного маскирования (Prime). Вместо бинарного маскирования Prime позволяет токенам принимать промежуточные состояния, маскируя части их закодированного представления. Такой постепенный процесс раскрытия информации улучшает качество предсказаний и снижает избыточные вычисления.

Архитектура и улучшения в обучении

MDM-Prime разбивает токены на последовательности суб-токенов с помощью обратимой функции, что обеспечивает более плавные промежуточные состояния диффузии и уменьшает количество простоя. Обратный процесс обучения строится на вариационной верхней оценке по суб-токенам. Для учета зависимостей между суб-токенами и предотвращения некорректных выходов модель обучается совместному распределению вероятностей с фильтрацией неконсистентных последовательностей. Архитектура включает эффективный энкодер-декодер, оптимизированный для обработки суб-токенов.

Результаты на задачах генерации текста и изображений

Оценка MDM-Prime на генерации текста с использованием OpenWebText показывает значительное улучшение перплексии (15.36) и снижение доли простоя, особенно при гранулярности суб-токенов ℓ ≥ 4. Модель превосходит предыдущие методы без применения авторегрессии и хорошо обобщается на различных zero-shot задачах. В генерации изображений на CIFAR-10 и ImageNet-32 MDM-Prime (ℓ = 2) достигает лучшего качества сэмплов и более низких значений FID, опережая базовые модели и сохраняя эффективность. Она также хорошо справляется с условной генерацией, предсказывая маскированные суб-токены по частично наблюдаемым изображениям.

Влияние и перспективы

MDM-Prime представляет собой новый этап в генерации дискретных данных с более детальным и эффективным подходом благодаря частичному размаскированию токенов. Это снижает повторные вычисления и повышает выразительность моделей генерации. Высокие показатели на текстовых и визуальных задачах демонстрируют потенциал MDM-Prime как мощного инструмента для различных применений.

Дополнительную информацию можно найти в статье, проекте и репозитории GitHub. Следите за исследователями в Twitter и присоединяйтесь к сообществу для обновлений.