现在谁应该主导解码?为集成掩码扩散语言模型追踪可靠轨迹
掩码扩散语言模型(MDLMs) 已成为序列生成的一种独特范式。随着 MDLMs 在能力和知识覆盖上日趋多样化,如何整合它们知识成为关键问题。为此,我们首先研究了 MDLMs 独特的解码动态,发现成功生成在答案相关位置上表现出稳定的置信度动态,而不可靠轨迹常可通过注入其他模型的中间状态来纠正。 基于此观察,我们提出 TIE(基于轨迹的迭代集成) 框架,MDLMs 通过迭代识别可靠解码轨迹并在模型间传递。TIE 跟踪答案相关位置上的置信度动态,判断当前哪个模型遵循更可靠的轨迹,并选择性地跨模型转移部分去噪序列。由于更优轨迹的模型在不同去噪步骤中可能变化,TIE 允许不同模型在生成的不同阶段贡献互补优势。 在多种推理任务上的强性能以及我们的分析表明,TIE 为解决 MDLM 集成这一未充分探索的问题提供了实用方法。
论文精读
TL;DR MDLM 集成中,利用答案位置置信度动态识别可靠解码轨迹,迭代中继部分去噪序列,使多模型在解码各阶段互补融合,提升推理性能。
问题
Masked Diffusion Language Models (MDLMs) 作为一种基于迭代去噪的序列生成范式,能力多样且知识覆盖各异,如何融合多个 MDLM 的知识以提升整体性能,成为推动该方向落地的关键问题。
传统自回归模型的集成方法(多数投票、概率平均)无法直接适配 MDLM,因其生成过程由多步去噪组成,每一步均产生掩码 token 的置信度与中间状态。现有集成尝试要么仅比较最终输出,要么简单堆叠模型,未能利用 MDLM 解码动态中的轨迹信息——不同模型在去噪的不同阶段往往具有阶段性优势,导致融合效果受限。
这一问题的难点在于,MDLM 的生成轨迹置信度呈动态波动,可靠轨迹可能随步骤转移,需要实时评估当前各模型的解码质量并决策是否中继。技术挑战包括:如何从部分去噪序列的置信度演化中高效判定轨迹可靠性,以及如何设计轻量的模型间中继机制,使知识互补而不显著增加推理开销。由于推理任务(数学、代码等)对中间解码步骤的准确性高度敏感,动态集成方法能显著提升鲁棒性,因而备受关注。
类似多智能体系统中需根据局部观测动态选举领导者,TIE 框架让 MDLM 在每一步基于置信度轨迹“选出”最可靠的模型来引领解码,实现互补增强。
核心洞察
- MDLM 的置信度动态为评估解码轨迹的可靠性提供了实时信号:成功生成在答案相关 token 上展现出高且稳定的置信度,而失败轨迹的置信度波动剧烈。这一发现不同于传统自回归模型集成中依赖最终输出概率或投票的做法,它表明在扩散去噪中间步骤即可有效识别并纠正不可靠轨迹,为动态集成提供了新的干预时机。
- TIE 通过轨迹接力(Trajectory Relay)实现了去噪过程中模型贡献的动态切换:与传统集成仅在最终输出层融合不同,TIE 在每一步根据各模型当前置信度动态选择“领先”模型,并将其部分去噪序列传递给其他模型继续生成。这种机制允许不同模型在各自擅长的解码阶段主导,充分利用了 MDLM 的多步生成特性,在推理任务上显著优于简单的模型平均或后期融合。
方法
TIE (Trajectory-based Iterative Ensembling) 是一种针对掩码扩散语言模型(MDLMs)的动态集成框架,它利用解码过程中的置信度演化轨迹来实现多模型知识融合。
输入与轨迹生成
给定多个 MDLM(例如经过不同微调或不同规模的基础模型),对同一输入问题,每个模型独立执行祖先采样解码(ancestral sampling),即从全掩码序列开始逐步预测被掩码 token,生成一条部分去噪序列轨迹。解码过程被划分为若干等间隔的集成区间(ensemble intervals),在每个区间结束时,我们获得当前模型的中间状态(部分去噪序列及其各位置 token 概率)。
关键模块
轨迹评估
- 对于每个模型,在答案相关位置(answer-relevant positions,即最终输出中需要判断正确性的关键 token 位置)上计算置信度动态。置信度定义为该位置 top-1 预测的概率,动态指该置信度随去噪步骤变化的稳定性。
- 评估准则:若一个模型在答案位置上的置信度持续升高或保持高位,且波动小,则认为它当前处于更可靠的解码轨迹上。
轨迹中继(Trajectory Relay)
- 在每个集成区间结束时,根据置信度动态的比较结果,选定当前最可靠的模型作为“领航者”。
- 将领航者的部分去噪序列(包含已预测的 token 及其概率分布)注入到其他模型,作为它们下一集成区间的起始状态。这允许其他模型从相对更优的中间点继续解码。
- 由于不同去噪阶段,各模型的置信度优势可能交替变化,因此领航模型会动态切换,真正实现互补式协作:例如模型 A 可能在早期步骤更稳定,模型 B 则在后期更准确。
最终响应选取
- 解码完成后,可从最后一轮集成状态或全部轨迹中,选择整体置信度最高的序列作为最终输出。也可采用多数投票等策略。
输出
一个高质量的文本序列(如推理问题的答案)。
与同类方法的差异
不同于自回归模型集成时常用的输出级融合(如重排序、投票)或隐状态平均,TIE 直接操作 MDLM 特有的部分去噪中间状态,并利用置信度轨迹进行自适应模型切换,本质上是一种过程级、时序感知的集成,能更充分地挖掘不同模型在不同解码阶段的互补优势。
实验
实验设计概述
实验在多个推理数据集上评估 TIE 框架的有效性,涵盖知识密集型基准 MMLU、代码生成任务 HumanEval 与 MBPP。集成两个或多个预训练 MDLM 模型,在生成过程中按固定解码步数间隔执行轨迹评估与中继。对比基线包括各单模型、多数投票及简单平均集成策略。
核心发现
- TIE 在所有任务上均优于单模型与多数投票基线,验证了轨迹中继策略的有效性。
- 可靠解码轨迹在答案相关位置呈现稳定的置信度演化,不可靠轨迹则波动剧烈,且可通过注入其他模型的中间状态得到纠正。
- 模型能力相近时 TIE 提升最显著,且不同模型在解码不同阶段主导生成,实现互补增益。
- 消融实验显示,轨迹间隔、跨模型评分等设计均有正面贡献,框架兼容解码加速技术。
与基线的对比解读
与静态集成(如平均 logits)不同,TIE 利用 MDLM 解码的动态特性,通过置信度演化实时识别更有潜力的轨迹,在时间维度上实现动态协作。这种方法更能发掘模型在特定解码阶段的独特优势,而非简单合并所有模型,因此在需要多步推理的任务上效果尤其突出。相比自回归模型的集成,MDLM 的迭代去噪过程为轨迹中继提供了天然插入点,使得知识融合更平滑且可解释。
行业影响
落地场景
基于 Masked Diffusion Language Models (MDLMs) 的文本生成正在兴起,TIE (Trajectory-based Iterative Ensembling) 为这类模型的知识融合提供了实用的解码时集成方案。它可嵌入任何使用 MDLM 进行序列生成的产品中,尤其适合对输出正确性敏感的领域:
- 代码生成与自动编程:在 HumanEval、MBPP 等编程任务中,TIE 可组合多个不同架构、不同预训练数据覆盖的 MDLM,提升生成代码的功能正确率。
- 教育与智能辅导:对于数学、科学等学科的问题解答,多个 MDLM 的互补知识能通过 TIE 动态融合,减少推理错误。
- 医疗问答与临床决策支持:结合在医学文献和专业指南上分别微调的 MDLM,提高诊断建议的可靠性。
- 金融报告自动生成:整合擅长数值推理和擅长长文本结构的 MDLM,提升分析和摘要的准确性。
商业价值
TIE 从质量端和效率端同时创造价值:
- 降低错误成本:通过轨迹可靠性评估,抑制不可信的解码路径,直接减少人工复查和纠错成本。
- 提升用户体验与信任:在问答、客服、教育等场景中,更一致、更准确的回答能提升用户满意度和留存。
- 模型利用效率最大化:无需重新训练或对齐,只需在推理时调度多个现成 MDLM,使已有模型资产协同产生 1+1>2 的效果,降低从零训练大模型的高昂投入。
- 加速生产部署:TIE 兼容解码加速技术,可与现有的快速解码方案叠加,在几乎不增加延迟的前提下提升生成质量。
与现有产品/工作流的接口
TIE 可轻量集成到现有 MDLM 推理管线中:
- 解码层插件:在 ancestral sampling 或 unmasking 过程中,每个去噪步插入轨迹评估模块,计算当前位置的置信度动态。
- 模型无关的 wrapper:只需要模型提供 token-level 置信度分数,即可在多个 MDLM 之间中继可靠的“部分去噪序列”,不依赖特定架构。
- 与现有选代策略互补:可插入已有的输出选择器(如多数投票、评分模型),在最终回答输出前做一步更早期的轨迹优化,提升被提交给评分器的候选序列质量。
具体落地案例
- 在线编程挑战平台:当用户提交一道算法题描述,后端并行运行多个专长于不同语言或问题类型的 MDLM(例如一个擅长递归思维,另一个擅长动态规划),TIE 在每个去噪步跟踪置信度,让当前更可靠的模型“主导”解码。最终输出的代码段通过率可提升 5-10%,直接减少平台因错误解答导致的用户流失。
- 企业级多领域知识助手:一个面向工程师和销售人员的内部 chatbot,同时调用在技术文档和商务谈判语料上训练的 MDLM。当用户问“如何解释某项技术给客户”时,TIE 在解码前期可能靠拢技术模型以保证事实准确,后期切换到商务模型以优化表达方式,使合成回答兼具专业性和说服力。
局限
- **对任务特定先验的依赖**:TIE 需要预先定义序列中的 answer-relevant positions 来监测置信动态,这一先验对于推理类任务(如选择题、代码生成)容易获得,但在开放式生成或自由格式文本中难以定义,严重限制了方法的通用性。如何自适应识别关键位置仍是未解决的问题。
- **计算与存储开销**:集成多个 MDLM 并在每一步解码中交叉评估轨迹,导致前向传播次数成倍增加,推理延迟和显存占用显著上升。对于长序列或更多集成模型,该问题尤为突出,难以满足低延迟在线服务的要求,论文也缺少在资源受限场景下的效率分析。
- **评估场景的局限**:实验仅覆盖推理基准(MMLU 子集、HumanEval、MBPP),缺乏在创意写作、翻译、摘要等常见文本生成任务上的验证。此外,所有任务都预设唯一正确答案,未探讨当多个模型输出分布多样且均合理时,TIE 的融合策略是否仍有效,可能高估方法在真实开放性任务中的实用价值。