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

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

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

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

    Sep 28, 2026Shengqin Wang, Jie Jin, Yu Cheng +4Molecular OptimizationTool-Augmented Language Model Agents

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

    Sep 27, 2026Yaoyao Xu, Xinjian Zhao, Xiaozhuang Song +2Chemical Reasoning in LLMsReasoning Distillation

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

    May 8, 2026Yuyang Wu, Yue Huang, Shuaike Shen +8AI Agents for Scientific DiscoveryLLM Agent Evaluation