cs.AIOct 9, 2026

Masked Diffusion Language Agents for Tool-Integrated Chemical Reasoning

Authors: Mengdi Liu, Chenghao Jia, Hong Chang, Shiguang Shan

Organizations: State Key Laboratory of AI Safety, Institute of Computing Technology, CAS, China · University of Chinese Academy of Sciences (CAS), China

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Designing tool-integrated reasoning agent systems for complex chemical tasks remains a fundamental challenge. Although recent chemical agents have demonstrated promising performance through the integration of external tools, they are still built on the autoregressive paradigm, which is inherently constrained by causal attention and left-to-right generation. As a result, they often struggle with bidirectional molecular understanding, global coordinated tool planning and multi-source evidence integration. To address this challenge, we introduce ChemDiffAgent, the first masked diffusion language agent for tool-integrated chemical reasoning. It reformulates chemical reasoning as iterative denoising over interleaved reasoning and tool-use trajectories, enabling tool-use decisions to be jointly determined under full-trajectory context and thereby better aligning the agent's decision boundary with its knowledge boundary. We further analyze how this paradigm benefits agentic reasoning through harder infilling subproblems during training and flexible arbitrary-order decoding during inference. Then we develop a two-stage post-training framework, comprising agentic supervised fine-tuning and variance-reduced preference optimization, to enhance tool-use and reasoning capabilities. We also construct a comprehensive benchmark covering six chemistry tasks across both single-turn and multi-turn tool-use scenarios. Under matched budgets and evaluation, ChemDiffAgent consistently outperforms autoregressive agents, achieving an overall improvement of approximately 20%. Notably, our 8B model achieves performance comparable to, in some cases surpassing, significantly larger frontier models, such as GPT 5.5 and Claude Opus 4.7. Code is available at https://github.com/bilemond/ChemDiffAgent.

Figures & tables

Figure 1: Compared with autoregressive agents that commit to tool-use trajectories sequentially, ChemDiffAgent jointly refines chemical modality thinking, tool-use decisions, arguments, and evidences to better handle bidirectional chemical modalities and globally coordinated tool use.
Figure 2: Overview of ChemDiffAgent . In training, both Agentic SFT and Agentic VRPO use Block Attention and Agentic Noising to compute the Agentic ELBO. It learns from complex block-infilling tasks at training time and enables flexible order-agnostic reasoning at inference time.
MethodEasyMed.HardTotal
Strong LLMs w/ ReAct
GLM-5.115.15.71.910.0
Gemini 2.5 Flash21.822.16.517.8
DeepSeek-V3.281.653.36.546.8
GPT-5.592.063.331.861.1
Claude Opus 4.798.976.444.972.8
Table 2: Success rate on multi-turn tool-use setting. Easy, Medium, Hard correspond to 2, 3–5, 6–7 tool-calling rounds.
Figure 3: Multi-turn tool selection accuracy across difficulty levels. Easy, Medium, and Hard correspond to 2, 3–5, and 6–7 tool-calling rounds, respectively.
Method234567Avg.
Claude Opus 4.71.993.294.165.505.656.623.90
Gemini 2.5 Flash5.745.815.637.006.856.366.00
DeepSeek-V3.22.784.285.117.337.097.324.84
GPT-5.51.902.884.056.004.765.523.42
Qwen3-8B(SFT)3.675.355.267.337.507.585.60
ChemDiffAgent(SFT)2.183.443.894.835.976.934.07
Table 3: Efficiency analysis of tool-use trajectories. Left: average executed tool rounds under different expected rounds. Right: average output length in single-turn and multi-turn settings.
Figure 4: Heatmap showing the relative decoding order of each token within each block of an example response. The last token of each block is highlighted with a red box. For brevity, only a single turn is shown; the complete response is provided in Fig 5 .
Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Token-level denoising order of a representative ChemDiffAgent tool-calling block. Lighter colors indicate tokens resolved earlier, while darker colors indicate tokens resolved later.
More in this topic

Chemical Reasoning in LLMs

See all 13 papers

Explore similar work

CardsList
Sep 28, 2026

TMCS: Tool-Grounded Multi-Agent Reasoning for Compositional Chemical Problem Solving

