Ускорение диффузионных языковых моделей с помощью адаптивного параллельного декодирования

Accelerating Diffusion LLMs via Adaptive Parallel Decoding

May 31, 2025

Авторы: Daniel Israel, Guy Van den Broeck, Aditya Grover

cs.AI

Аннотация

Скорость генерации больших языковых моделей (LLM) ограничена авторегрессивным декодированием, при котором токены предсказываются последовательно один за другим. В качестве альтернативы, диффузионные большие языковые модели (dLLM) теоретически позволяют генерировать токены параллельно, но на практике сталкиваются с трудностями в достижении скорости авторегрессивных моделей без значительного ухудшения качества. Поэтому мы представляем адаптивное параллельное декодирование (APD) — новый метод, который динамически регулирует количество токенов, выбираемых параллельно. Мы достигаем этого, определяя мультипликативную смесь между маргинальными вероятностями dLLM и совместной вероятностью последовательностей, полученной с помощью небольшой вспомогательной авторегрессивной модели. Это инвертирует стандартную схему спекулятивного декодирования, где цель состоит в выборке из большой авторегрессивной модели с использованием черновиков, созданных меньшей моделью. Мы дополнительно оптимизируем APD, включая кэширование ключей и значений (KV) и ограничивая размер маскированного ввода. В целом, наш метод предлагает три настраиваемых параметра для гибкого баланса между пропускной способностью и качеством. Мы показываем, что APD обеспечивает значительно более высокую пропускную способность с минимальным ухудшением качества на тестовых наборах данных.

English

The generation speed of LLMs are bottlenecked by autoregressive decoding, where tokens are predicted sequentially one by one. Alternatively, diffusion large language models (dLLMs) theoretically allow for parallel token generation, but in practice struggle to achieve the speed of autoregressive models without significantly sacrificing quality. We therefore introduce adaptive parallel decoding (APD), a novel method that dynamically adjusts the number of tokens sampled in parallel. We achieve this by defining a multiplicative mixture between the dLLM marginal probabilities and the joint probability of sequences under a small auxiliary autoregressive model. This inverts the standard setup of speculative decoding, where the goal is to sample from a large autoregressive verifier by drafting from a smaller model. We further optimize APD by enabling KV caching and limiting the size of the masked input. Altogether, our method puts forward three tunable parameters to flexibly tradeoff throughput and quality. We show that APD provides markedly higher throughput with minimal quality degradations on downstream benchmarks.