O Set Block Decoding é um Acelerador de Inferência para Modelos de Linguagem.

Set Block Decoding is a Language Model Inference Accelerator

September 4, 2025

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

cs.AI

Resumo

Modelos de linguagem de previsão autoregressiva do próximo token oferecem capacidades poderosas, mas enfrentam desafios significativos na implantação prática devido aos altos custos computacionais e de memória durante a inferência, particularmente na etapa de decodificação. Introduzimos o Decodificação em Blocos de Conjunto (Set Block Decoding - SBD), um paradigma simples e flexível que acelera a geração ao integrar a previsão padrão do próximo token (Next Token Prediction - NTP) e a previsão de tokens mascarados (Masked Token Prediction - MATP) em uma única arquitetura. O SBD permite que o modelo amostre múltiplos tokens futuros, não necessariamente consecutivos, em paralelo, uma distinção crucial em relação aos métodos de aceleração anteriores. Essa flexibilidade possibilita o uso de solucionadores avançados da literatura de difusão discreta, oferecendo ganhos significativos de velocidade sem sacrificar a precisão. O SBD não requer alterações arquiteturais ou hiperparâmetros de treinamento adicionais, mantém compatibilidade com o armazenamento exato de chaves-valores (KV-caching) e pode ser implementado por meio do ajuste fino de modelos existentes de previsão do próximo token. Ao ajustar finamente os modelos Llama-3.1 8B e Qwen-3 8B, demonstramos que o SBD permite uma redução de 3 a 5 vezes no número de passagens diretas necessárias para a geração, mantendo o mesmo desempenho obtido com o treinamento equivalente em 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.