Uno ускоряет вывод LLM до 3 раз без потери качества

Uno ускоряет вывод LLM до 3 раз без потери качества

Обычные большие языковые модели (LLM) строятся на предсказании следующего токена (next-token prediction, NTP): генерация идёт последовательно, токен за токеном, и это узкое место замедляет вывод. Сахо с соавторами предложили новый класс моделей, «диффузионно-расширенные LLM» (diffusion-augmented LLMs), которые сохраняют авторегрессионное распределение модели, но используют диффузию, чтобы вытягивать из этого распределения сразу несколько токенов параллельно.

Параметры такой модели разделены на два набора: авторегрессионные веса обучаются как обычно, по цели NTP, а лёгкие «диффузионные» веса, на генерацию нескольких токенов одновременно, через отдельную стадию обучения, которую авторы называют дистилляцией диффузии (Diffusion Distillation). Эта стадия добавляет, по их словам, незначительные накладные расходы к уже существующим пайплайнам обучения LLM. Дополнительно авторы представили Ψ-Spec, семейство сэмплеров, которое даёт лосслес-ускорение (без потери качества) и позволяет наращивать вычисления на этапе вывода при фиксированной длине контекста.

В отличие от спекулятивного декодирования, метод не требует отдельной модели-черновика. В отличие от диффузионных LLM (d-LLM), он ускоряет генерацию, не жертвуя качеством исходной авторегрессионной модели. Получившиеся модели, названные Uno, можно обучать с нуля или получать, дообучая уже существующие открытые авторегрессионные LLM.

По заявлению авторов, Uno обходит по пропускной способности ведущие методы спекулятивного декодирования при любом размере пакета (батча) запросов и даёт ускорение до 3 раз относительно базовой авторегрессионной модели, в том числе при максимальном размере пакета, который поддерживает устройство. Версия Uno на 8 млрд параметров, по их данным, превосходит по всем протестированным бенчмаркам, в агентном использовании инструментов, кодинге и рассуждениях на длинном контексте, как ведущую открытую диффузионную LLM DiffusionGemma на 26 млрд параметров, так и проприетарную модель Mercury 2. Код и веса модели (чекпоинты) авторы выложили в открытом доступе.

Ключевые факты

  • Diffusion-augmented LLM «Uno» разделяет параметры на авторегрессионные веса (обучены обычным предсказанием следующего токена) и лёгкие диффузионные веса, обученные генерировать несколько токенов параллельно
  • Диффузионные веса добавляются отдельной стадией, дистилляцией диффузии, которая, по словам авторов, вносит незначительные накладные расходы в существующий пайплайн обучения
  • Метод даёт ускорение до 3 раз относительно базовой авторегрессионной модели при любом размере пакета запросов, включая максимальный, и обходит по пропускной способности ведущие методы спекулятивного декодирования
  • В отличие от спекулятивного декодирования, отдельная модель-черновик не нужна; в отличие от диффузионных LLM, качество исходной модели не теряется
  • Версия Uno на 8 млрд параметров превосходит по всем заявленным бенчмаркам открытую DiffusionGemma (26 млрд параметров) и проприетарную Mercury 2; код и веса модели выложены авторами в открытом доступе

Почему это важно

Последовательная, токен-за-токеном генерация, фундаментальное узкое место авторегрессионных LLM, из-за которого инференс остаётся медленным даже на мощном железе. Uno предлагает обойти это ограничение, не переходя целиком на диффузионную архитектуру (которая обычно жертвует качеством) и не заводя отдельную модель-черновик, как это делает спекулятивное декодирование. По утверждению авторов, ускорение при этом лосслес, то есть без потери качества исходной модели.

Кому это важно

В первую очередь тем, кто разворачивает и обслуживает LLM в проде и упирается в стоимость и задержку инференса: инфраструктурным командам, разработчикам инференс-серверов и всем, кто уже использует открытые авторегрессионные модели и хочет ускорить их без переобучения с нуля, Uno можно получить, дообучив существующую открытую AR-модель.

Как это применить

Авторы выложили код и веса модели (чекпоинты) в открытом доступе по ссылке в публикации. Модель Uno можно как обучать с нуля, так и получать из уже существующих открытых авторегрессионных LLM через отдельную, по описанию лёгкую, стадию дистилляции диффузии, то есть встроить метод в уже существующий пайплайн обучения, а не строить архитектуру заново.

Можно ли доверять

Все цифры и выводы в тексте, это заявления самих авторов препринта, без независимой проверки. В самом описании не названы ни конкретные бенчмарки, ни их баллы, сказано лишь, что Uno «превосходит» перечисленные модели «по всем протестированным бенчмаркам». Не указаны также абсолютные значения задержки или конкретное устройство, на котором измерялось ускорение, есть только относительная цифра «до 3 раз».

Риски и подводные камни

«До 3 раз», это верхняя граница, а не гарантированное ускорение на любой задаче и любом железе; конкретное устройство измерения не названо. Детали бенчмарков, объём данных обучения и вычислительная стоимость стадии дистилляции диффузии в тексте не раскрыты, есть лишь оценка авторов, что накладные расходы «незначительны». Утверждения о превосходстве над DiffusionGemma и Mercury 2 также нельзя проверить без доступа к самим бенчмаркам.