Learning Unmasking Policies for Diffusion Language Models

Diffusion Large Language Models (dLLMs) are considered a promising alternative to Autoregressive models in terms of inference efficiency, as they can decode multiple tokens in parallel. However, the sampling process that determines which tokens to restore (unmask) first at each step acts as a key design factor for performance.

Previously, heuristic strategies such as Confidence Thresholding have been used, but these have limitations in that they require manual tuning and their performance degrades as block size increases. In this study, we formalize this sampling process as a Markov Decision Process (MDP) and propose a method for learning an optimal sampling policy through reinforcement learning.

The characteristics of the proposed method are as follows:

• Lightweight Policy: Based on a single-layer Transformer, it maps the dLLM's token confidence to unmasking decisions.
• Proven Performance: Experimental results show that the proposed policy achieved performance on par with existing SOTA heuristics in Semi-autoregressive (Block) generation, and outperformed them in the Full-diffusion setting.