Diffusion language models enable editable planning states for better search

DiffuPlan: Editable Planning State for Diffusion Language Model Agents

Computation and LanguageArtificial Intelligence

Summary

Most language models generate text step-by-step and can only add new information at the end, making it hard to fix past mistakes. The authors show that diffusion language models can rewrite earlier parts of their plan as new information comes in, leading to better and more flexible answers. They built a system called DiffuPlan that keeps one editable plan instead of just appending text, improving performance on complex question-answering tasks. This approach shows that diffusion models can do things autoregressive models cannot, like revising their previous reasoning steps.

What this means in practice

  • For ai system builders: Create AI agents that revise earlier reasoning steps when new input arrives, improving decision-making accuracy on complex tasks.
  • For customer support software teams: Implement chatbots that update their previous responses dynamically as new information from conversations emerges, enhancing user experience.

Authors

Qisheng Zhou, Zhen Xiong, Qiaoyu Tan

Abstract

Diffusion language models (dLLMs) are emerging as an alternative to autoregressive (AR) language models, yet dLLM agents largely inherit an interaction protocol designed for AR generation: once intermediate state is written, later observations can only be incorporated by appending more text. This is a poor match to a defining capability of dLLMs -- reconstructing earlier text conditioned on context that appears both before and after it. We ask whether this capability can support a different agent-state paradigm in which intermediate planning state is directly revisable as evidence arrives. We introduce DiffuPlan, a diffusion search agent that maintains a single editable planning state and reconstructs only the affected steps when later evidence resolves an entity or changes the search route. To support variable-length revisions, we introduce Expansion-Partition Marginalization (EPM), which marginalizes over equivalent expansion paths. We further adapt diffusion trajectory RL to multi-turn search. Across four multi-hop QA benchmarks, DiffuPlan achieves 50.24 average F1, outperforming strong AR search agents. More importantly, under identical data and training, editable planning state improves the diffusion agent by 2.44 F1 versus only 0.47 for its AR source model. These results show that dLLMs need not imitate append-only AR agents: revising committed state is a distinct agentic advantage.