Diffusion-State Policy Optimization for Masked Diffusion Language Models

• URL: http://arxiv.org/abs/2602.06462v2
• Date: Mon, 09 Feb 2026 03:24:51 GMT
• Title: Diffusion-State Policy Optimization for Masked Diffusion Language Models
• Abstract summary: Masked diffusion language models generate by iteratively filling masked tokens over multiple denoising steps.<n>We propose DiSPO, a plug-in credit-assignment layer that directly optimize intermediate filling decisions.
• Score: 23.931675509910047
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Masked diffusion language models generate by iteratively filling masked tokens over multiple denoising steps, so learning only from a terminal reward on the final completion yields coarse credit assignment over intermediate decisions. We propose DiSPO (Diffusion-State Policy Optimization), a plug-in credit-assignment layer that directly optimizes intermediate filling decisions. At selected intermediate masked states, DiSPO branches by resampling fillings for the currently masked positions from rollout-cached logits, scores the resulting completions, and updates only the newly filled tokens -- without additional multi-step diffusion rollouts. We formalize a fixed-state objective for branched completions and derive a policy-gradient estimator that can be combined with terminal-feedback policy optimization using the same rollouts. On LLaDA-8B-Instruct, DiSPO consistently improves over the terminal-feedback diffu-GRPO baseline on math and planning benchmarks under matched rollout compute and optimizer steps. Our code will be available at https://daioba.github.io/dispo .

Related papers

• Adaptive Multi-Step Lookahead Decoding for Diffusion Language Models [21.088701659053466]
  Masked diffusion language models (DLMs) enable parallel text generation by iteratively refining masked tokens.<n>Recent lookahead-based decoding methods improve the accuracy--efficiency trade-off by exploring future decoding states.<n>We propose AdaLook, an adaptive lookahead framework for DLM decoding.
  arXiv  Detail & Related papers  (2026-07-17T06:04:04Z)
• Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models [59.51249894128724]
  Masked diffusion language models can reduce inference steps by revealing multiple tokens per denoising iteration.<n>We propose ADAS, a training-free reranking rule for parallel masked diffusion decoding.
  arXiv  Detail & Related papers  (2026-06-09T13:17:27Z)
• Stochastic MeanFlow Policies: One-Step Generative Control with Entropic Mirror Descent [53.828537014796574]
  Online off-policy reinforcement learning (RL) is shaped by two choices: the policy class and the update rule.<n>We propose MeanFlow Policies, a one-step generative policy class that maps noise to actions through a MeanFlow transformation.<n>Across seven MuJoCo benchmarks, S improves over Gaussian and generative baselines while retaining single-step inference efficiency.
  arXiv  Detail & Related papers  (2026-05-20T15:14:14Z)
• PSD: Pushing the Pareto Frontier of Diffusion LLMs via Parallel Speculative Decoding [32.667256256847246]
  Diffusion large language models (dLLMs) generate text by iteratively denoising masked token sequences.<n>We propose Parallel Speculative Decoding (PSD), a training-free framework that jointly improves inference along both axes.
  arXiv  Detail & Related papers  (2026-05-15T04:43:02Z)
• F-GRPO: Factorized Group-Relative Policy Optimization for Unified Candidate Generation and Ranking [79.49893545611779]
  Large Language Models (LLMs) can generate a subset and order it within a single autoregressive pass.<n>This flexibility introduces a new optimization challenge: the model must search an output space while receiving utility feedback only after the full ranked list is generated.<n>This credit assignment gap makes end-to-end optimization unstable and sample-inefficient.<n>We propose a unified framework that performs both within a single autoregressive rollout and optimize them end-to-end via factorized group-relative policy optimization.
  arXiv  Detail & Related papers  (2026-05-13T04:52:33Z)
• DyLLM: Efficient Diffusion LLM Inference via Saliency-based Token Selection and Partial Attention [2.7422645382944935]
  We present DyLLM, a training-free inference framework that accelerates decoding by selectively computing only salient tokens.<n>Across diverse reasoning and code-generation benchmarks, DyLLM achieves up to 9.6x higher throughput.
  arXiv  Detail & Related papers  (2026-03-09T07:02:01Z)
• Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching [66.39914384073145]
  We propose a self-consistency framework that turns cheap diffusion-sampled reasoning into a reusable pool of step-level candidates.<n>We find that step-level recombination is most beneficial on harder problems.<n>Our training-free framework improves average accuracy by up to 2 across six math and coding tasks.
  arXiv  Detail & Related papers  (2026-02-26T11:08:39Z)
• Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models [96.0074341403456]
  Inference-time compute has re-emerged as a practical way to improve LLM reasoning.<n>Most test-time scaling (TTS) algorithms rely on autoregressive decoding.<n>We propose Prism, an efficient TTS framework for dLLMs.
  arXiv  Detail & Related papers  (2026-02-02T09:14:51Z)
• Co-GRPO: Co-Optimized Group Relative Policy Optimization for Masked Diffusion Model [74.99242687133408]
  Masked Diffusion Models (MDMs) have shown promising potential across vision, language, and cross-modal generation.<n>We introduce Co-GRPO, which reformulates MDM generation as a unified Markov Decision Process (MDP) that jointly incorporates both the model and the inference schedule.
  arXiv  Detail & Related papers  (2025-12-25T12:06:04Z)
• OptPO: Optimal Rollout Allocation for Test-time Policy Optimization [11.375209834858135]
  Test-time policy optimization enables large language models to adapt to distribution shifts by leveraging feedback from self-generated rollouts.<n>We propose Optimal Rollout Allocation for Test-time Policy Optimization (OptPO), a principled framework that adaptively allocates inference budgets.
  arXiv  Detail & Related papers  (2025-12-02T15:38:52Z)
• Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies [47.6755955972232]
  We cast denoising as a KL-regularized Markov decision process (MDP) with an explicit reference policy and optimize a regularized objective.<n>We prove that the optimized policy under this framework generates samples that more closely match the data distribution than schedules.
  arXiv  Detail & Related papers  (2025-10-07T09:44:24Z)
• GRAPE: Let GPRO Supervise Query Rewriting by Ranking for Retrieval [19.73916326078242]
  The CLIP model has become a cornerstone of large-scale retrieval systems by aligning text and image data in a unified embedding space.<n>To avoid costly retraining, existing methods mainly adopt query-rewriting strategies with large language models (LLMs)<n>We address this challenge with GRAPE, a plug-and-play enhancement approach that incorporates ranking signals into retrieval-guided query rewriting.
  arXiv  Detail & Related papers  (2025-09-27T15:36:59Z)
• Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models [13.575063025878208]
  Masked diffusion language models promise fast, non-autoregressive text generation.<n>Existing samplers, which pick tokens to unmask based on model confidence, ignore interactions when unmasking multiple positions in parallel.<n>We propose the Dilated Unmasking Scheduler (DUS), an inference-only, planner-model-free method that partitions sequence positions into non-adjacent dilated groups and unmasked them in parallel.
  arXiv  Detail & Related papers  (2025-06-23T18:49:23Z)
• Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum [78.27945336558987]
  Decentralized server (DFL) eliminates reliance on client-client architecture.<n>Non-smooth regularization is often incorporated into machine learning tasks.<n>We propose a novel novel DNCFL algorithm to solve these problems.
  arXiv  Detail & Related papers  (2025-04-17T08:32:25Z)
• Aligning CodeLLMs with Direct Preference Optimization [44.34483822102872]
  This work first identifies that the commonly used PPO algorithm may be suboptimal for the alignment of CodeLLM. Based on only preference data pairs, DPO can render the model rank data automatically, giving rise to a fine-grained rewarding pattern. Studies show that our method significantly improves the performance of existing CodeLLMs on benchmarks such as MBPP and HumanEval.
  arXiv  Detail & Related papers  (2024-10-24T09:36:13Z)
• COrAL: Order-Agnostic Language Modeling for Efficient Iterative Refinement [80.18490952057125]
  Iterative refinement has emerged as an effective paradigm for enhancing the capabilities of large language models (LLMs) on complex tasks. We propose Context-Wise Order-Agnostic Language Modeling (COrAL) to overcome these challenges. Our approach models multiple token dependencies within manageable context windows, enabling the model to perform iterative refinement internally.
  arXiv  Detail & Related papers  (2024-10-12T23:56:19Z)
• Preference-Based Planning in Stochastic Environments: From Partially-Ordered Temporal Goals to Most Preferred Policies [25.731912021122287]
  We consider systems modeled as Markov decision processes, given a partially ordered preference over a set of temporally extended goals. To plan with the partially ordered preference, we introduce order theory to map a preference over temporal goals to a preference over policies for the MDP. A most preferred policy under a ordering induces a nondominated probability distribution over the finite paths in the MDP.
  arXiv  Detail & Related papers  (2024-03-27T02:46:09Z)
• Amortizing intractable inference in large language models [56.92471123778389]
  We use amortized Bayesian inference to sample from intractable posterior distributions. We empirically demonstrate that this distribution-matching paradigm of LLM fine-tuning can serve as an effective alternative to maximum-likelihood training. As an important application, we interpret chain-of-thought reasoning as a latent variable modeling problem.
  arXiv  Detail & Related papers  (2023-10-06T16:36:08Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.