Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL [R]

Reddit r/MachineLearning Papers
masked-diffusion language-models world-models reinforcement-learning agentic-rl steerable

Summary

This paper proposes using Masked Diffusion Language Models (MDLMs) as text-based world models for agentic reinforcement learning, showing that their any-order denoising objective avoids prefix mode collapse and leads to stronger performance than autoregressive baselines.

Autoregressive LLM world models factorize next-state generation left-to-right, preventing them from conditioning on globally interdependent anchors (tool schemas, trailing status fields, expected outcomes) and yielding prefix-consistent but globally incoherent rollouts. MDLMs' any-order denoising objective sidesteps this by learning every conditional direction from the same training signal. Empirically, fine-tuned MDLMs (SDAR-8B, WeDLM-8B) surpass AR baselines up to 4x their total parameter count on BLEU-1, ROUGE-L, and MAUVE across in- and out-of-domain splits, with lower Self-BLEU and higher Distinct-N confirming reduced prefix mode collapse. GRPO training on MDLM-generated rollouts shows up to +15% absolute task-success gains over AR generated training on held-out ScienceWorld, ALFWorld, and AppWorld across 1.2B–7B backbones (LFM2.5, Qwen3, Mistral) in a zero-shot transfer setting.
Original Article

Similar Articles

Masked Language Flow Models

arXiv cs.CL

This paper introduces Masked Language Flow Models (MLFMs), which incorporate masking into flow-based language models to enable continuous flow for conditional generation and allow pretrained Masked Diffusion Models to be converted. The authors propose a novel sampler that alternates continuous denoising with discrete unmasking, demonstrating for the first time that flow-based language models can scale to downstream reasoning and instruction-following tasks.

Trace-Based On-Policy Distillation for Masked Diffusion Language Models

arXiv cs.CL

A paper proposing Trace-Based On-Policy Distillation (TOPD), a teacher-supervised framework for transferring reasoning abilities to masked diffusion language models without reward estimation, achieving comparable accuracy to RL-trained counterparts with significant compute speedup.

Masked Diffusion Decoding as $x$-Prediction Flow

arXiv cs.CL

This paper reinterprets masked diffusion language model decoding as continuous clean-state prediction, introducing a flow-based framework where tokens are updated continuously and asynchronously based on confidence, achieving 97% of LLaDA's performance with 25% of the decoding budget.