Set Block Decoding — это ускоритель вывода языковых моделей.

Set Block Decoding is a Language Model Inference Accelerator

September 4, 2025

Авторы: Itai Gat, Heli Ben-Hamu, Marton Havasi, Daniel Haziza, Jeremy Reizenstein, Gabriel Synnaeve, David Lopez-Paz, Brian Karrer, Yaron Lipman

cs.AI

Аннотация

Авторегрессионные языковые модели, предсказывающие следующий токен, обладают мощными возможностями, но сталкиваются с существенными трудностями при практическом внедрении из-за высоких вычислительных затрат и требований к памяти, особенно на этапе декодирования. Мы представляем Set Block Decoding (SBD) — простую и гибкую парадигму, которая ускоряет генерацию, интегрируя стандартное предсказание следующего токена (NTP) и предсказание маскированных токенов (MATP) в рамках единой архитектуры. SBD позволяет модели параллельно выбирать несколько, не обязательно последовательных, будущих токенов, что является ключевым отличием от предыдущих методов ускорения. Эта гибкость позволяет использовать продвинутые решатели из литературы по дискретной диффузии, обеспечивая значительное ускорение без потери точности. SBD не требует изменений архитектуры или дополнительных гиперпараметров обучения, сохраняет совместимость с точным кэшированием ключей и значений (KV-caching) и может быть реализована путем тонкой настройки существующих моделей предсказания следующего токена. На примере тонкой настройки моделей Llama-3.1 8B и Qwen-3 8B мы демонстрируем, что SBD позволяет сократить количество прямых проходов, необходимых для генерации, в 3–5 раз, сохраняя при этом ту же производительность, что и эквивалентное обучение NTP.

English

Autoregressive next token prediction language models offer powerful capabilities but face significant challenges in practical deployment due to the high computational and memory costs of inference, particularly during the decoding stage. We introduce Set Block Decoding (SBD), a simple and flexible paradigm that accelerates generation by integrating standard next token prediction (NTP) and masked token prediction (MATP) within a single architecture. SBD allows the model to sample multiple, not necessarily consecutive, future tokens in parallel, a key distinction from previous acceleration methods. This flexibility allows the use of advanced solvers from the discrete diffusion literature, offering significant speedups without sacrificing accuracy. SBD requires no architectural changes or extra training hyperparameters, maintains compatibility with exact KV-caching, and can be implemented by fine-tuning existing next token prediction models. By fine-tuning Llama-3.1 8B and Qwen-3 8B, we demonstrate that SBD enables a 3-5x reduction in the number of forward passes required for generation while achieving same performance as equivalent NTP training.