Multi-Token Residual Prediction¶
会议: NeurIPS 2026
arXiv: 2605.18817
代码: https://github.com/heavyball-research/multi-token-residual-prediction
领域: LLM效率
关键词: 扩散语言模型、残差预测、投机解码、重新掩码、块级扩散
一句话总结¶
本文提出多词元残差预测（MRP），利用扩散语言模型相邻去噪步间输出分布高相似度的特性，以轻量级模块直接预测步间对数概率残差，在静态去噪下实现高达 1.4 倍的无损投机加速，在动态低阈值去噪下通过重掩码机制恢复多达 22.6 点的下游任务准确率。
研究背景与动机¶
语言模型的设计长期面临自回归顺序依赖与非自回归并行度之间的核心权衡。传统的自回归（AR）模型保证了生成质量，但逐词生成的解码范式受限于访存带宽与严格串行瓶颈；扩散语言模型（DLM）通过对掩码序列进行多步去噪实现词元级并行生成，却常常以牺牲生成质量为代价。近期的块级扩散（Block Diffusion）与基于置信度阈值的自适应解码算法虽然在一定程度上拓展了帕累托前沿，但依然困于质量与吞吐量的单调权衡：在单步内强行揭示更多词元必然导致质量断崖式下跌，其核心症结在于单次前向中同时解码的多个词元彼此缺乏依赖感知。
面对这一矛盾，现有的投机解码或多步重审机制通常需要调用完整骨干网络进行二次前向验证或重新预测完整输出分布，带来了巨大的计算开销。然而，深入观察扩散去噪轨迹可以发现一个关键物理事实：当相邻两个去噪步骤之间仅揭示少数几个词元时，骨干网络的隐状态与 logits 分布高度相似，两次输出之间的残差幅值远小于状态本身的绝对幅值。直接预测全量词元分布需要模型拟合庞大的动态范围且极易导致误差累积，而步间残差天然是一个低动态范围、收敛稳定的预测目标，零函数本身就是一个优秀的初始逼近。
基于上述洞察，本文提出了轻量级步间残差学习框架，旨在不增加沉重计算负担的前提下打通多词元依赖建模。核心 idea：利用扩散去噪相邻步输出分布差异极小的结构特性，设计仅含 3 层的极轻量 Transformer 模块直接从骨干隐状态预测步间 logits 残差，并分别作为投机草稿生成器与动态重掩码校验器无缝赋能静态与动态去噪。
方法详解¶
整体框架¶
MRP 建立在以 SDAR 为代表的块级掩码扩散语言模型之上。在生成长度为 \(L\) 的词元块时，序列从全掩码状态 \(\bm{x}_T\) 开始，通过双向 Transformer 骨干网络 \(f\) 进行多轮去噪。骨干网络在第 \(t\) 步输出隐藏状态 \(\bm{h}_t\) 与未归一化 logits \(\bm{\ell}_t = f(\bm{x}_t)\)，并通过 softmax 得到每个位置的置信度。MRP 模块 \(g\) 冻结原骨干网络、语言模型头（LM Head）与词表嵌入，仅引入极小规模的附加网络层，接收骨干网络当前隐状态 \(\bm{h}_t\) 与新揭示位置构成的序列表示，直接预测下一步的状态增量 \(\bm{\Delta}_t\) 与 logits 残差 \(\bm{\delta}_t\)。
在推理阶段，MRP 根据不同的部署需求提供协同互补的工作机制：在强调无损质量的静态去噪中充当投机解码草稿模型（或直接贪心外推），而在强调超高吞吐量的低阈值动态去噪中充当重新掩码（Remasking）审校器。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入掩码状态 x_t"] --> B["步间残差收缩建模<br/>骨干网络单次前向输出 h 与 l_t"]
B --> C["轻量级 MRP 模块架构<br/>输入 h 预测步间增量 Δ 与残差 δ"]
C --> D["单向 KL 散度蒸馏训练<br/>对齐教师双步差异并反传更新 MRP"]
D --> E["双运行范式灵活部署<br/>静态投机解码草稿 / 动态重掩码过滤"]关键设计¶
1. 步间残差收缩建模：奠定残差学习优于全量预测的理论与经验基础 为了克服单次前向无法建模多词元协同依赖的瓶颈，本文深入刻画了去噪马尔可夫链状态转移的扰动特性。定义相邻去噪状态间的 logit 残差为 \(\bm{\delta}_t = \bm{\ell}_{t-1} - \bm{\ell}_t\)。由于每步仅揭示集合 \(\mathcal{R}_t\) 中的少量词元，对于其余掩蔽位置 \(i \notin \mathcal{R}_t\)，其概率分布的变化完全由注意力机制引入的间接扰动驱动。假定骨干网络关于输入嵌入满足 \(\kappa\)-Lipschitz 连续，本文证明了单步全变差距离（Total Variation Distance）收缩界： $\(D_{\mathrm{TV}}(\bm{\pi}_{t-1}^{i},\,\bm{\pi}_{t}^{i}) \le \kappa \cdot \frac{|\mathcal{R}_t|}{L} \cdot \max_{j\in\mathcal{R}_t} \|\bm{e}_{v_j} - \bm{e}_{\texttt{[MASK]}}\|_2\)$ 由于在实际解码中扰动比例 \(|\mathcal{R}_t| / L \ll 1\)，且随着去噪推进预测概率趋于集中，残差幅值随着去噪步呈现单调衰减。经验均方根（RMS）测量进一步表明，揭示 1 个词元时残差仅占全量 logits 强度的约 20%，证明了残差预测比直接拟合复杂分布在数值与优化上更具优越性。
2. 轻量级 MRP 模块架构：复用骨干表征的低开销增量头 MRP 架构吸纳了 EAGLE 与多词元预测（MTP）的设计精髓，但关键变革在于其输出为残差偏置而非绝对分布。该模块由仅 \(D=3\) 层的 Transformer 组成，参数量远小于骨干模型。在具体计算中，MRP 接收骨干网络顶层隐状态 \(\bm{h}_t\) 以及包含最新揭示词元嵌入的序列，通过可学习线性投影融合后送入 3 层 Transformer 处理，最终复用骨干网络自身已冻结的 LM Head 投影得到 logits 修正项 \(\bm{\delta}_t\)。校正后的预测分布直接通过线性加和计算：\(\bm{\ell}_{t-1} \approx \bm{\ell}_t + \bm{\delta}_t\)。这种残差参数化使得小模型无需记忆完整的词表高维先验，而仅专注于学习上下文因果联动的微小修正量。
3. 单向 KL 散度蒸馏训练：自监督模拟真实去噪演进 针对训练目标，MRP 采用基于干净语料自监督构造的蒸馏策略。给定真实序列 \(\bm{x}_0\)，首先按均匀分布采样掩码比例破坏响应词元生成带噪输入 \(\bm{x}_t\)；随后每块揭示 \(r\) 个真实词元模拟单步去噪后得到 \(\bm{x}_{t-1}\)。冻结的骨干网络分别处理两份输入以生成学生输入底座与无梯度的教师目标 \(\bm{\ell}_{t-1}\)。MRP 仅在仍处于掩蔽状态的位置上计算与教师输出的向前 KL 散度： $\(\mathcal{L} = \mathrm{KL}\Big(\mathrm{softmax}(\bm{\ell}_{t-1}) \;\Big\|\; \mathrm{softmax}(\bm{\ell}_t + \bm{\delta}_t)\Big)\Big|_{\text{masked}}\)$ 由于 logits 本质为非归一化对数概率，softmax 归一化项天然相消，使该目标严格等价于条件概率拟合。训练阶段展开 \(K=2\) 步反向传播，仅更新轻量 MRP 模块的权重，骨干网络完全保持冻结。
4. 双运行范式灵活部署：静态投机加速与动态重新掩码 训练完成的统一 MRP 模块无需微调即可自适应两种推理模式： - 静态去噪（投机解码模式）：骨干网络揭示高置信词元后，MRP 模块快速迭代 \(K\) 步生成未来位置的草稿预测；骨干网络在下一步单次并行验证所有草稿位置（\(\arg\max \bm{\ell}'^i = \arg\max \bm{\ell}^i\)），验证通过者直接采纳，分歧位置撤销重新掩码。由于采纳分布严格与骨干网络一致，实现 100% 无损加速，同时验证步的中间激活可直接复用为下一轮基座，有效分摊了验证开销。 - 动态去噪（重新掩码模式）：在低置信度阈值 \(\tau\) 下，骨干网络激进揭示大量候选集合 \(\mathcal{R}\)。紧接着调用单次 MRP 前向生成校正置信度，对于集合 \(\mathcal{R}\) 中二次置信度跌落至 \(\tau\) 以下的位置执行重新掩码撤回。通过低成本双重确认过滤偶发假阳性，以极小延迟代价挽回激进解码造成的语义崩塌。
损失函数 / 训练策略¶
模型训练基于开源 SFTDatasetV3（涵盖 OpenHermes-2.5、GenQA 与 Infinity-Instruct 的 12.4M 多轮对话数据）。超参数设置严格遵循稳定蒸馏准则：采用 AdamW 优化器，学习率峰值为 \(1\times 10^{-3}\) 并配合余弦衰减；蒸馏温度 \(T_{\mathrm{KD}}=1.0\)；每个序列长度为 4096，块大小为 16；训练展开步数 \(K=2\)，权重初始化标准差 \(\sigma_{\mathrm{init}}=0.2\)。所有主干模型参数、LM Head 及嵌入层在微调全过程中严格锁定。
实验关键数据¶
主实验¶
在配备单张 NVIDIA H100 GPU 的 SGLang 推理引擎上，评估覆盖了 SDAR-Chat（1.7B、4B、8B）三款不同规模骨干网络，对比基线包含原始单步揭示基线（Base \(r=1\)）、无保护的多词元基线（Base \(r=2\)）、MRP 直接外推模式（Direct）以及无损投机解码模式（Speculative）。测试基准涵盖数学推理（GSM8K、MATH500）与代码生成（HumanEval、MBPP）。
| 模型规模 | 解码配置 | GSM8K Acc. | GSM8K 相对吞吐 | HumanEval Acc. | HumanEval 相对吞吐 | MATH500 Acc. | MATH500 相对吞吐 |
|---|---|---|---|---|---|---|---|
| SDAR-1.7B | Base (\(r=1\)) | 77.6% | 1.00× (339 TPS) | 53.7% | 1.00× (314 TPS) | 55.0% | 1.00× (358 TPS) |
| Base (\(r=2\)) | 70.8% | 1.79× | 29.9% | 1.70× | 45.8% | 1.79× | |
| Direct (\(K=1\)) | 76.7% | 1.50× | 48.2% | 1.45× | 54.2% | 1.45× | |
| Direct (\(K=2\)) | 73.9% | 1.71× | 41.5% | 1.53× | 50.6% | 1.65× | |
| Spec (\(K=3\)) | 78.3% | 1.25× | 55.5% | 1.23× | 58.4% | 1.23× | |
| SDAR-4B | Base (\(r=1\)) | 90.3% | 1.00× (188 TPS) | 67.1% | 1.00× (170 TPS) | 71.4% | 1.00× (204 TPS) |
| Base (\(r=2\)) | 86.4% | 1.84× | 45.7% | 1.74× | 64.0% | 1.81× | |
| Direct (\(K=1\)) | 88.6% | 1.56× | 62.8% | 1.55× | 64.4% | 1.54× | |
| Direct (\(K=2\)) | 87.6% | 1.84× | 53.0% | 1.79× | 60.4% | 1.80× | |
| Spec (\(K=3\)) | 90.0% | 1.36× | 67.7% | 1.35× | 68.0% | 1.26× | |
| SDAR-8B | Base (\(r=1\)) | 90.9% | 1.00× (126 TPS) | 73.8% | 1.00× (114 TPS) | 72.2% | 1.00× (132 TPS) |
| Base (\(r=2\)) | 87.7% | 1.80× | 42.1% | 1.74× | 68.4% | 1.81× | |
| Direct (\(K=1\)) | 90.1% | 1.59× | 67.1% | 1.53× | 71.4% | 1.61× | |
| Direct (\(K=2\)) | 89.2% | 1.89× | 64.0% | 1.78× | 70.8% | 1.91× | |
| Spec (\(K=3\)) | 90.4% | 1.40× | 72.6% | 1.34× | 74.8% | 1.39× |
在动态去噪机制下，针对低阈值激进解码，MRP Remask 展现出了出色的精度挽救能力：
| 模型规模 | 阈值 \(\tau\) | GSM8K (Dyn.) | GSM8K (Remask) | 增益 | HumanEval (Dyn.) | HumanEval (Remask) | 增益 |
|---|---|---|---|---|---|---|---|
| SDAR-1.7B | 0.5 | 41.6% | 59.1% | +17.5% | 17.7% | 28.7% | +11.0% |
| 0.7 | 65.4% | 71.8% | +6.4% | 40.9% | 45.1% | +4.2% | |
| 0.9 | 76.2% | 77.3% | +1.1% | 49.4% | 52.4% | +3.0% | |
| SDAR-4B | 0.5 | 63.4% | 81.1% | +17.7% | 32.3% | 53.1% | +20.8% |
| 0.7 | 84.6% | 88.5% | +3.9% | 60.4% | 62.2% | +1.8% | |
| 0.9 | 88.5% | 90.1% | +1.6% | 67.1% | 65.9% | -1.2% | |
| SDAR-8B | 0.5 | 67.9% | 82.3% | +14.4% | 32.3% | 54.9% | +22.6% |
| 0.7 | 85.9% | 89.0% | +3.1% | 64.6% | 72.6% | +8.0% | |
| 0.9 | 90.8% | 91.4% | +0.6% | 75.0% | 75.0% | 0.0% |
消融实验¶
残差建模（Residual Modeling） vs 直接全量蒸馏（Direct Modeling）
在相同模型容量（3层 Transformer）与相同训练集下，对比在不同自回归展开步数 \(K \in \{1,2,3,4\}\) 下 GSM8K 的零样本推理精度：
| 骨干规模 | 建模方式 | \(K=1\) | \(K=2\) | \(K=3\) | \(K=4\) | 关键分析 |
|---|---|---|---|---|---|---|
| SDAR-1.7B | Direct Modeling | 76.5% | 73.1% | 62.4% | 48.0% | 全量蒸馏误差随步数快速累积崩塌 |
| Residual Modeling (MRP) | 76.7% | 73.9% | 66.7% | 60.9% | 残差建模外推鲁棒性显著（\(K=4\) 领先 +12.9%） | |
| SDAR-4B | Direct Modeling | 84.8% | 19.6% | 5.9% | 1.9% | 深步数下全量分布彻底发散崩溃 |
| Residual Modeling (MRP) | 88.6% | 87.6% | 70.9% | 57.2% | \(K=2\) 保持 87.6%，相对领先 +68.0% | |
| SDAR-8B | Direct Modeling | 90.5% | 72.0% | 34.2% | 6.7% | 无法适应多步自回归递归生成 |
| Residual Modeling (MRP) | 90.1% | 89.2% | 84.8% | 79.8% | \(K=4\) 仍维持 79.8%，领先 +73.1% |
关键发现¶
- 残差建模是防止自回归误差爆炸的核心基石：消融实验直接印证了本文的理论推演——直接预测下一步的完整词表分布会使预测误差随着展开步数 \(K\) 呈指数级累积；而残差预测因目标动态范围极小，即使在 \(K=4\) 步时仍能保持高精度（8B 模型下 79.8% vs 6.7%）。
- 投机验证成本在大模型上高度摊薄：投机解码在 8B 骨干上实现了最显著的端到端吞吐提升（GSM8K 上达到 1.40×），而在 1.7B 骨干上为 1.25×。这是因为轻量 MRP 头在整体前向中的耗时占比随骨干规模增大而显著缩小，且草稿的高接受率有效抵消了验证步的开销。
- 重掩码增益与解码激进程度严格呈反比：在 \(\tau=0.5\) 的超激进解码配置下，MRP Remask 带来的精度回升极其巨大（8B HumanEval 提升 22.6 点）；而当 \(\tau=0.9\) 时骨干网络本就保守输出，MRP 增益收敛至约 0 点，证明其作为自适应校准器的针对性极强。
亮点与洞察¶
- 将扩散步间相似性具象为残差学习：巧妙指出相邻去噪步骤输出分布具有高相似性与低变差特征，将复杂的多词元联合分布建模降维为学习微小的步间扰动残差，用极轻的计算代价突破了非自回归并行去噪的依赖屏障。
- 单一轻量模块适配两种截然不同的解码范式：同一个 3 层 MRP 头，在静态低吞吐模式下充当投机解码的“提速前推器”，在动态高吞吐模式下变身重掩码的“防错刹车片”，展示了高度统一的系统工程美感。
- 对代码与推理任务特性的敏锐洞察：实验发现重掩码在自然语言与数学推理上表现出单调正收益，但在保守阈值下的代码任务中偶尔出现轻微回退。论文指出代码符号对步间扰动极其敏感，展现了严谨的经验自洽性。
局限与展望¶
- 固定展开步数缺乏动态适配性：当前 MRP 采用固定的残差预测步数 \(K\)。对于容易预测的平凡词元或难以把握的复杂拐点，缺乏依据当前去噪难度动态调节步数的自适应机制。
- 代码领域的残差动态建模存在挑战：代码语料中的结构缩进与变量名存在较为剧烈、离散的分布跳变，导致残差估计难度高于连贯自然语言，未来需结合语法敏感先验设计专用的残差预测头。
- 阈值规划与残差调度的协同设计空间：目前动态阈值策略与静态投机截然分开，未来若能将连续自适应阈值调度（如进度感知阈值）与残差草稿生成一体化协同优化，有望探索出更极致的帕累托边界。
相关工作与启发¶
- vs Fast-dLLM / WINO: Fast-dLLM 依赖单一静态置信度阈值并行解码且不可逆；WINO 引入了双向重掩码但需要调用昂贵的主干网络进行重复验证。MRP 则直接以极低开销的残差头替代了二次骨干前向，在保留甚至超越 WINO 撤回能力的同时大幅削减了延迟开销。
- vs SSD / Learn2PD: SSD 在扩散模型中探索自投机解码但依赖主干自身的层次树结构，Learn2PD 学习二值接受/拒绝判别器。MRP 创新性地采用残差参数化设计，理论与实验均证明残差学习比二值判别或完整分布蒸馏具备更强的高阶展开鲁棒性。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [首次揭示并理论证明扩散去噪相邻步的残差收缩性质，提出残差投机与重掩码统一框架]
- 实验充分度: ⭐⭐⭐⭐⭐ [跨 3 个尺寸骨干模型、4 大推理与代码基准，包含完整端到端 SGLang GPU 吞吐实测与深度消融]
- 写作质量: ⭐⭐⭐⭐⭐ [理论推导、经验测量与算法流程逻辑环环相扣，图表详实且论证清晰]
- 价值: ⭐⭐⭐⭐⭐ [已集成至主流推理引擎 SGLang，为扩散大语言模型实用化推理部署提供了通用即插即用的加速利器]