SpecFold ускоряет диффузионные языковые модели до 1,99x

SpecFold ускоряет диффузионные языковые модели до 1,99x

Диффузионные большие языковые модели (DLLM) генерируют текст через итеративное поблочное удаление шума. Мультиветвевое спекулятивное декодирование ускоряет этот процесс: за один прямой проход проверяются основная ветвь и несколько ветвей-черновиков.

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

На этом наблюдении построен SpecFold, совместная разработка алгоритма и системной реализации (algorithm-system co-design). На алгоритмическом уровне метод выполняет пороговое управление остатками на уровне токенов (token-level residual gating) и выборочно переиспользует вычисления родительской ветви через «свёрнутые» (folded) слои внимания и FFN, сохраняя при этом остаточные скрытые состояния. На системном уровне реализация на ядре Triton превращает такое мелкозернистое переиспользование в реальный выигрыш по сквозной пропускной способности за счёт эффективного разреженного выполнения нескольких ветвей.

По утверждению авторов, SpecFold ортогонален кэшированию по времени и совместим с существующими стратегиями спекуляции для DLLM. В экспериментах на двух семействах DLLM, пяти моделях и пяти стандартных бенчмарках метод даёт ускорение по пропускной способности до 1,64x относительно Spiffy и до 1,99x относительно обычного (vanilla) декодирования, сохраняя сопоставимое качество решения задач.

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

  • SpecFold снижает стоимость мультиветвевой спекулятивной проверки в диффузионных языковых моделях за счёт переиспользования вычислений родительской ветви.
  • Идея: ветви-черновики наследуют большинство токенов родителя и раскрывают лишь немного новых позиций, поэтому их скрытые состояния во многом похожи.
  • Реализация: пороговое управление остатками на уровне токенов, «свёрнутые» внимание и FFN, а также ядро на Triton для разреженного выполнения ветвей.
  • Результат по заявлению авторов: до 1,64x по пропускной способности относительно Spiffy и до 1,99x относительно обычного декодирования при сопоставимом качестве.
  • Оценка проведена на двух семействах DLLM, пяти моделях и пяти стандартных бенчмарках; метод совместим с кэшированием по времени и существующими стратегиями спекуляции.

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

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

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

Прежде всего исследователям и инженерам, которые работают над выводом (инференсом) диффузионных языковых моделей и спекулятивным декодированием. Работа пригодится и тем, кто пишет ядра на Triton и ищет способы разреженного выполнения похожих по содержимому веток вычислений.

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

Практически метод встраивается поверх существующих стратегий спекуляции для DLLM: авторы заявляют совместимость с ними и с кэшированием по времени. Для внедрения нужна реализация «свёрнутых» внимания и FFN с управлением остатками на уровне токенов и ядро Triton для разреженного выполнения ветвей. Выпуск кода, лицензия и площадка публикации в аннотации не упомянуты, так что готовность к использованию придётся проверять по полному тексту статьи.

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

Все цифры взяты из аннотации и приведены самими авторами. Значения 1,64x и 1,99x заявлены как максимумы («до»), средний или медианный выигрыш не указан. Сопоставимость качества задач не выражена числами, а аппаратная конфигурация, размер батча, длина последовательности и абсолютная скорость в токенах в секунду в аннотации не названы. Не названы также авторы, организации и конкретные модели и бенчмарки.

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

Выигрыш зависит от доли общих токенов и схожих скрытых состояний у ветвей: метод опирается на это наблюдение авторов. Заявленные максимумы получены на конкретном наборе из двух семейств моделей и пяти бенчмарков, поэтому на других моделях и нагрузках результат может отличаться. Spiffy в аннотации описан лишь как базовая линия для сравнения, так что оценить, насколько она сильна, без полного текста статьи нельзя.