Set Block Decoding is een versneller voor taalmodelinferentie.

Set Block Decoding is a Language Model Inference Accelerator

September 4, 2025

Auteurs: Itai Gat, Heli Ben-Hamu, Marton Havasi, Daniel Haziza, Jeremy Reizenstein, Gabriel Synnaeve, David Lopez-Paz, Brian Karrer, Yaron Lipman

cs.AI

Samenvatting

Autoregressieve taalmodellen voor next token prediction bieden krachtige mogelijkheden, maar worden in de praktijk geconfronteerd met aanzienlijke uitdagingen bij implementatie vanwege de hoge reken- en geheugenkosten van inferentie, met name tijdens het decoderingstraject. Wij introduceren Set Block Decoding (SBD), een eenvoudig en flexibel paradigma dat de generatie versnelt door standaard next token prediction (NTP) en masked token prediction (MATP) te integreren binnen een enkele architectuur. SBD stelt het model in staat om meerdere, niet noodzakelijk opeenvolgende, toekomstige tokens parallel te bemonsteren, een belangrijk onderscheid ten opzichte van eerdere versnellingsmethoden. Deze flexibiliteit maakt het gebruik van geavanceerde oplossers uit de discrete diffusieliteratuur mogelijk, wat aanzienlijke snelheidswinsten biedt zonder in te leveren op nauwkeurigheid. SBD vereist geen architectuurwijzigingen of extra trainingshyperparameters, behoudt compatibiliteit met exacte KV-caching, en kan worden geïmplementeerd door bestaande next token prediction-modellen te fine-tunen. Door Llama-3.1 8B en Qwen-3 8B te fine-tunen, demonstreren we dat SBD een reductie van 3-5x in het aantal forward passes die nodig zijn voor generatie mogelijk maakt, terwijl dezelfde prestaties worden behaald als bij equivalente NTP-training.

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.