Команда из MIT показала, что непрерывная диффузия для текста может быть не просто теоретической выкладкой на бумаге - модель на 105M параметров по их замерам бьёт более крупные (~170M) дискретные и непрерывные DLMs, тратя на порядок меньше данных на обучение и меньше шагов на генерацию.
В отличие от картинок и видео, где диффузионки рулят, с текстом всегда была беда из-за дискретности токенов. В ресёрче популярны Diffusion Language Models, а непрерывные подходы буксовали.
ELF предлагает пересобрать архитектуру так, чтобы непрерывная диффузия наконец поехала.
Идея методики в том, чтобы модель перестала мучить дискретные токены на каждом шаге денойзинга:
🟠Берём замороженный энкодер (T5) и переводим дискретные токены в непрерывное пространство эмбеддингов. В инференсе его не будет, он нужен только на обучении.
🟠Запускаем Flow Matching. На протяжении всего процесса инфернса модель сидит в непрерывном пространстве, решая ODE/SDE.
🟠Дискретизация обратно в текст происходит строго на финальном шаге, при этом отдельный декодер не нужен, веса шарятся с основным денойзером. А ещё в процесс замечательно встаёт классический CFG.
В дискретных моделях CFG был малоизучен и работал заметно хуже, а здесь нормально рулит балансом качества и разнообразия генерации.
🟡Результаты тестов
🟢Метод обходит топовые дискретные (MDLM, Duo) и непрерывные (FLM, LangFlow) DLMs по генеративной перплексии и делает это, будучи меньше по размеру.
🟢На обучение понадобилось всего 45B токенов. Конкурентам нужно примерно в 10 раз больше.
🟢ELF бьёт даже дистиллированные версии конкурентов (few-step версия FLM, Duo с дистилляцией DCD) на малом числе шагов - хватает 32 шагов с SDE.
🟢Хорошие метрики на задачах с условием (то есть когда генерация опирается на вход): BLEU 26.4 в машинном переводе (WMT14 De-En) и лучшие ROUGE-цифры в XSum среди сравниваемых вариантов.
🟡Дисклеймер
Это пруф-оф-концепт. Самая большая протестированная модель ELF-L содержит всего 652M параметров.
Внутри тестового диапазона (105M → 652M) скейлинг работает и улучшает метрики, но как поведёт себя метод на 7B–70B - неизвестно.
Так что, если под рукой есть пара свободных кластеров H100 и интерес попробовать что-то новое - код открыт, метод описан, можно смело заниматься. Доступны и тестовые чекпоинты из пейпера.
📌Лицензирование: MIT License
🟡Arxiv
🟡Набор моделей
🖥GitHub
@ai_machinelearning_big_data
#AI #ML #Diffusion #ELF #MIT