Beyond Hard Masks: Progressive Token Evolution for Diffusion Language Models

TL;DR

EvoToken-DLM tackles inefficiencies in masked diffusion language models by replacing hard mask transitions with evolving soft token distributions, enabling progressive and revisable decoding. The approach introduces continuous trajectory supervision to align training with iterative probabilistic updates in a continuous embedding space, while preserving architectural compatibility with KV-caching and blockwise diffusion. It defines a progressive inference process with four token states and blockwise decoding, plus a training scheme that simulates refinement trajectories. Empirical results across multiple backbones and benchmarks show substantial improvements over strong masked DLM baselines, demonstrating improved reasoning and generation quality with minimal latency overhead. Overall, EvoToken-DLM offers a general, practical enhancement to diffusion language models, enabling robust token evolution and revisable decoding in realistic deployment settings.

Abstract

Diffusion Language Models (DLMs) offer a promising alternative for language modeling by enabling parallel decoding through iterative refinement. However, most DLMs rely on hard binary masking and discrete token assignments, which hinder the revision of early decisions and underutilize intermediate probabilistic representations. In this paper, we propose EvoToken-DLM, a novel diffusion-based language modeling approach that replaces hard binary masks with evolving soft token distributions. EvoToken-DLM enables a progressive transition from masked states to discrete outputs, supporting revisable decoding. To effectively support this evolution, we introduce continuous trajectory supervision, which aligns training objectives with iterative probabilistic updates. Extensive experiments across multiple benchmarks show that EvoToken-DLM consistently achieves superior performance, outperforming strong diffusion-based and masked DLM baselines. Project webpage: https://aim-uofa.github.io/EvoTokenDLM.

Figures (15)

• Figure 1: Inefficient utilization of predictions in masked diffusion language models, where distributions are computed for all positions but only a subset are used for decoding. $[M_1, M_2, \dots, M_n]$ denote the initial mask tokens following prompt $P$, and $dist_i$ represents the predicted probability distribution for the $i$-th token in the generation sequence. In this example, the total sequence of $542$ tokens consists of $30$ prompt tokens and $512$ generated tokens, while only two positions are updated per step.
• Figure 2: Comparison between MDLMs and EvoToken-DLM. (a) Standard MDLMs employ only two token states, alternating between <mask> and discrete decoded tokens, leading to abrupt mask-to-token transitions. (b) EvoToken-DLM introduces soft tokens represented by probability distributions and four token states, enabling tokens to evolve progressively through iterative refinement. The top-right panel illustrates a quantitative comparison between the two approaches under the same settings based on LLaDA-Instruct-8B.
• Figure 3: Progressive step-wise token update with blockwise decoding in EvoToken-DLM.
• Figure 4: Continuous trajectory supervision by performing $\Delta \tau$ consecutive refinement steps during training and applying supervision at each step, aligning the training objective with the inference process.
• Figure 5: Ablation study on the presence of intermediate refinement states in EvoToken-DLM.