Set Block Decoding è un acceleratore per l'inferenza nei modelli linguistici.

Set Block Decoding is a Language Model Inference Accelerator

September 4, 2025

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

cs.AI

Abstract

I modelli linguistici autoregressivi per la previsione del token successivo offrono capacità potenti, ma affrontano sfide significative nella distribuzione pratica a causa degli elevati costi computazionali e di memoria durante l'inferenza, in particolare nella fase di decodifica. Introduciamo il Set Block Decoding (SBD), un paradigma semplice e flessibile che accelera la generazione integrando la previsione standard del token successivo (NTP) e la previsione del token mascherato (MATP) all'interno di un'unica architettura. SBD consente al modello di campionare più token futuri, non necessariamente consecutivi, in parallelo, una distinzione chiave rispetto ai metodi di accelerazione precedenti. Questa flessibilità permette l'uso di risolutori avanzati tratti dalla letteratura sulla diffusione discreta, offrendo accelerazioni significative senza sacrificare l'accuratezza. SBD non richiede modifiche architetturali o iperparametri di addestramento aggiuntivi, mantiene la compatibilità con la KV-caching esatta e può essere implementato tramite il fine-tuning di modelli esistenti per la previsione del token successivo. Mediante il fine-tuning di Llama-3.1 8B e Qwen-3 8B, dimostriamo che SBD consente una riduzione di 3-5x nel numero di passaggi in avanti necessari per la generazione, mantenendo le stesse prestazioni dell'addestramento NTP equivalente.

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.