Command Palette
Search for a command to run...
ReFusion:一种具有并行自回归解码的扩散型大语言模型
ReFusion:一种具有并行自回归解码的扩散型大语言模型
Jia-Nan Li Jian Guan Wei Wu Chongxuan Li
Abstract
自回归模型(Autoregressive Models, ARMs)受限于缓慢的串行推理过程。尽管掩码扩散模型(Masked Diffusion Models, MDMs)提供了一种并行推理的替代方案,但其存在关键缺陷:由于无法使用键值(Key-Value, KV)缓存,导致计算开销显著增加;同时,由于在不可行的token组合空间中学习依赖关系,生成结果常出现不连贯问题。为解决上述局限,我们提出ReFusion——一种新型掩码扩散模型,通过将并行解码从token层面提升至更高层级的“槽位”(slot)层面,实现了性能与效率的双重突破。每个槽位为固定长度、连续的子序列。该方法采用迭代式“规划-填充”解码机制:首先通过基于扩散的规划步骤识别出一组弱相关槽位,随后由自回归填充步骤并行解码这些选定槽位。该槽位级设计在保持统一因果框架的同时,实现了完整的KV缓存复用,并将学习复杂度从难以处理的token组合空间显著降低至可管理的槽位级排列空间。在七个多样化基准上的大量实验表明,ReFusion不仅在性能上相较先前MDMs平均提升34%,推理速度平均提升超过18倍,同时显著缩小了与强自回归模型之间的性能差距,并保持了平均2.33倍的加速优势。
一句话总结
中国人民大学与蚂蚁集团的研究人员提出REFUSION,一种基于槽位的掩码扩散模型,通过扩散引导的"规划-填充"流程替代词元级解码。该模型并行生成固定长度子序列,实现完整的键值缓存复用并降低学习复杂度,有效克服了现有扩散模型效率低下与生成不连贯的问题。REFUSION在性能上提升34%,速度较现有扩散模型快18倍,同时比自回归模型快2.33倍。
核心贡献
- 自回归模型面临顺序推理速度慢的问题,而现有掩码扩散模型因键值缓存不兼容及学习不可行词元组合导致生成不连贯。REFUSION通过将并行解码从词元级提升至固定长度连续槽位,采用迭代式"规划-填充"流程解决上述问题。
- 创新的槽位设计在自回归填充过程中通过统一因果框架实现完整键值缓存复用,同时将学习复杂度从词元组合空间降至可控的槽位排列空间。该设计同步解决了键值缓存的架构瓶颈与掩码扩散模型的生成不连贯问题。
- 在七项多样化基准测试中,REFUSION较先前掩码扩散模型性能提升34%、速度提升18倍以上,在性能上逼近强自回归模型的同时保持2.33倍平均加速。
引言
掩码扩散模型(MDM)旨在通过并行词元解码加速文本生成,相比传统自回归模型(ARM)具有优势,但存在关键瓶颈:标准双向注意力机制阻碍推理过程中的高效键值缓存复用,削弱其速度优势。先前方法尝试部分解决方案——近似缓存复用、在固定块中混合因果/双向注意力或纯因果注意力——但这些方法要么牺牲全局生成灵活性,要么依赖不可靠的置信度启发式规则,或引入独立验证模型,或导致难以处理的训练复杂度。作者提出基于新型"槽位"框架的ReFusion,通过槽间并行规划与槽内自回归解码的融合设计,独特地实现了每个解码词元的完整键值缓存复用,同时保留灵活的任意顺序生成能力,消除了传统MDM架构固有的权衡。
方法
作者提出名为REFUSION的创新架构,将并行解码过程从词元级重构至更高层的槽位级,以克服传统掩码扩散模型(MDM)的效率与连贯性限制。核心创新在于将序列划分为固定长度、非重叠的槽位,并执行两步迭代解码循环:基于扩散的槽位规划与自回归槽位填充。该设计基于经验观察:词元间依赖性随距离快速衰减,允许槽内词元通过串行生成保持局部连贯性,而槽位本身可并行处理以提升计算效率。
推理过程始于完全掩码的响应序列,该序列被划分为 K 个连续槽位,每槽包含 k 个词元。每次迭代中,模型首先执行槽位规划:并行计算所有掩码槽位的草案,并为每个槽位分配确定性分数(通常基于槽内首词元概率)。选择超过预设阈值的一批槽位进行填充。此步骤识别出受上下文强约束但槽间弱依赖的槽位,使其适合并行处理。随后在第二步验证并完成所选草案槽位:模型首先尝试全局验证,将草案连接后检查条件概率超过词元级阈值的前缀;若该前缀覆盖一个或多个完整槽位,则整体接受。否则,模型回退至并行迭代完成流程,通过验证与预测循环独立完善每个槽位直至完成。关键创新在于:每次填充后,新生成槽位在物理顺序上被移至剩余掩码槽位之前,使所有已生成词元的键值缓存得以完全复用。此重排序通过为所有词元保持不变的绝对位置ID实现,结合RoPE机制使模型能正确计算相对距离,即使缓冲区动态重排。该过程持续至所有槽位填满,最终响应通过恢复原始槽位顺序构建。

