Reward-Weighted Sampling: Enhancing Non-Autoregressive Characteristics in Masked Diffusion LLMs

• URL: http://arxiv.org/abs/2509.00707v2
• Date: Sat, 20 Sep 2025 05:43:25 GMT
• Title: Reward-Weighted Sampling: Enhancing Non-Autoregressive Characteristics in Masked Diffusion LLMs
• Authors: Daehoon Gwak, Minseo Jung, Junwoo Park, Minho Park, ChaeHun Park, Junha Hyung, Jaegul Choo,
• Abstract summary: Masked diffusion models (MDMs) offer a promising non-autoregressive alternative for large language modeling.<n>Standard decoding methods for MDMs select tokens independently based on individual token confidences at each diffusion step.<n>We propose Reward-Weighted Sampling (RWS) to provide a principled global signal during the iterative diffusion process.
• Score: 44.55861996331439
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Masked diffusion models (MDMs) offer a promising non-autoregressive alternative for large language modeling. Standard decoding methods for MDMs, such as confidence-based sampling, select tokens independently based on individual token confidences at each diffusion step. However, we observe that this independent token selection often results in generation orders resembling sequential autoregressive processes, limiting the advantages of non-autoregressive modeling. To mitigate this pheonomenon, we propose Reward-Weighted Sampling (RWS), a novel decoding strategy that leverages an external reward model to provide a principled global signal during the iterative diffusion process. Specifically, at each diffusion step, RWS evaluates the quality of the entire intermediate sequence and scales token logits accordingly, guiding token selection by integrating global sequence-level coherence. This method selectively increases the confidence of tokens that initially have lower scores, thereby promoting a more non-autoregressive generation order. Furthermore, we provide theoretical justification showing that reward-weighted logit scaling induces beneficial rank reversals in token selection and consistently improves expected reward. Experiments demonstrate that RWS significantly promotes non-autoregressive generation orders, leading to improvements across multiple evaluation metrics. These results highlight the effectiveness of integrating global signals in enhancing both the non-autoregressive properties and overall performance of MDMs.

Related papers

• Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models [58.946955321428845]
  This work presents self-rewarding sequential Monte Carlo (SMC)<n>Our algorithm stems from the observation that most existing MDLMs rely on a confidence-based sampling strategy.<n>We introduce the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights.
  arXiv  Detail & Related papers  (2026-02-02T09:21:45Z)
• Learning Unmasking Policies for Diffusion Language Models [33.44995119635116]
  Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks.<n>One particularly successful variant is masked discrete diffusion, in which a buffer filled with special mask tokens is progressively replaced with tokens sampled from the model's vocabulary.<n>In this work, we propose to train sampling procedures using reinforcement learning.
  arXiv  Detail & Related papers  (2025-12-09T20:44:33Z)
• STDD:Spatio-Temporal Dynamics-Driven Token Refinement in Diffusion Language Models [12.172699141988728]
  diffusion language models (DLMs) generate text by iteratively denoising all token positions in parallel.<n>We propose a novel remasking approach that dynamically detects Temporal Variance and Spatial Deviance of each token.<n>Our approach significantly improves the operational efficiency of DLMs across mainstream datasets, achieving speedups of up to 8.9 times.
  arXiv  Detail & Related papers  (2025-12-07T12:53:48Z)
• Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective [85.06838178922791]
  Reinforcement Learning (RL) has proven highly effective for autoregressive language models.<n>But adapting these methods to diffusion large language models (dLLMs) presents fundamental challenges.<n>We propose a principled RL framework that treats entire sequence generation as a single action and uses the ELBO as a tractable sequence-level likelihood proxy.
  arXiv  Detail & Related papers  (2025-12-03T13:05:32Z)
• Improving Deepfake Detection with Reinforcement Learning-Based Adaptive Data Augmentation [60.04281435591454]
  CRDA (Curriculum Reinforcement-Learning Data Augmentation) is a novel framework guiding detectors to progressively master multi-domain forgery features.<n>Central to our approach is integrating reinforcement learning and causal inference.<n>Our method significantly improves detector generalizability, outperforming SOTA methods across multiple cross-domain datasets.
  arXiv  Detail & Related papers  (2025-11-10T12:45:52Z)
