PaperlayerView exact revision on arXiv

Source-linked AI summary

Masked Diffusion Language Agents for Tool-Integrated Chemical Reasoning

Mengdi Liu, Chenghao Jia, Hong Chang, Shiguang Shan

arXiv:2610.13330v1cs.AIcs.LG

TL;DR

Chemical agents built on autoregressive generation face challenges with bidirectional molecular understanding, coordinated tool planning, and multi-source evidence integration. CHEMDIFFAGENT uses masked diffusion over interleaved reasoning and tool-use trajectories, and the authors report gains over autoregressive agents and competitive performance against larger frontier models.

  • Problem

    Autoregressive chemical agents are constrained by causal attention and left-to-right generation despite chemical reasoning requiring bidirectional modality understanding, coordinated tool use, and evidence integration.

    Source: §1 · p. 2

  • Method

    CHEMDIFFAGENT iteratively denoises interleaved reasoning and tool-use trajectories and uses agentic SFT followed by variance-reduced preference optimization.

    Source: §1 · pp. 2–3

  • Results

    Under matched supervised fine-tuning conditions, CHEMDIFFAGENT gains approximately 20 percentage points over comparable-scale autoregressive agents, while its 8B model is competitive with frontier models.

    Source: §6 · p. 10

  • Takeaways & Limitations

    Behavioral analyses report more accurate tool selection and better evidence integration across multi-step workflows.

    Source: §6 · p. 10

    Evaluation is limited to an 8B scale, a fixed toolset, and computational chemistry workflows; larger backbones, broader scientific domains, and wet-lab workflows remain future directions.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

CHEMDIFFAGENT reframes chemical reasoning as iterative denoising over interleaved reasoning and tool-use trajectories, addressing the limits of sequential autoregressive agents. The paper reports improved benchmark performance and develops training and inference strategies for coordinating tool use.

1 Introduction

  • Tool-integrated chemical reasoning requires agents to interpret chemical modalities, coordinate tool calls, and synthesize returned evidence into subsequent reasoning.The introduction frames effective tool-use decisions as aligned with the agent’s knowledge boundary.Source: §1 · pp. 1–2
  • CHEMDIFFAGENT jointly denoises interleaved reasoning and tool-use trajectories, determining thoughts, tool choices, arguments, and evidence under bidirectional context.This contrasts with autoregressive agents’ sequential commitments and targets chemical modalities and tool-use decisions that are globally coupled.Source: §1 · p. 2, Figure 1 · p. 2
  • Agentic reasoning is analyzed through harder trajectory-infilling tasks during training and arbitrary-order decoding during inference; the resulting post-training framework combines SFT and VRPO.The framework is intended to enhance tool-use decision-making and chemical reasoning capabilities.Source: §1 · pp. 2–3
  • Approximately 20% overall improvement under matched budgets and constraints, with CHEMDIFFAGENT outperforming autoregressive agents and its 8B model matching or surpassing larger frontier models.The benchmark covers six chemistry tasks across single-turn and multi-turn tool-use scenarios.Source: §1 · pp. 2–3

2 Related Work

Prior chemistry agents integrate external tools and develop varied planning and collaboration strategies, but remain autoregressive. Diffusion language models and block-diffusion architectures offer alternatives, while existing diffusion agents focus on web tasks rather than scientific reasoning.

2.1 Chemistry Agents

  • Chemistry agents use structured reasoning, expert tools, tool graphs, multi-agent workflows, and self-updating memory, yet remain constrained by autoregressive generation.The paper positions its approach as the first diffusion language model chemistry agent for bidirectional chemical understanding, global tool planning, and multi-evidence integration.Source: §2.1 · p. 3

2.2 Diffusion Large Language Models and DLLM Agents

  • Diffusion language models include from-scratch training and AR-initialized approaches, while block diffusion combines bidirectional within-block attention with causal inter-block attention and arbitrary-order parallel decoding.The cited block-diffusion architectures retain non-autoregressive generation within blocks while supporting KV caching and variable-length generation.Source: §2.2 · p. 3
  • Existing diffusion language model agents address web search, whereas this work explores scientific reasoning through multi-step, multi-modal, multi-tool chemical coordination.The overview describes training with Agentic SFT and Agentic VRPO, and flexible order-agnostic reasoning at inference.Source: §2.2 · p. 3, Figure 2 · p. 4

