Masked Diffusion Language Agents for Tool-Integrated Chemical Reasoning¶

会议: NeurIPS 2026
arXiv: 2610.13330
代码: https://github.com/bilemond/ChemDiffAgent
领域: LLM Agent
关键词: 化学推理智能体、掩码扩散语言模型、工具调用、双向建模、块级扩散

一句话总结¶

本文提出首个面向工具整合化学推理的掩码扩散语言智能体 ChemDiffAgent，将化学交织推理与工具调用建模为块级迭代去噪过程，通过两阶段后训练（智能体 SFT 与低方差偏好优化 VRPO）克服自回归模型的单向因果瓶颈，在单轮与多轮化学基准上较同规模自回归模型取得约 20% 的性能跃升。

研究背景与动机¶

将大语言模型与外部专业化学工具（如分子性质预测器、反应引擎及化学信息学库）结合构建化学智能体，已成为解决复杂分子合成与性质推断的重要路径。然而，现有的化学智能体全部基于自回归（AR）架构构建。自回归模型依赖因果注意力与严格的从左到右单向生成机制，其基本假设与复杂化学推理任务之间存在内在矛盾。首先，SMILES 字符串等化学领域特定模态本质上具有强双向依赖性，原子、分支结构与环闭合标记在空间和逻辑上前后相互约束，单向解码容易在未掌握全局结构前过早固化局部标记；其次，化学工具调用是全局耦合的决策链，工具的选择、参数的合法性以及调用背后的推理意图紧密互依；第三，多步化学任务往往需要将来自异构工具的多源返回证据整合成自洽的结论，自回归的局部逐步展开极易导致次优的局部工具选择偏见与错误级联。

这种局限性的核心矛盾在于：有效利用工具要求智能体将其工具调用决策边界与内部知识边界精准对齐——即在自身参数化知识不确定时调用外部计算，在参数化知识足够时自主推理。而自回归模型的单向标记承诺机制剥夺了模型在生成完整动作前进行全局协调的能力。一旦模型输出了某个工具名称或前缀参数，后续生成就必须无条件受制于该早期承诺，极易掉入幻觉陷阱或参数格式崩溃。

本文的切入角度是打破单向因果生成枷锁，引入块级扩散语言模型（BDLM）。在每个生成块内部，模型保持全双向注意力，并将中间思考、工具触发决策、工具名、参数及多源证据整合转化为多步迭代去噪轨迹。核心 idea：将交织推理与工具调用的轨迹重构为块级掩码扩散去噪过程，利用块内任意序解码与全局双向上下文实现工具调用的联合规划与后验对齐，并设计智能体专属的 SFT 与方差缩减偏好优化两阶段后训练管线。

方法详解¶

整体框架¶

ChemDiffAgent 的核心目标是将化学交织推理与外部工具调用统一建模为条件掩码扩散过程。智能体的交互轨迹表示为交织序列 \(\mathcal{H}=(r_1, c_1, o_1, \ldots, r_T, c_T, o_T, a)\)，其中 \(r_t\) 为中间思考推理片段，\(c_t=(\tau_t, \alpha_t)\) 表示调用工具 \(\tau_t\) 并传入参数 \(\alpha_t\)，\(o_t=\tau_t(\alpha_t)\) 为外部工具环境返回的确定性观察结果，最后 \(a\) 为最终答案。整个模型建立在块级扩散架构（BDLM）之上，块间保持因果关系以天然支持 KV Cache 和工具环境的真实物理交互，块内则维持全双向注意力以支持并行与任意序迭代去噪。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入查询与化学上下文<br/>(Query Q)"] --> B["块级因果注意力与上下文拼接<br/>[Q, H_teacher, H_noisy]"]
    B --> C["智能体掩码加噪机制<br/>(仅扰动思考与工具调用，冻结工具返回观察)"]
    C --> D["智能体 ELBO 目标监督优化<br/>(Agentic ELBO SFT)"]
    D --> E["低方差偏好优化<br/>(Agentic VRPO 对比成败轨迹)"]
    E --> F["测试时任意序去噪解码<br/>(优先固化动作脚手架，延迟确定化学参数)"]
    F --> G["外部化学工具环境协同执行<br/>(属性预测 / 反应检索 / 结构规范化)"]
    G --> H["多源证据插槽级融合与最终答案输出"]

ChemDiffAgent 的方法管线涵盖理论支撑、智能体监督微调（Agentic SFT）以及智能体偏好对齐（Agentic VRPO）三个核心支柱。

关键设计¶

1. 任意序解码的自适应路径与非均匀风险决策理论：证明非自回归去噪对工具调用的数学优势

