Исследователи Yandex Research предложили способ ускорить генерацию текста с помощью диффузионных языковых моделей. Новый метод помогает таким нейросетям выполнять больше вычислений параллельно, не теряя в точности ответов. В экспериментах с моделями на 16 млрд параметров качество на некоторых задачах даже улучшилось.

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

Но добиться хорошего качества при таком способе генерации непросто. Для этого модель дополнительно обучают после основной тренировки. Обычно такие методы могут требовать сложных вычислений или дополнительных моделей.

Команда Yandex Research предложила более простой подход. Он сравнивает тексты, которые генерирует нейросеть, с реальными примерами из обучающих данных. Для сравнения используется математический метод Maximum Mean Discrepancy (MMD). Он оценивает, насколько отличаются два набора данных.

Сравнение происходит не напрямую по словам, а по внутренним признакам, которые извлекает уже обученная языковая модель. Её параметры при этом остаются неизменными. Это позволяет использовать её как своеобразный измерительный инструмент, не обучая ещё одну большую модель специально для оценки результатов.

Авторы проверили подход на двух типах диффузионных моделей — работающих с дискретными токенами и с непрерывными представлениями данных.

Что показали эксперименты

На текстах из набора OpenWebText новый метод помог модели лучше воспроизводить закономерности человеческой речи, сохранив разнообразие ответов. На математических задачах GSM8K модель стала находить более удачное соотношение между точностью ответов и количеством вычислительных шагов. Чем меньше шагов нужно для получения правильного ответа, тем быстрее модель потенциально может справляться с задачей.

Метод также проверили на моделях с 16 млрд параметров. Их дополнительно обучали на восьми ускорителях NVIDIA H100, и, по данным исследователей, это занимало около 1,3–1,9 минуты. На четырёх математических тестах модели за один вычислительный шаг обрабатывали на 10–16% больше частей текста, сохраняя прежнюю точность или немного улучшая её. В задачах по программированию точность ответов выросла на 2,4–3,8 процентного пункта.

Диффузионные нейросети могут изменить способ генерации текста. Однако они  пока распространены меньше, чем нейросети, которые последовательно предсказывают следующий токен.