3 Preliminaries

The preliminaries describe masked diffusion language models, their reverse prediction objective, and blockwise inputs used by block diffusion models.

3 Preliminaries

  • Masked diffusion models gradually corrupt clean completion tokens by replacing them with a mask token, conditioned on the prompt and diffusion time.The forward process is denoted q_t(y_t|y, x).Source: §3 · p. 4
  • In the reverse process, the model predicts masked tokens’ original values, and an Evidence Lower Bound serves as a surrogate objective for approximating exact log-likelihood.The text denotes the reverse probability as p_θ(·|y_t, x) and the exact log-likelihood as log π_θ(y|x).Source: §3 · p. 4
  • Block diffusion models partition the input into K continuous blocks, each of length B, and define a corresponding ELBO.Source: §3 · p. 4

4 CHEMDIFFAGENT

CHEMDIFFAGENT models tool-integrated chemical reasoning as iterative denoising over interleaved reasoning and tool-use trajectories. The section develops the rationale for flexible decoding and describes post-training methods for adapting the agent to chemistry tools and multi-step reasoning.

4.1 Theoretical Foundations of Flexible Tool Reasoning

  • Masked-block training teaches reconstruction from varied visible contexts, supporting decoding in flexible orders rather than only left to right.The objective masks contiguous blocks and trains the model to reconstruct missing tokens from arbitrary context combinations.Source: §4.1 · p. 5
  • Flexible paths can exploit differences in subproblem difficulty to reduce accumulated approximation error; when a lower-risk decision informs a higher-risk one, resolving it first can reduce weighted risk.The risk reduction is strict when the conditional mutual information between the decisions is positive.Source: §4.1 · p. 5

4.2 Agentic Supervised Fine-Tuning

  • Agentic SFT adapts diffusion language models to executable chemistry tool calling and multi-step reasoning using filtered teacher trajectories and loss on masked agent decisions, not external tool observations.Teacher trajectories are retained only when the answer is correct, reasoning is clear and complete, and tool choices and arguments follow the required format.Source: §4.2 · pp. 5–6

4.3 Agentic Variance-Reduced Preference Optimization

  • Agentic VRPO contrasts successful and failed trajectories for the same query using Agentic ELBO scores, with shared corrupted contexts for the policy and reference model.Aggregating scores reduces reliance on individual context samples.Source: §4.3 · pp. 6–7

5 Experiments

The experiments evaluate six chemistry tasks in single-turn and multi-turn settings, comparing ChemDiffAgent with matched-scale autoregressive and larger frontier agents. ChemDiffAgent achieves the best reported single-turn overall score, while the analyses examine tool-use efficiency and flexible decoding.

5.1 Dataset and Experimental Setup

  • The benchmark covers six chemistry tasks; multi-turn examples are built from tool graphs and grouped by gold-chain length, while evaluation measures answer correctness and tool-selection accuracy.Easy, Medium, and Hard correspond to 2, 3–5, and 6–7 tool-calling rounds; baselines include a matched-scale Qwen3-8B agent and ReAct-based LLMs.Source: §5.1 · pp. 7–8

5.2 Benchmarking Chemistry Agents Performance

  • Under matched conditions, ChemDiffAgent reaches a 91.8% single-turn overall score, improving 19.9% over Qwen3-8B with SFT and surpassing the compared frontier models and ChemCrow.The RL stage further raises the overall score from 89.0% to 91.8% (+2.8%).Source: §5.2 · p. 8
  • ChemDiffAgent’s executed tool rounds closely track expected steps, and it produces the shortest single-turn output among compared agents, avoiding redundant calls and long-winded trajectories.Gemini 2.5 Flash often uses extra rounds, whereas GPT-5.5 often makes too few tool calls.Source: §5.2 · pp. 8–9, Table 3 · p. 9

