Диффузионные большие языковые модели (LLM) исследуются как многообещающая альтернатива традиционным авторегрессионным моделям, предлагая потенциал для одновременной генерации нескольких токенов. Используя механизмы двунаправленного внимания, эти модели стремятся ускорить декодирование, теоретически обеспечивая более быстрый вывод, чем авторегрессионные системы. Однако, несмотря на свой потенциал, на практике диффузионные модели часто не могут обеспечить конкурентоспособную скорость вывода, что ограничивает их способность соответствовать реальным показателям авторегрессионных LLM.
Основные проблемы
Основная проблема заключается в неэффективности вывода в диффузионных LLM. Эти модели обычно не поддерживают механизмы KV-кэширования, которые необходимы для ускорения вывода за счёт повторного использования ранее вычисленных состояний внимания. Без KV-кэширования каждый новый шаг генерации в диффузионных моделях повторяет полные вычисления внимания, что делает их ресурсоёмкими.
Кроме того, при одновременном декодировании нескольких токенов — ключевой особенности диффузионных моделей — качество генерации часто ухудшается из-за нарушений зависимостей между токенами в соответствии с предположением об условной независимости. Это делает диффузионные модели ненадёжными для практического применения, несмотря на их теоретические преимущества.
Попытки улучшения
Попытки улучшить диффузионные LLM сосредоточены на таких стратегиях, как блочное генерирование и частичное кэширование. Например, такие модели, как LLaDA и Dream, используют техники маскированного диффузионного генерирования для облегчения многомаркерной генерации. Однако им всё ещё не хватает эффективной системы KV-кэширования, а параллельное декодирование в этих моделях часто приводит к несогласованным результатам.
Хотя некоторые подходы используют вспомогательные модели для аппроксимации зависимостей между токенами, эти методы вносят дополнительную сложность, не решая полностью основные проблемы производительности. В результате скорость и качество генерации в диффузионных LLM продолжают отставать от авторегрессионных моделей.
Fast-dLLM
Исследователи из NVIDIA, Гонконгского университета и MIT представили Fast-dLLM — фреймворк, разработанный для устранения этих ограничений без необходимости переобучения. Fast-dLLM вносит в диффузионные LLM две инновации:
1. Блочно-приблизительный механизм KV-кэширования.
2. Стратегию параллельного декодирования с учётом уверенности.
Приблизительное KV-кэширование адаптировано для двунаправленной природы диффузионных моделей, позволяя эффективно повторно использовать активации из предыдущих шагов декодирования. Параллельное декодирование с учётом уверенности выборочно декодирует токены на основе порога уверенности, снижая ошибки, возникающие из-за предположения о независимости токенов.
Этот подход обеспечивает баланс между скоростью и качеством генерации, делая его практическим решением для задач генерации текста на основе диффузии.
Реализация
Метод KV-кэширования Fast-dLLM реализован путём разделения последовательностей на блоки. Перед генерацией блока вычисляются и сохраняются KV-активации для других блоков, что позволяет повторно использовать их на последующих этапах декодирования. После генерации блока кэш обновляется для всех токенов, что минимизирует избыточность вычислений при сохранении точности.
Версия DualCache расширяет этот подход, кэшируя как префиксные, так и суффиксные токены, используя высокое сходство между соседними шагами вывода, как показано тепловыми картами косинусного сходства в статье.
Для компонента параллельного декодирования система оценивает уверенность каждого токена и декодирует только те, которые превышают установленный порог. Это предотвращает нарушения зависимостей при одновременном сэмплировании и обеспечивает более высокое качество генерации даже при одновременном декодировании нескольких токенов.
Результаты тестирования
Fast-dLLM продемонстрировал значительные улучшения производительности в тестовых испытаниях. Например, на наборе данных GSM8K он достиг ускорения в 27,6 раз по сравнению с базовыми моделями в 8-шаговых конфигурациях при длине генерации 1024 токена с точностью 76,0%. На бенчмарке MATH было достигнуто ускорение в 6,5 раз с точностью около 39,3%. В бенчмарке HumanEval ускорение составило до 3,2 раз при сохранении точности на уровне 54,3%, а на MBPP система достигла ускорения в 7,8 раз при длине генерации 512 токенов.
Во всех задачах и моделях точность оставалась в пределах 1–2 пунктов от базовой, показывая, что ускорение Fast-dLLM не приводит к существенному снижению качества вывода.
Исследовательская группа эффективно устранила основные узкие места в диффузионных LLM, представив новую стратегию кэширования и механизм декодирования, основанный на уверенности. Устранив неэффективность вывода и повысив качество декодирования, Fast-dLLM демонстрирует, как диффузионные LLM могут приблизиться или даже превзойти авторегрессионные модели по скорости, сохраняя при этом высокую точность. Это делает их пригодными для использования в реальных приложениях генерации языка.