Set Block Decoding is a Language Model Inference Accelerator
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.
68beb7f7c123124955ef61b5
Itai Gat
68beb7f7c123124955ef61b6
Heli Ben-Hamu
68beb7f7c123124955ef61b7
Marton Havasi
68beb7f7c123124955ef61b8
Daniel Haziza
68beb7f7c123124955ef61b9
Jeremy Reizenstein
68beb7f7c123124955ef61ba
Gabriel Synnaeve
68beb7f7c123124955ef61bb
David Lopez-Paz
68beb7f7c123124955ef61bc
Brian Karrer
68beb7f7c123124955ef61bd
Yaron Lipman
Set Block Decoding is a Language Model Inference Accelerator
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.
Set Block Decoding accelerates language model generation by integrating next token prediction and masked token prediction, enabling parallel sampling of future tokens and reducing computational cost without sacrificing accuracy.
autoregressive next token prediction masked token prediction Set Block Decoding discrete diffusion forward passes Llama-3.1 8B Qwen-3 8B