5.3 Case Study

  • In the illustrated block, action structure is established before selected SMILES tokens are decoded, allowing chemical details to follow a more stable global action structure.Figure 4 shows relative token decoding order within blocks, with each block’s final token highlighted.Source: §5.3 · p. 9, Figure 4 · p. 10
  • Qualitative case analyses describe coordinated tool selection, slot-level fusion of heterogeneous evidence, and refinement of uncertain arguments before execution.The authors contrast these behaviors with redundant tool use, lossy evidence aggregation, and error propagation in autoregressive agents.Source: §5.3 · p. 9

6 Conclusion and Limitations

The conclusion presents masked diffusion as a tool-integrated reasoning approach with gains over comparable-scale autoregressive agents. It also identifies limits in evaluation scope and practical deployment.

6 Conclusion and Limitations

  • Under matched supervised fine-tuning, ChemDiffAgent gains approximately 20 percentage points over comparable-scale AR agents, while its 8B model is competitive with GPT-5.5 and Claude Opus 4.7.The authors also report more accurate tool selection and better evidence integration across multi-step workflows.Source: §6 · p. 10
  • The approach reformulates reasoning as iterative denoising over interleaved trajectories, supporting bidirectional understanding and globally coordinated planning through two-stage post-training.The conclusion presents masked diffusion as a promising foundation for scientific agents coordinating reasoning, tool use, and evidence integration.Source: §6 · p. 10
  • Evaluation is limited to an 8B scale, fixed toolset, and computational chemistry; dLLM deployment infrastructure and training also remain challenging as model size and reasoning horizons increase.Scaling to larger backbones, broader domains, and wet-lab-grounded workflows is identified as future work.Source: §6 · p. 10

Broader Impacts

The authors describe potential benefits of tool-grounded, inspectable scientific-agent workflows and a shared benchmark for further research.

Broader Impacts

  • The authors suggest inspectable tool-use trajectories could help researchers trace results and reproduce computational workflows, while the benchmark may support research on reliable diffusion-based agents.They also identify potential applications in chemical discovery and access to computational chemistry tools.Source: Broader Impacts · p. 10

A Benchmark Construction Details

The benchmark construction converts chemistry resources into single- and multi-turn tool-use tasks, then filters generated examples for executable, faithful trajectories.

A Benchmark Construction Details

  • Four chemistry-oriented resource types supply six task families: composition inference, name conversion, property prediction, retrosynthesis, molecular optimization, and forward synthesis.The resources include SMolInstruct, ChemTool-Bench, MT-Mol, and T5Chem-style reaction data.Source: A Benchmark Construction Details · p. 14
  • For the single-turn split, normalized seeds become queries with a unified task schema, and restricted-tool rollouts are retained when the anchor tool is invoked and the answer matches reference output.Task-specific equivalence rules canonicalize structures, normalize formulas or labels, and allow small scalar tolerances.Source: A Benchmark Construction Details · p. 14
  • The multi-turn split samples executable subgraphs from a tool dependency graph, encouraging branching or converging structures that require coordination and evidence aggregation.Graph edges represent tool-output compatibility, and sampled subgraphs must include the seed’s anchor tool.Source: A Benchmark Construction Details · p. 14
  • Shortcut filtering uses bounded breadth-first search to remove or rewrite paths when a shorter executable graph covers the required nodes; the resulting benchmark stores verified answers and multi-turn calling chains or trees.The benchmark’s single-turn split tests canonical tool execution, while its multi-turn split tests coordination, intermediate-state tracking, and answer synthesis.Source: A Benchmark Construction Details · pp. 14–15
  • A teacher rewrites each seed query to preserve its semantic core, require the sampled tools, and ground the final answer in the anchor tool’s reference output.Generated tasks can combine representation conversion, multi-step reasoning, and aggregation of tool results.Source: A Benchmark Construction Details · p. 15
  • Generated multi-turn tasks are retained only when execution invokes the anchor tool, matches its normalized input and reference output, and keeps the final answer consistent with verified evidence.Rollouts run in a restricted environment containing only the sampled tools.Source: A Benchmark Construction Details · p. 15

