Learning Unmasking Policies for Diffusion Language Models

2025-12-09

AI Summary

Dynamic unmasking of tokens in diffusion language models (dLLMs) relies on hand-crafted heuristics—e.g., confidence-based thresholds—leading to tedious hyperparameter tuning and degraded performance under large buffer sizes. Method: We formulate unmasking as a Markov decision process and propose a lightweight, single-layer Transformer-based reinforcement learning policy network for end-to-end, hyperparameter-free, layer-wise unmasking decisions. Our approach integrates token confidence mapping within a unified semi-autoregressive–full-diffusion framework. Contribution/Results: In the full-diffusion setting, our method significantly outperforms state-of-the-art heuristic baselines. The learned policy exhibits strong generalization across diverse dLLM architectures and sequence lengths, simultaneously improving generation quality and inference throughput. To our knowledge, this is the first work to cast iterative token unmasking in dLLMs as a learnable sequential decision problem, eliminating reliance on fixed rules while maintaining efficiency and scalability.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: (Large) Language Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Diffusion (Large) Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks, while holding the promise of being more efficient during inference. 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. Efficiency can be gained by unmasking several tokens in parallel, but doing too many at once risks degrading the generation quality. Thus, one critical design aspect of dLLMs is the sampling procedure that selects, at each step of the diffusion process, which tokens to replace. Indeed, recent work has found that heuristic strategies such as confidence thresholding lead to both higher quality and token throughput compared to random unmasking. However, such heuristics have downsides: they require manual tuning, and we observe that their performance degrades with larger buffer sizes. In this work, we instead propose to train sampling procedures using reinforcement learning. Specifically, we formalize masked diffusion sampling as a Markov decision process in which the dLLM serves as the environment, and propose a lightweight policy architecture based on a single-layer transformer that maps dLLM token confidences to unmasking decisions. Our experiments show that these trained policies match the performance of state-of-the-art heuristics when combined with semi-autoregressive generation, while outperforming them in the full diffusion setting. We also examine the transferability of these policies, finding that they can generalize to new underlying dLLMs and longer sequence lengths. However, we also observe that their performance degrades when applied to out-of-domain data, and that fine-grained tuning of the accuracy-efficiency trade-off can be challenging with our approach.

Problem

• Optimizing token unmasking strategies in diffusion language models
• Replacing heuristic methods with learned policies for sampling
• Balancing generation quality and efficiency in parallel token generation

Innovation

• Reinforcement learning trains unmasking policies for diffusion
• Lightweight transformer maps token confidences to unmasking decisions
• Policies generalize to new models and longer sequences