论文从理论上严格论证了扩散模型相比固定自回归顺序在智能体推理中的两大优势。首先是在子问题条件难度不均情况下的自适应路径优势：对于块内标记序列 \(U=(U_1,\ldots,U_B)\) 及上下文 \(C\)，自回归只能按固定排列 \(\pi_{\mathrm{AR}}=(1,\ldots,B)\) 生成，而扩散模型容许在所有合法排列空间 \(\Pi_{\mathrm{flex}}\) 中搜索。因模型累积逼近误差取决于已解析位置集合 \(S_t\)，其总负对数似然满足：

\[\inf_{\pi\in\Pi_{\mathrm{flex}}}\mathcal{R}_{\theta}(\pi)\leq\mathcal{R}_{\theta}(\pi_{\mathrm{AR}})\]

其次是在非均匀决策风险下的严格排序优势：在化学工具调用中，不同位置标记的下游错误代价具有不对称性。若 \(Z\) 代表高风险决策（如选定不可逆的反应分析工具），\(R\) 代表同块内提供信息的低风险辅助决策（如规范化的结构标识符），其错误代价权重满足 \(w_Z > w_R > 0\)。自回归顺序强制先决断 \(Z\) 再解析 \(R\)，而逆向或自适应路径可先解析低风险的 \(R\) 再填补 \(H(Z|C,R)\)。两者风险加权对数损失的差值严格满足：

\[\mathcal{J}_{\mathrm{AR}} - \mathcal{J}_{R\rightarrow Z} = (w_Z - w_R)I(Z;R\mid C)\]

只要两者在给定上下文下具有互信息 \(I(Z;R|C)>0\)，灵活解码顺序即可严格降低加权决策风险。这从数学上证实：扩散智能体能够先建立稳定的全局工具动作脚手架与语法槽位，再依据充分的上下文回填精细且易错的化学 SMILES 参数，有效避免过早确定不可逆参数。

2. 智能体掩码加噪与 Agentic ELBO 监督微调：对齐环境观察与自主动作边界

现成的扩散语言模型无法直接用于交互式智能体，因为外部工具返回的真实观察数据由环境提供，不能作为模型的生成目标。为此，本文设计了智能体掩码加噪（Agentic Noising）与专属的 Agentic ELBO 训练机制。在训练中，输入由查询 \(Q\)、无噪教师轨迹 \(\mathcal{H}_{\mathrm{teacher}}\) 与加噪轨迹 \(\mathcal{H}_{\mathrm{noisy}}\) 拼接组成单次前向序列。Agentic Noising 只对智能体自主决策的区域（即 <think> 推理思考块和 <tool_call> 动作块）施加独立掩码扰动，而对工具返回内容 <tool_response> 保持完全冻结且不被破坏。

为了彻底消除块内双向注意力可能带来的训练与推理数据泄漏，若某个块内包含跨越至工具返回的信息，该部分在块内被完全遮蔽。形式上，设 \([R]\) 表示工具响应的标记区间，Agentic ELBO \(\hat{\mathcal{L}}^{block}_{\theta}\) 定义为仅在被掩码且非工具返回位置上计算的对数似然下界：

\[\hat{\mathcal{L}}^{block}_{\theta}(y\mid x)\triangleq\mathbb{E}_{t\sim\mathcal{U}[0,1],y_t\sim\hat{q}_t}\Bigg[\frac{1}{t}\sum_{k=1}^K\sum_{i=1}^B\mathbf{1}\Big[(y_t^{k,i}=[\mathrm{M}])\land(y^{k,i}\notin[R])\Big]\log p_{\theta}(y^{k,i}\mid y_t^k,y^{<k},x)\Bigg]\]

通过最小化 \(\mathcal{L}_{SFT}=\mathbb{E}[-\hat{\mathcal{L}}^{block}_{\theta}(y|x)]\)，模型精确学会了在给定部分上下文和外部环境真实反馈的情况下，联合反推中间思考与规范工具调用，从而实现了决策边界与真实知识边界的高度对齐。

3. 智能体低方差偏好优化：消除多轮交互长程决策漂移

单靠 SFT 能够教会模型标准的工具交互协议，但无法让模型在面对复杂长程多轮化学推理时，辨别不同解题路径的优劣。为此，本文提出了 Agentic VRPO。传统针对长序列扩散模型的强化学习通常面临 Monte Carlo 采样方差过大的问题。本文利用 Agentic Noising 采样 \(M\) 个受控扰动上下文 \(\xi_j=(t_j, y_{t_j})\)，定义仅评估智能体生成标记的重构得分 \(\ell_{\theta}(y\mid x;\xi_j)\)，并对比策略模型与固定参考模型 \(\ell_{\mathrm{ref}}\) 的期望差异：