B Evaluation Metrics

The evaluation combines GPT-4o-judged final-answer correctness with tool-selection diagnostics based on expected tool use. Multi-turn results are grouped by gold calling-chain length, with totals computed over all examples.

B Evaluation Metrics

  • Answer correctness is judged against the reference answer by GPT-4o and remains the primary success criterion.A model can still be considered correct if it finds a valid solution with fewer tool calls than the stored gold structure.Source: B Evaluation Metrics · pp. 15–16
  • Tool selection accuracy compares invoked tools with the expected tools in a minimal executable dependency graph.The score checks whether each expected tool name appears at least once; it does not assess arguments, outputs, or call order.Source: B Evaluation Metrics · p. 16
  • Multi-turn examples are grouped by gold calling-chain length: Easy has 2 rounds, Medium 3–5, and Hard 6–7.The total score is computed over all evaluated examples rather than as an unweighted average of group scores.Source: B Evaluation Metrics · p. 16

C Implementation Details

The implementation uses SDAR-8B for ChemDiffAgent and evaluates it alongside tool-enabled baselines. Agentic SFT and VRPO use distinct trajectory-training setups, while evaluation applies a shared protocol to non-agent baselines.

C Implementation Details

  • Experiments use SDAR-8B on 8 NVIDIA A100 GPUs, with block size and denoising steps set to 64 and a default maximum of 7 tool-calling rounds.The round limit is changed only when otherwise specified.Source: C Implementation Details · p. 16
  • Agentic SFT trains on constructed chemistry-agent trajectories with Agentic ELBO loss only on agent-generated think and tool_call regions.External tool_response regions are excluded from the loss; training uses AdamW with learning rate 1 × 10−5.Source: C Implementation Details · p. 16
  • Agentic VRPO pairs failed SFT trajectories with correct trajectories verified by tool execution and estimates likelihoods using Agentic ELBO.Preference scaling is β = 0.1, and likelihood estimation uses 2 Monte Carlo samples.Source: C Implementation Details · p. 16
  • Non-agent baselines use the same ReAct-style tool-use protocol and required chemistry tools; open-source chemistry agents retain their original execution settings.GPT-4o judges answer correctness, and tool-selection accuracy uses the expected calling sequence.Source: C Implementation Details · p. 17

D Theoretical Analysis of Flexible Ordering

The analysis contrasts flexible DLLM decoding with fixed autoregressive order, identifying training support for block infilling and conditions under which alternative ordering can improve prediction or risk-weighted loss.

D.1 Block-Level Infilling Support

  • The BDLM objective supplies reconstruction supervision across non-empty block-mask patterns, supporting completion from partially observed blocks.For a clean completion partitioned into blocks, the masking process samples a non-empty subset of blocks to replace with mask tokens.Source: §D.1 · p. 17

D.2 Adaptive Path Advantage under Non-Uniform Subproblem Difficulty

  • Under non-uniform subproblem difficulty, DLLMs can resolve easier or informative positions first, potentially reducing accumulated approximation error relative to fixed AR order.The advantage arises because intermediate contexts affect approximation error, while the entropy term is invariant to generation order.Source: §D.2 · p. 18

D.3 Strict Ordering Advantage under Non-Uniform Decision Risk

  • Flexible ordering has a strict risk-weighted advantage when a lower-risk decision informs an earlier, higher-risk decision.The analysis requires positive conditional information between the decisions and a greater expected downstream error cost for the high-risk decision.Source: §D.3 · p. 18
  • In chemistry reasoning, tool calls and molecular strings contain decisions with differing difficulty and consequences, so useful ordering must account for both information and error costs.The adaptive-path analysis concerns approximation error, whereas the risk-weighted analysis concerns information value under unequal error costs; an easier decision is not necessarily lower-risk.Source: §D.3 · p. 19

E Denoising Visualization of Tool-use Decisions

The visualization depicts tokens resolving at different denoising steps within a tool-calling block. It shows a non-left-to-right pattern in which action structure is resolved before some chemical argument tokens, enabling later refinement within the established tool-use context.

