Masked Diffusion Decoding as $x$-Prediction Flow

2026-06-27

AI Summary

This work addresses a key limitation of conventional masked diffusion language models, which rely on binary masking decisions during decoding and thus struggle to represent intermediate confidence states, ultimately constraining performance under limited decoding budgets. To overcome this, the authors reformulate masked token prediction as clean-token prediction (x-prediction), establishing a continuous diffusion decoding framework that enables progressive and revisable token updates at each position throughout the diffusion process. The approach introduces two core innovations: a confidence-based asynchronous update mechanism that replaces global synchronous scheduling, and a lightweight policy network trained via reinforcement learning to optimize the decoding trajectory. Experimental results demonstrate that the proposed method achieves 97% of the performance of the pretrained LLaDA model on HumanEval while using only 25% of the decoding budget.

Abstract

Masked diffusion language models (MDLMs) generate text by iteratively unmasking tokens, but their standard decoder reduces each step to a binary action: a position is either committed to a single token or left fully masked, with no representation of partial belief in between. This all-or-nothing regime discards rich predictive information and forces premature, irrevocable commitments, leading to poor performance under a limited decoding budget. In this paper, we reinterpret mask prediction as clean-state prediction ($x$-prediction) and show that it can be used to induce a continuous flow in input embedding space. Building on this view, we propose a continuous decoding framework for MDLMs where tokens can accumulate partial progress at each diffusion step and remain revisable. To match the uneven contextual constraints across positions in language, we replace the globally synchronous schedule in image diffusion with a confidence-based asynchronous update in which the diffusion progress is token-wise accumulated. Additionally, we introduce a lightweight policy network and formulate its training as a reinforcement learning problem. Applied to pretrained LLaDA, our continuous decoder reaches 97% of its performance on the HumanEval dataset with 25% of decoding budget.

Problem

• masked diffusion language models
• decoding budget
• token commitment
• partial belief
• x-prediction

Innovation

• masked diffusion language models
• x-prediction
• continuous decoding
• asynchronous update
• reinforcement learning