Despite the promise of Large Language Models (LLMs) in computational chemistry, rigorous combinatorial chemistry problems remain difficult because they require quantitatively constrained molecular modification, candidate validation, and systematic revision after failed attempts. Existing tool-augmented chemical agents demonstrate useful planning and tool use, but they rarely provide a unified loop for property-driven molecular optimization and workflow-level composition. To bridge this gap, we propose Tool-Grounded Multi-Agent Reasoning for Compositional Chemical Problem Solving (TMCS), a step-by-step multi-agent framework that formalizes chemical problem solving as an interpretable, tool-augmented workflow. At the task level, specialized agents leverage external tools, few-shot trajectory memory, and structured reflection to iteratively refine solutions. At the workflow level, TMCS chains generation, understanding, editing, description, and optimization into a closed-loop pipeline. Evaluations across multiple chemical tasks demonstrate that TMCS consistently enhances chemical reasoning across both open- and closed-source base models, achieving state-of-the-art performance.

Shengqin Wang, Jie Jin, Yu Cheng +4

East China Normal University, Shanghai Innovation Institute, University College London, Huawei Noah’s Ark Lab

Molecular OptimizationTool-Augmented Language Model AgentsComputational Quantum ChemistryLLM AgentsMulti-Agent ReasoningLarge Language Model-Guided Scientific Discovery
AskPDF
Sep 27, 2026

ChemOPD: Multi-Teacher On-Policy Distillation for Multi-Task Chemical Reasoning

Large language models are increasingly expected to support diverse chemical reasoning capabilities within a unified model. One approach is to develop specialized capabilities separately and consolidate them through multi-teacher on-policy distillation, but this raises two questions: how should specialization be organized, and how should specialist guidance be integrated? We introduce ChemOPD, which addresses both. We estimate task affinities from supervised fine-tuning gradients and solve a constrained mixed-integer program(MIP) to construct partially overlapping specialist groups. During distillation, we retain a generalist teacher trained on all tasks so that specialist guidance supplements rather than replaces its supervision. Our anchor-residual objective gradually increases the routed specialist's contribution on student-generated responses. On ChemCoTBench, affinity-guided specialization produces task-dependent gains over the generalist teacher and improves several capabilities beyond semantic task grouping. Yet stronger teacher-side performance does not automatically yield stronger students: with the same specialists and routes, anchor-residual OPD improves most reported metrics over specialist-only distillation and realizes a larger share of the available teacher gains. These results highlight specialization and capability integration as connected but distinct design problems in chemical reasoning.

Yaoyao Xu, Xinjian Zhao, Xiaozhuang Song +2

School of Data Science, The Chinese University of Hong Kong, Shenzhen · Shanghai Artificial Intelligence Laboratory

Chemical Reasoning in LLMsReasoning DistillationMulti-Teacher Knowledge DistillationOn-Policy Distillation
AskPDF
May 8, 2026

Can Agents Price a Reaction? Evaluating LLMs on Chemical Cost Reasoning

Large Language Models (LLMs) have become increasingly capable as tool-using agents, with benchmarks spanning diverse general agentic tasks. Yet rigorous evaluation of scientific tool use remains limited. In chemistry, recent agents can plan syntheses and invoke domain-specific tools, but evaluations often rely on curated demonstrations, expert assessment, or LLM-as-judge scoring rather than exact, judge-free ground truth. We address this gap with chemical procurement cost estimation, a practical task in which an agent must ground chemical identities, retrieve supplier quotes, select valid purchasable packs, normalize quantities, and compute cost from a reaction description. We introduce ChemCost, a benchmark of 1,427 evaluable reactions grounded to a frozen pricing snapshot covering 2,261 chemicals and 230,775 supplier quotes, supporting scalar scoring and stage-level diagnosis of grounding, retrieval, procurement, and arithmetic failures. To evaluate robustness, we further construct controlled noise-injected views that perturb chemical aliases, quantity expressions, missing fields, and input formatting. Experiments with frontier, open-weight, and chemistry-specialized LLM agents show that tool access is necessary but insufficient for solving the task. The strongest agents reach only 50.6% accuracy within 25% relative error on clean inputs and degrade substantially with realistic noise. Stage-level analysis further shows that failures arise from brittle parsing, ineffective evidence integration, invalid pack selection, and non-convergent tool use.

Yuyang Wu, Yue Huang, Shuaike Shen +8

Carnegie Mellon University · University of Notre Dame · University of North Carolina, Chapel Hill +1

AI Agents for Scientific DiscoveryLLM Agent EvaluationAI Agent BenchmarksLLM Tool UseLLM Agent ReliabilityChemical Reasoning in LLMs
AskPDF