让扩散语言模型一次前向传播多猜几个词,又快又准。
Multi-Token Residual Prediction
- 用残差预测技术,在一次模型前传中实现多词并行去噪。
- 在代码生成任务上准确率提升22.6点,推理速度最高快1.4倍。
- 适合追求高精度或高吞吐的文本生成场景,尤其适合扩散模型应用。
扩散语言模型(DLM)通过迭代去噪掩码词元序列生成文本,相比自回归模型在并行性和生成质量间取得平衡。当前实践中,每步解码的词元数量由置信度阈值控制,随着每步解码词元增多,质量持续下降。本文提出多词残差预测(MRP),一个轻量级模块,可在单次主干网络前向传播中实现依赖感知的多词去噪。MRP利用去噪过程中相邻步骤的对数概率分布高度相似这一特性:不需再次运行主干网络获取下一步对数概率,而是从主干隐藏状态预测步骤间的残差,从而以极低成本实现每前向传播更多词元的去噪。我们将MRP应用于DLM解码的两种模式。在高精度低吞吐的静态去噪模式下,MRP作为推测解码的草稿器,其提议经主干验证,实现无损加速,最高达1.4倍(SGLang)。在低精度高吞吐的动态去噪模式下,MRP驱动一种重掩码机制,撤销过激的词元暴露,恢复因激进低阈值解码丢失的准确性,在代码生成任务HumanEval上准确率提升22.6点,在推理任务GSM8K上提升17.7点。
原文摘要 · Abstract (English)
Diffusion Language Models (DLMs) generate text by iteratively denoising masked token sequences, offering a tradeoff between parallelism and quality compared to autoregressive models. In current practice, the number of tokens decoded per step is controlled by a confidence threshold, and quality degrades monotonically as more tokens are denoised per step. We introduce Multi-token Residual Prediction (MRP), a lightweight module that enables dependency-aware multi-token denoising within a single backbone forward pass. MRP exploits a key property of the denoising process: the logit distributions at adjacent denoising steps are remarkably similar. Rather than running the backbone a second time to obtain the next-step logits, MRP predicts the residual between steps from the backbone's hidden states, effectively denoising more tokens per backbone forward at a fraction of the cost. We apply MRP across the two operating regimes of DLM decoding. In the high-quality-low-throughput static denoising regime, MRP serves as a drafter for speculative decoding: its proposals are verified against the backbone, yielding lossless acceleration of up to 1.4x in SGLang. In the low-quality-high-throughput dynamic denoising regime, MRP instead drives a remasking scheme that revokes over-eager reveals, recovering most of the accuracy lost to aggressive low-threshold decoding and improving accuracy by up to 22.6 points on code generation task HumanEval and 17.7 points on reasoning task GSM8K.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。