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

摘要

自回归式下一个词元预测语言模型具备强大的能力，但在实际部署中面临显著挑战，尤其是在解码阶段，其推理过程存在较高的计算与内存开销。本文提出一种简单且灵活的范式——集合块解码（Set Block Decoding, SBD），通过在单一架构中融合标准的下一个词元预测（Next Token Prediction, NTP）与掩码词元预测（Masked Autoregressive Token Prediction, MATP），实现生成过程的加速。SBD 允许模型并行采样多个非连续的未来词元，这一特性区别于以往的加速方法，具有显著优势。该灵活性使得可直接引入离散扩散模型领域中的先进求解器，从而在不牺牲准确率的前提下实现显著的提速。SBD 无需修改模型架构，也不需额外训练超参数，兼容精确的键值缓存（KV-caching）机制，且可通过微调现有下一个词元预测模型即可实现。通过对 Llama-3.1 8B 和 Qwen-3 8B 模型进行微调，我们验证了 SBD 能够在保持与传统 NTP 训练相当性能的前提下，将生成过程所需的前向传播次数减少 3 至 5 倍。