ChemDiffAgent improves chemical reasoning with masked diffusion models
Masked Diffusion Language Agents for Tool-Integrated Chemical Reasoning
Artificial IntelligenceMachine Learning
Summary
Figuring out complex chemical problems often requires smart tools and careful planning, but existing AI systems usually think step-by-step from start to finish, which can limit their understanding. The authors created ChemDiffAgent, a new kind of AI that looks at the whole problem and tools it uses all at once, helping it make better decisions and understand molecules more deeply. They trained it with special steps to improve its reasoning and tested it on many chemical tasks. Their system performed about 20% better than older models and even matched much bigger AI systems.
What this means in practice
- •For chemical software developers: Integrate ChemDiffAgent’s approach to improve AI-assisted molecular design tools with better reasoning and tool planning capabilities.
- •For pharmaceutical data teams: Use ChemDiffAgent to enhance chemical analysis workflows, enabling more accurate multi-step reasoning with integrated computational tools.
Authors
Mengdi Liu, Chenghao Jia, Hong Chang, Shiguang Shan
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.