\[\widehat{\Delta}_M(y\mid x)=\frac{1}{M}\sum_{j=1}^M\left[\ell_{\theta}(y\mid x;\xi_j)-\ell_{\mathrm{ref}}(y\mid x;\xi_j)\right]\]

针对同一化学问题配对的成功轨迹 \(y_w\) 与因工具选择错误或参数幻觉失败的轨迹 \(y_l\)，通过 Bradley-Terry 偏好模型优化目标：

\[\mathcal{L}_{\mathrm{VRPO}}(\theta)=-\mathbb{E}_{(x,y_w,y_l)}\left[\log\sigma\left(\beta\left[\widehat{\Delta}_M(y_w\mid x)-\widehat{\Delta}_M(y_l\mid x)\right]\right)\right]\]

共享掩码上下文的差分估计极大地压低了方差，强化了智能体在关键多步决策点上的自我修正与长期规划能力。

损失函数 / 训练策略¶

训练分为两阶段： - 阶段一（Agentic SFT）：基座采用基于 Qwen3-8B 持续预训练得到的块级扩散模型 SDAR-8B。采用优质教师模型通过工具子集拒绝采样构建高质量化学轨迹（保证最终答案正确、推理完备、工具参数格式绝对合规）。在 8 张 NVIDIA A100 GPU 上采用 AdamW 优化器，学习率设为 \(1\times 10^{-5}\)，最大提示长度 4096，最大生成长度 8192，按 Agentic ELBO 损失进行全监督拟合。 - 阶段二（Agentic VRPO）：由 SFT 检查点针对多轮化学任务生成采样子轨迹，自动比对执行结果构建正负样本偏好对 \((y_w, y_l)\)。设置采样数 \(M=2\)，温度缩放系数 \(\beta=0.1\)，学习率为 \(5\times 10^{-7}\)，序列总长上限为 6144，在极低学习率下完成长程决策对齐。 - 推理配置：固定块大小（Block size）为 64，块内去噪步数设定为 64，最大允许交互轮次为 7 轮。

实验关键数据¶

主实验¶

评估构建了覆盖分子组成推断（Comp.）、名称转换（Name Conv.）、性质预测（Prop. Pred.）、逆合成分析（Retro.）、分子优化（Mol. Opt.）以及正向合成预测（Fwd. Synth.）六大家族的单轮基准，以及涵盖简单（2 轮）、中等（3-5 轮）和困难（6-7 轮）三个难度梯度的多轮依赖图基准。

表 1：六大化学任务单轮成功率对比（%）

方法	模型规模 / 底座	组成推断	名称转换	性质预测	逆合成	分子优化	正向合成	总分
GLM-5.1 (ReAct)	商业 API	39.8	19.6	40.6	38.5	0.0	30.0	32.8
Gemini 2.5 Flash	商业 API	77.9	40.2	67.7	0.0	79.2	75.0	59.8
DeepSeek-V3.2	商业 API	83.5	71.6	90.6	50.8	88.7	78.3	77.0
GPT-5.5	商业 API	71.4	56.9	64.6	96.9	64.2	91.7	70.6
Claude Opus 4.7	商业 API	90.9	88.2	87.5	100.0	100.0	100.0	90.9
ChemCrow	Claude Opus 4.7	91.2	92.2	83.3	95.4	100.0	86.7	89.6
ChemCrow	DeepSeek-V3.2	55.0	53.2	53.0	72.6	43.1	28.3	72.4
S1-NexusAgent	GPT-5.5	71.4	78.8	77.5	71.8	38.5	33.3	76.7
Qwen3-8B + SFT	8B (AR)	76.6	53.9	87.5	78.5	74.5	63.3	69.1
SDAR-8B (Zero-shot)	8B (Diffusion)	49.4	1.0	17.7	7.7	11.3	5.0	23.3
ChemDiffAgent (SFT)	8B (Diffusion)	91.3	78.4	94.8	96.9	75.5	98.3	89.0
ChemDiffAgent (SFT+RL)	8B (Diffusion)	91.8	97.1	95.8	96.9	79.2	98.3	91.8

消融实验与多轮能力¶

多轮工具调用实验测试了智能体在图约束下的长程工具编排与推理能力，各模型在不同调用轮次难度下的成功率如下表所示。

表 2：多轮工具调用难度阶梯成功率（%）与轨迹效率对比