• Consolidating Reinforcement Learning for Multimodal Discrete Diffusion Models [40.82263997290613]
  We introduce MaskGRPO, the first viable approach to enable scalable multimodal reinforcement learning in discrete diffusion.<n>MaskGRPO brings more stable and efficient updates, leading to stronger reasoning performance and better generation quality.
  arXiv  Detail & Related papers  (2025-10-03T10:36:24Z)
• Multi-Metric Preference Alignment for Generative Speech Restoration [15.696247605348383]
  We propose a multi-metric preference alignment strategy for generative models.<n>We observe consistent and significant performance gains across three diverse generative paradigms.<n>Our aligned models can serve as powerful ''data annotators'', generating high-quality pseudo-labels.
  arXiv  Detail & Related papers  (2025-08-24T07:05:10Z)
• PC-Sampler: Position-Aware Calibration of Decoding Bias in Masked Diffusion Models [33.98279129315148]
  Masked diffusion models (MDMs) are powerful non-autoregressive alternatives for sequence generation.<n>In this work, we introduce Position-Aware Confidence-Calibrated Sampling (PC-Sampler), a novel decoding strategy.<n>PC-Sampler consistently outperforms existing MDM decoding strategies by more than 10% on average.
  arXiv  Detail & Related papers  (2025-08-18T15:38:37Z)
• Dynamic and Generalizable Process Reward Modeling [74.36829922727026]
  We propose Dynamic and Generalizable Process Reward Modeling (DG-PRM), which features a reward tree to capture and store fine-grained, multi-dimensional reward criteria.<n> Experimental results show that DG-PRM achieves stunning performance on prevailing benchmarks, significantly boosting model performance across tasks with dense rewards.
  arXiv  Detail & Related papers  (2025-07-23T18:17:22Z)
• Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling [70.8832906871441]
  We study how to steer generation toward desired rewards without retraining the models.<n>Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement.<n>We introduce particle Gibbs sampling for diffusion language models (PG-DLM), a novel inference-time algorithm enabling trajectory-level refinement while preserving generation perplexity.
  arXiv  Detail & Related papers  (2025-07-11T08:00:47Z)
• Reviving Any-Subset Autoregressive Models with Principled Parallel Sampling and Speculative Decoding [55.2480439325792]
  In arbitrary-order language models, it is an open question how to sample tokens in parallel from the correct joint distribution.<n>We find that a different class of models, any-subset autoregressive models (AS-ARMs), holds the solution.<n>We show that AS-ARMs achieve state-of-the-art performance among sub-200M parameter models on infilling benchmark tasks, and nearly match the performance of models 50X larger on code generation.
  arXiv  Detail & Related papers  (2025-04-29T06:33:13Z)
• Steering Masked Discrete Diffusion Models via Discrete Denoising Posterior Prediction [88.65168366064061]
  We introduce Discrete Denoising Posterior Prediction (DDPP), a novel framework that casts the task of steering pre-trained MDMs as a problem of probabilistic inference. Our framework leads to a family of three novel objectives that are all simulation-free, and thus scalable. We substantiate our designs via wet-lab validation, where we observe transient expression of reward-optimized protein sequences.
  arXiv  Detail & Related papers  (2024-10-10T17:18:30Z)
• Utilizing Multiple Inputs Autoregressive Models for Bearing Remaining Useful Life Prediction [3.448070371030467]
  We introduce a novel multi-input autoregressive model to address this challenge in RUL prediction for bearings. Through autoregressive iterations, the model attains a global receptive field, effectively overcoming the limitations in generalization. Empirical evaluation on the PMH2012 dataset demonstrates that our model, compared to other backbone networks using similar autoregressive approaches, achieves significantly lower Root Mean Square Error (RMSE) and Score.
  arXiv  Detail & Related papers  (2023-11-26T09:50:32Z)
• 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)
• Sparse Model Soups: A Recipe for Improved Pruning via Model Averaging [24.64264715041198]
  We introduce Sparse Model Soups (SMS), a novel method for merging sparse models by initiating each prune-retrain cycle with the averaged model from the previous phase. SMS preserves sparsity, exploits sparse network benefits, is modular and fully parallelizable, and substantially improves IMP's performance.
  arXiv  Detail & Related papers  (2023-06-29T08:49:41Z)

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.