E Denoising Visualization of Tool-use Decisions

  • The pattern is non-left-to-right: tool-call delimiters and argument formatting resolve relatively early, while some SMILES tokens resolve later.This allows individual token predictions to use information resolved elsewhere in the same block.Source: E Denoising Visualization of Tool-use… · p. 19, Figure 5 · p. 20
  • Figure 5 encodes relative token-resolution timing within a representative tool-calling block, with lighter colors earlier and darker colors later.The denoising process exposes the temporal structure of token decisions within the block.Source: E Denoising Visualization of Tool-use… · p. 19, Figure 5 · p. 20
  • Resolving the action scaffold before uncertain argument content lets later refinement condition on the established tool-use format.This is especially relevant to SMILES, where small token mistakes can change the molecule or make the input invalid.Source: E Denoising Visualization of Tool-use… · pp. 19–20, Figure 5 · p. 20
  • The observed order is consistent with forming an action-level structure before refining uncertain chemical details, rather than committing token by token.Successful execution depends on consistency between reasoning intent, tool selection, argument format, and exact chemical input.Source: E Denoising Visualization of Tool-use… · pp. 20–21

F Case Studies

This section contrasts autoregressive agents’ local continuation with ChemDiffAgent’s global planning for coordinating tool dependencies, preserving heterogeneous outputs, and correcting inputs before downstream computation. It also reports checklist responses on ethics, societal impacts, safeguards, limitations, experimental details, reproducibility, statistical significance, and planned release of code and pretrained models.

F Case Studies

  • In tool selection, the AR agent repeats entity-lookup calls after obtaining an identifier instead of moving to formula conversion and numerical calculation; DLLM plans the three stages and selects tools accordingly.The required dependency chain is compound name → CID → molecular formula → element percentage.Source: F Case Studies · p. 21
  • For heterogeneous property summaries, AR’s lossy compression can cause schema drift and type mismatches, whereas DLLM’s latent slot-level aggregation maintains required fields and output consistency.The example requires collecting properties including biological activity, solubility, molecular descriptors, functional groups, and a long side-effect profile into fixed-key JSON.Source: F Case Studies · p. 22
  • In product-property computation, noisy initial products can produce unstable downstream calls; DLLM inserts validation or canonicalization before using the corrected product for subsequent properties.The case contrasts retrying similar property queries with correcting the product representation before continuing.Source: F Case Studies · p. 23
  • Across three cases, autoregressive agents’ reliance on the latest visible state leads to tool-selection traps, lossy aggregation, and error propagation, while DLLM global planning supports coordination, structured outputs, and correction.The authors say these qualitative observations are consistent with stronger tool-selection accuracy and more stable multi-turn performance, especially on tasks requiring long-range coordination across tools.Source: F Case Studies · pp. 21, 23

NeurIPS Paper Checklist

  • The checklist marks limitations, complete theoretical proofs, reproducibility information, training and test details, and statistical-significance reporting as addressed, and says code and pretrained models would be released upon acceptance.It points to the Appendix for limitations, Section 5.1 for reproducibility and training/test details, and Section 5 for statistical significance.Source: NeurIPS Paper Checklist · pp. 24–26
  • The authors report sufficient compute-resource information and say external assets are credited and their terms respected.The compute-resources response points to Section 5.Source: NeurIPS Paper Checklist · pp. 27–28
  • The authors report ethics compliance and say the paper discusses both positive and negative societal impacts, with details in the appendix.They answer Yes to both checklist questions.Source: NeurIPS Paper Checklist · pp. 27–28
  • The authors answer No on release safeguards, explaining that the paper poses no such risks.The checklist question concerns safeguards for data or models with high misuse risk.Source: NeurIPS Paper Checklist · p. 28
  • The authors state that they release no new assets, conduct no human-subjects or crowdsourcing research, and do not use LLMs in the core methodology.The checklist marks these items N/A and gives those explanations.Source: NeurIPS Paper Checklist · pp. 29–30
Loading 2610.13330v1…