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

arXiv: https://arxiv.org/abs/2509.00707 · HTML: https://mlsanity.com/paper/2509.00707/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 24 of its 42 references that are in the map.
- First posted on arXiv 31 Aug 2025.

## Placed in

- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)
- LLM alignment / RLHF / preference optimisation / [Diffusion and multimodal model alignment](https://mlsanity.com/alignment/diffusion-and-multimodal-model-alignment/)

## Builds on (most cited first, 24 of 24)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [A Diversity-Promoting Objective Function for Neural Conversation Models](https://mlsanity.com/paper/1510.03055/) (2015)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks](https://mlsanity.com/paper/1506.03099/) (2015)
- [Plug and Play Language Models: A Simple Approach to Controlled Text Generation](https://mlsanity.com/paper/1912.02164/) (2019)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Sequence Level Training with Recurrent Neural Networks](https://mlsanity.com/paper/1511.06732/) (2015)
- [FUDGE: Controlled Text Generation With Future Discriminators](https://mlsanity.com/paper/2104.05218/) (2021)
- [Skywork-Reward: Bag of Tricks for Reward Modeling in LLMs](https://mlsanity.com/paper/2410.18451/) (2024)
- [Mercury: Ultra-Fast Language Models Based on Diffusion](https://mlsanity.com/paper/2506.17298/) (2025)
- [Advancing LLM Reasoning Generalists with Preference Trees](https://mlsanity.com/paper/2404.02078/) (2024)
- [BERT has a Mouth, and It Must Speak: BERT as a Markov Random Field Language Model](https://mlsanity.com/paper/1902.04094/) (2019)
- [Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs](https://mlsanity.com/paper/2406.10216/) (2024)
- [Quantile Regression for Distributional Reward Models in RLHF](https://mlsanity.com/paper/2409.10164/) (2024)
- [BiasDPO: Mitigating Bias in Language Models through Direct Preference Optimization](https://mlsanity.com/paper/2407.13928/) (2024)

## Built on by (most cited first, 10 of 10)

- [dParallel: Learnable Parallel Decoding for dLLMs](https://mlsanity.com/paper/2509.26488/) (2025)
- [ReFusion: A Diffusion Large Language Model with Parallel Autoregressive Decoding](https://mlsanity.com/paper/2512.13586/) (2025)
- [DMax: Aggressive Parallel Decoding for dLLMs](https://mlsanity.com/paper/2604.08302/) (2026)
- [Consistent Diffusion Language Models](https://mlsanity.com/paper/2605.00161/) (2026)
- [Knowledge Editing for Masked Diffusion Language Models](https://mlsanity.com/paper/2606.03924/) (2026)
- [Enabling KV Caching of Shared Prefix for Diffusion Language Models](https://mlsanity.com/paper/2606.07571/) (2026)
- [From Interface to Inference: Eliciting Any-Order Inference from Any-Order Models](https://mlsanity.com/paper/2607.26504/) (2026)
- [FourierSampler: Unlocking Non-Autoregressive Potential in Diffusion Language Models via Frequency-Guided Generation](https://mlsanity.com/paper/2601.23182/) (2026)
- [When to Plan, When to Polish: Noise Level as a Granularity Axis for Diffusion Language Models](https://mlsanity.com/paper/2606.21802/) (2026)
- [Visual Information-Guided Parallel Decoding for Diffusion Multimodal Large Language Models](https://mlsanity.com/paper/2608.26580/) (2026)

## Abstract (arXiv)

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.