🌟 Embedded Language Flows: диффузия для текста оживает

Команда из 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