训练流程精心设计以模拟推理动态。对每个训练样本,响应被划分为槽位。随机掩码部分槽位,未掩码(干净)槽位经置换以模拟推理中遇到的非顺序生成顺序。最终训练输入通过连接置换后的干净槽位与掩码槽位构建。模型采用混合目标函数联合训练规划与填充能力:干净槽位使用标准自回归损失(模型基于槽内前置词元预测当前词元),掩码槽位使用去噪损失(模型基于干净上下文及部分可见掩码词元重建原始词元)。该双重目标确保序列中每个词元提供直接监督,显著提升数据效率(传统MDM仅从掩码位置学习)。最终损失为ARM与MDM组件的加权和,模型从现成自回归骨干网络初始化。

模型架构基于标准因果Transformer构建,关键改进在于支持显式非连续位置ID列表,使模型能正确关注缓冲区中所有逻辑前置词元。槽位设计本质上将学习复杂度从不可行的词元组合空间降至可控的槽位排列空间。作者证明模型性能对块大小变化具有鲁棒性,在32-128词元块大小区间表现最佳,吞吐量与准确率均持续超越强自回归基线。

作者还指出,模型设计兼容半自回归重掩码:大块序列按序解码,槽级算法在块内应用。这种灵活性使模型能适应不同计算约束,同时保持核心优势.
实验
- 在七项基准测试中,REFUSION以最高34%性能提升与18倍加速超越所有掩码扩散模型(MDM)基线(如HumanEval上pass@1达78.66%),并在GSM8K和MBPP上以2.33倍平均加速超越Qwen3-8B自回归模型3.68分。
- 使用Qwen3-8B骨干网络的受控对比中,REFUSION在HumanEval上以1.9倍加速领先重训练基线约16分,证实架构优势无数据依赖。
- 对比Dream-7B-Instruct,REFUSION在跳过预训练情况下仍实现2.23%平均性能提升与11.05倍加速,在GSM8K、HumanEval和MBPP上表现优异。
- 键值缓存复用消融实验显示,基准测试中吞吐量提升1.16–1.33倍且精度稳定或微升,验证效率提升无质量折衷。
作者在Dream-7B-Instruct骨干网络上重训练REFUSION以进行受控对比,结果表明:尽管缺乏预训练且样本更少,REFUSION仍实现2.23%平均性能提升与11.05倍加速。在GSM8K、HumanEval和MBPP等推理与编码任务上显著超越Dream-7B-Instruct,但在MMLU-Pro和ARC-C等知识密集型基准上表现较弱,符合其缺乏知识注入预训练的特性。

作者基于相同Qwen3-8B检查点在120K数据集上重训练REFUSION与基线模型以隔离架构影响。结果显示:REFUSION在平均性能与速度上显著超越LLaDA和BD3-LMs(平均分53.34 vs 14.43与12.11),且在HumanEval上比重训练Qwen3-8B基线高约16分。这证实REFUSION的设计优势使其在数据受限条件下仍具备卓越学习效率。

REFUSION在七项基准测试中全面超越所有对比掩码扩散模型与强自回归基线,性能与吞吐量均获最高平均分。结果表明:REFUSION以2.33倍加速超越Qwen3-8B,同时在GSM8K和MBPP等关键任务上表现更优,证明其槽位并行解码打破了传统速度-质量权衡。该设计通过完整键值缓存复用与学习复杂度降低,持续大幅超越先前MDM(平均性能提升34%、加速18倍以上)。

REFUSION作为槽位掩码扩散模型,在保持高吞吐量的同时实现与自回归模型相当的性能。结果表明:尽管未使用系统优化,REFUSION在HumanEval和MBPP分数上匹配或超越多个自回归基线,且在精度与速度上均优于其他掩码扩散模型。

作者对比REFUSION启用/禁用键值缓存重计算的效果,显示默认缓存复用方法在基准测试中加速1.16–1.33倍且性能稳定或微升。结果表明:并行解码中避免完整上下文化不会降低精度,甚至可能提供隐式正则化收益。