模型 / 架构	易 (2轮)	中 (3-5轮)	难 (6-7轮)	多轮总成功率	单轮平均生成长度	多轮平均生成长度
GLM-5.1	15.1	5.7	1.9	10.0	-	-
Gemini 2.5 Flash	21.8	22.1	6.5	17.8	1012.8	2370.3
DeepSeek-V3.2	81.6	53.3	6.5	46.8	2821.6	3684.2
GPT-5.5	92.0	63.3	31.8	61.1	573.0	1515.4
Claude Opus 4.7	98.9	76.4	44.9	72.8	456.8	1228.2
ChemCrow (Opus 4.7)	100.0	75.4	48.6	73.5	-	-
Qwen3-8B + SFT (AR基准)	74.5	42.8	20.5	44.0	648.8	1769.4
SDAR-8B (无SFT扩散基座)	14.9	11.0	8.4	11.2	-	-
ChemDiffAgent (SFT)	75.9	66.8	43.0	62.3	571.2	1746.8
ChemDiffAgent (SFT+RL)	78.2	69.9	44.9	64.6	548.2	1747.7

关键发现¶

扩散机制相对自回归的结构性代差：在同为 8B 规模、同源预训练起点（SDAR 持续预训练自 Qwen3）且享有完全相同 SFT 数据的公平对比下，ChemDiffAgent (SFT) 在单轮任务上以 89.0% 压倒 Qwen3-8B 的 69.1%（提升 19.9%）；多轮场景上以 62.3% 碾压 Qwen3-8B 的 44.0%（提升 18.3%）。这证明了非因果双向建模对于化学分子式解析与工具编排的压倒性优势。
跨体量比肩顶尖前沿模型：经过 Agentic VRPO 强化后，仅 8B 的 ChemDiffAgent 取得了 91.8% 的单轮总成功率，超越了 GPT-5.5 (70.6%) 与 Claude Opus 4.7 (90.9%)，并在多轮困难分项（6-7 轮）中达到 44.9%，打平 Claude Opus 4.7 (44.9%)，大幅优于 GPT-5.5 (31.8%)。
去噪动态的可视化印证：热力图分析显示，ChemDiffAgent 在块内去噪时，总是先快速收敛外层的工具调用定界符、Schema 槽位和控制逻辑，而后在双向上下文约束下审慎填补内部复杂的 SMILES 标记，展现了自适应降低风险的生成范式，避免了无效调用与格式错误。
紧凑高效的交互步数控制：执行步数统计表明，ChemDiffAgent 的实际调用轮数与最小黄金依赖链高度吻合（平均 3.93 轮对应目标 3.90 轮），避免了 Gemini 2.5 Flash 的冗余试错（平均 6.0 轮）或 GPT-5.5 的过早终止（平均 3.42 轮），展现出精准的目标导向推理。

亮点与洞察¶

首次将扩散语言模型引入复杂科学智能体系统：打破了过去几年将扩散模型仅作为“单轮文本生成器”或简单“网络检索器”的狭窄定位，开创性地拓展至多轮、多工具、多模态的化学深度推理中。
非对称决策风险与任意序解码的优雅数学统一：通过严格的互信息与加权风险证明，解释了为何扩散去噪比从左至右自回归解码更不易发生工具调用幻觉，为 LLM 智能体架构演进提供了坚实的理论支撑。
环境感知的专属 Agentic ELBO 训练范式：精准剥离不可控的环境返回观察与自主推理动作，避免了在外部确定性证据上白白浪费模型容量与引入潜在的训练期注意力穿越。

局限与展望¶

模型规模与工具生态局限：目前验证主要局限在 8B 参数级别和计算化学范畴（性质预测、SMILES 转换等）；更大参数规模（如 70B+）的扩展性以及湿实验（Wet-lab）真实物理操作流的闭环验证仍待展开。
推理服务部署基础设施尚不成熟：相较于成熟且经过深度工业级优化的自回归生态系统（如 vLLM、SGLang、投机采样等），扩散语言模型的 KV 缓存重用机制、分块调度与推断加速生态仍在初期，单次生成的高频去噪步数在一定程度上拉高了延迟。
多轮交互的块大小静态约束：固定 64 标记的块大小与 64 去噪步长对于动态波动的复杂分子名或超长 SMILES 串可能不是最优分配，未来可探索动态自适应块长与自适应步长扩散调度。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [首创基于掩码扩散模型的化学工具推理智能体，理论性证明了任意序解码相比自回归因果采样的风险控制优势]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖单轮六大化学任务族与多轮图约束调用基准，与同源 AR、专用化学智能体及顶尖商业模型进行了详尽对标]
写作质量: ⭐⭐⭐⭐⭐ [理论推导严密自洽，问题剖析一针见血，架构图与去噪解码动力学分析清晰透彻]
价值: ⭐⭐⭐⭐⭐ [有力证明了扩散架构在科学复杂推理与智能体场景下的独特潜力，为摆脱自回归因果注意力的固有缺陷提供了突破口]