Set Block Decoding is a Language Model Inference Accelerator

2025-09-04

AI Summary

Autoregressive language models incur substantial computational and memory overhead during decoding, hindering efficient deployment. To address this, we propose Set Block Decoding (SBD), the first method enabling parallel generation of non-contiguous token blocks. SBD unifies standard autoregressive prediction with masked language modeling and introduces a discrete diffusion solver to optimize multi-step parallel sampling. Crucially, SBD requires no architectural modifications, preserves exact KV cache fidelity, and is directly deployable via lightweight fine-tuning on existing open-weight models (e.g., Llama-3.1 8B, Qwen-3 8B). Experiments demonstrate that SBD reduces forward pass count by 3–5× while maintaining generation quality comparable to standard autoregressive decoding. This yields significant improvements in inference throughput without compromising accuracy or compatibility.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Diffusion Models for Vision

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

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.

Problem

• Reduces computational and memory costs of language model inference
• Accelerates generation by combining next and masked token prediction
• Samples multiple future tokens in parallel without sacrificing accuracy

Innovation

• Integrates standard and masked token prediction
• Samples multiple future tokens in parallel
• Uses discrete diffusion solvers for speedup