扩散语言模型新突破:频率感知重掩码提升推理能力

两阶段框架结合频率感知重掩码,动态分配推理计算以提升扩散语言模型性能。
该论文提出将扩散语言模型的测试时扩展形式化为推理计算量和解码策略的自适应分配,并设计两阶段框架:先用快速基础采样器解决简单问题,再将困难问题交给频率感知重掩码策略。频率感知重掩码追踪token预测在去噪步骤中的变化频率,优先提交频繁变化的token以鼓励探索替代推理路径,避免局部置信度信号的误导。在LLaDA-8B-Instruct和Dream-7B-Instruct上,该方法在MATH-500上准确率最高提升3.73和7.33个百分点,在HumanEval上覆盖率最高提升11.39和10.98个百分点,并能恢复基线100次采样后仍未解决的3.75%至18.75%的案例。
扩散语言模型(Diffusion Language Models, DLMs)通过迭代去噪生成文本,其中哪些token被提交或重新掩码的决策直接影响输出质量。现有的确定性感知策略依赖瞬时置信度或熵,但这些局部信号在困难推理问题上可能产生误导。最近一篇论文提出将DLM的测试时扩展(test-time scaling)形式化为推理计算量和解码策略的自适应分配,并引入频率感知重掩码(frequency-aware remasking)策略,在多个基准上取得了显著提升。
两阶段自适应分配框架
该研究设计了一个两阶段方法:首先用一个快速基础采样器解决简单问题,然后将剩余困难问题交给频率感知重掩码策略处理。这种分配方式的核心洞察是:不同问题所需的计算资源差异巨大,简单问题无需过度消耗采样预算,而困难问题则需要更精细的解码控制。

与现有方法相比,该框架不是对所有问题一视同仁,而是根据问题难度动态调整策略。基础采样器快速过滤掉容易的案例,让频率感知重掩码能够集中资源攻克真正困难的推理任务。
测试时扩展(test-time scaling):指在推理阶段通过增加计算量(如更多采样次数、更长去噪步数或更复杂的解码策略)来提升模型性能,而不改变模型参数。在自回归语言模型中,常见形式包括链式思维提示、多数投票和树搜索等。对于扩散语言模型,测试时扩展还涉及解码策略的动态调整,例如自适应分配去噪步数或重掩码频率。本文的两阶段框架正是将测试时计算资源根据问题难度进行路由,属于推理计算自适应分配的典型实践。
频率感知重掩码的原理
传统确定性感知策略在每一步去噪时,根据当前token的置信度或熵决定是否提交。然而,这种局部信号在复杂推理中并不可靠——模型可能在某个步骤表现出高置信度,但最终答案仍是错误的。
频率感知重掩码转而追踪每个token的预测在去噪步骤中变化的频率。具体来说,它优先提交那些频繁变化的token,从而鼓励样本探索替代推理路径,而不是坍缩到模型偏好的轨迹上。这一机制背后的直觉是:如果一个token的预测反复摇摆,说明模型在该位置存在不确定性,强制提交可能锁定错误;相反,允许这些token继续变化,可以给模型更多机会找到正确路径。
重掩码(remasking)与确定性感知(deterministic perception):在扩散语言模型的迭代去噪过程中,每一步会生成所有未提交token的预测分布。重掩码指在下一步去噪前,将某些已提交的token重新标记为未确定状态(即重新掩码),使其有机会被重新预测。确定性感知策略通常基于当前步的置信度或熵,决定哪些token应当提交(固定)或重掩码。局部信号虽计算高效,但在多步推理中容易因模型过度自信而锁定错误。
频率感知信号与探索-利用权衡:频率感知重掩码不依赖单步置信度,而是统计每个token在多个去噪步骤中预测类别发生变化的次数(或频率)。高频变化意味着模型在该位置持续不确定,可能对应多条候选推理路径。优先提交高频token,相当于在解码早期强制模型“选择”一条路径,从而避免后期坍缩到单一轨迹;而保留低频token则允许稳定信息逐步固化。这种机制实质上是将探索(保留不确定token以尝试替代路径)与利用(提交稳定token以推进生成)进行动态平衡,与强化学习中的探索策略有相似直觉。
实验表现:准确率与覆盖率双提升
研究者在LLaDA-8B-Instruct和Dream-7B-Instruct两个模型上进行了实验。与强自适应分配基线相比,该方法在MATH-500上取得了最高3.73个百分点(LLaDA)和7.33个百分点(Dream)的绝对准确率提升,在HumanEval上取得了最高11.39个百分点(LLaDA)和10.98个百分点(Dream)的绝对覆盖率提升。
更有意思的是,对于基础确定性感知策略在100次采样后仍未解决的问题,频率感知重掩码能够恢复其中3.75%到18.75%的案例。这表明该方法不仅提升了平均性能,还能攻克基线方法完全无法解决的难题。
对DLM推理优化的启示
这项工作的意义在于揭示了测试时计算分配的重要性。随着扩散语言模型规模增大,如何高效利用推理计算成为关键问题。频率感知重掩码提供了一种轻量级的解码策略调整,无需重新训练模型即可提升推理能力。
此外,该方法与现有的确定性感知策略互补,可以结合使用。两阶段框架的设计也提醒我们,在实际部署中,根据问题难度动态路由不同解码策略可能比单一策略更有效。未来,这种频率感知的思路或许可以扩展到其他生成模型或更复杂的推理任务中。
相关推荐

Hermes Agent零基础教程:自进化AI智能体实战
Hermes Agent零基础教程:详解这个GitHub两月获41.7K星的开源自进化AI智能体,从Harness Engineering核心理念、学习循环、三层记忆系统到WSL安装与实战配置,小白也能轻松上手。

AI Agent落地实战:用WorkBuddy打造智能办公助理
本文通过实战案例演示如何用腾讯WorkBuddy完成文件批量转换与归档,对比聊天大模型与AI Agent的能力差异,分析WorkBuddy的积分体系、适用场景与产品定位,为办公自动化选型提供参考。

从单卡瓶颈到千亿参数:大模型分布式并行6大关键问题
基于B站教程梳理大模型分布式并行核心知识:从decoder-only架构的逐token生成与X乘W计算本质出发,解析单卡瓶颈下数据并行、张量并行、流水线并行等6个关键问题,适合私有化部署与模型训练场景的学习者建立系统认知。