PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Hard Masks: Progressive Token Evolution for Diffusion Language Models

Linhao Zhong, Linyu Wu, Bozhen Fang, Tianjian Feng, Chenchen Jing, Wen Wang, Jiaheng Zhang, Hao Chen, Chunhua Shen

arXiv:2601.07351v2cs.CLcs.AI

TL;DR

Existing masked diffusion language models make hard, irreversible token assignments and discard probabilistic information computed for positions not updated at a given step. EvoToken-DLM replaces these masks with progressively refined soft token distributions and trains them with continuous trajectory supervision; across benchmarks, it consistently outperforms strong baselines while retaining compatibility with KV-caching and blockwise architectures. Its main limitation is increased training difficulty and slower convergence when initialized from autoregressive priors.

  • Problem

    Most MDLMs use hard binary masks and discrete assignments that prevent revising decoded tokens and discard probabilistic information for positions not updated.

  • Method

    EvoToken-DLM iteratively refines soft token distributions through diffusion and uses continuous trajectory supervision to align training with probabilistic token evolution.

  • Results

    EvoToken-DLM consistently outperforms strong MDLM baselines across diverse backbones, datasets, and inference configurations.

  • Takeaways & Limitations

    Progressive soft-token evolution provides a general enhancement for diffusion language modeling that remains compatible with KV-caching and blockwise architectures.

  • Takeaways & Limitations

    Autoregressive initialization increases training difficulty and slows convergence because unidirectional pretraining differs from iterative bidirectional refinement.

Abstract

from arXiv · show

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.

1 Introduction

DLMs enable parallel generation through iterative refinement, but masked models make irreversible discrete decisions and discard many computed distributions. EvoToken-DLM addresses these limitations with evolving soft token distributions and trajectory-aligned training, and reports consistent improvements over strong MDLM baselines.

  • DLMs frame language generation as iterative refinement, enabling parallel decoding instead of strictly sequential autoregressive generation.
  • Hard binary masking finalizes decoded tokens, limiting revision of early decisions and creating an abrupt transition from uncertainty to determinism.
  • MDLMs compute distributions for all positions but update only a small subset, discarding remaining probabilistic information.
  • EvoToken-DLM replaces hard masks with evolving soft token distributions that are iteratively refined before producing discrete outputs.
  • Continuous trajectory supervision aligns optimization with iterative probabilistic refinement, while the method remains compatible with existing architectures, KV-caching, and blockwise diffusion.
  • EvoToken-DLM consistently improves over strong MDLM baselines across diverse backbones, datasets, and inference configurations.

2 Preliminaries on MDLMs

MDLMs generate by corrupting sequences into masked states and learning a reverse denoising process. Although predictions are made for all masked positions in parallel, only selected high-confidence tokens are finalized at each step, with others retained for later refinement.

  • MDLMs use a forward corruption process followed by a learned reverse denoising process.
  • The forward process gradually replaces tokens with a special mask token over a time schedule until the sequence is fully masked.
  • The reverse model predicts all masked tokens simultaneously, enabling parallel generation from the fully masked sequence toward the original text.
  • At each decoding step, MDLMs finalize a confidence-selected subset of masked tokens while retaining the rest for subsequent steps.
  • Sequences are partitioned into discrete blocks processed sequentially, and the next block begins only after the current block is fully refined.

3 From Discrete to Continuous: A Continuous Relaxation Perspective

EvoToken-DLM relaxes discrete token selection into a continuous embedding space of soft tokens represented by vocabulary distributions. Its transition function iteratively refines embeddings and auxiliary token states before mapping the terminal representation back to discrete outputs.

  • Continuous Relaxation: A soft token is a convex combination of vocabulary embeddings, obtained by multiplying the embedding matrix by a probability-simplex vector.
  • Iterative Refinement in Continuous Domain: The method models reverse diffusion as iterative refinement in the continuous embedding space rather than repeated prediction over the discrete vocabulary.
  • Iterative Refinement in Continuous Domain: A transition function recursively updates continuous embeddings and auxiliary token states until the process reaches terminal representations.
  • Iterative Refinement in Continuous Domain: The terminal continuous representations are mapped back to the discrete domain to produce the output sequence.

4 EvoToken-DLM

EvoToken-DLM replaces hard mask-to-token decisions with progressive soft-token evolution, then combines step-wise updates with blockwise decoding. Continuous trajectory supervision trains the model on the same iterative refinement behavior used during inference.

  • Progressive inference: Tokens transition through masked, soft-token, and decoded states as their representations are updated during inference.The state transitions preserve token distributions and allow selected tokens to advance while others continue refining.
  • Progressive inference: Step-wise token updates and blockwise decoding let tokens refine progressively within a block before the block is simultaneously decoded.Generation advances to subsequent blocks after all tokens in the current block reach the Soft(V) state.
  • Continuous trajectory supervision: Continuous trajectory supervision simulates consecutive refinement steps and applies supervised losses at every step rather than only at the final step.Each forward pass produces updated distributions, embeddings, and token states, while every step receives a backward pass.
  • Continuous trajectory supervision: The training objective is aligned with inference by explicitly modeling the trajectory of iterative probabilistic token refinement.This specialized training paradigm is needed because current states and embeddings depend on cumulative preceding refinements.
  • Blockwise diffusion: EvoToken-DLM extends to blockwise diffusion while preserving global autoregressive structure and enabling intra-block parallelism.Training adapts continuous trajectory supervision to simulate ∆τ refinement steps within each block.

5 Experiments

Experiments evaluate EvoToken-DLM across reasoning benchmarks, model backbones, inference strategies, and diffusion architectures. The method generally improves performance while retaining compatibility with KV-caching and incurring minimal latency overhead.

  • Main Performance Comparison: EvoToken-DLM improves average accuracy over the original LLaDA-Instruct-8B by 17.45% on Countdown, 3.08% on GSM8K, 2.06% on MATH500, and 3.23% on SVAMP at NFE Gen Len = 1.Table 1 compares the method with the original model and FT-baseline across multiple reasoning benchmarks.
  • Importance of Intermediate States: Removing intermediate refinement states reduces performance, supporting their role in gradual mask-to-soft-token prediction refinement.The ablation identifies gradual transition through intermediate states as essential for iterative prediction refinement.
  • Qualitative Visualization and Top-K Filtering: EvoToken-DLM shows progressive convergence from uncertain tokens to precise, coherent outputs while remaining robust across top-K settings.The inference visualization retains the top K = 3 tokens per position, while Figure 11 reports consistent performance across K values.
  • Compatibility with KV-Caching: EvoToken-DLM maintains superior performance over the baseline with KV-caching across computational budgets on Countdown.The integration is evaluated against the baseline under KV-caching in Table 2.
  • Robustness across Thresholds: With KV-caching and equal average tokens per step, EvoToken consistently outperforms binary masking across confidence thresholds on MATH500.Confidence thresholds replace fixed NFE settings to enable more flexible computational-resource allocation.
  • Generalization Across Models and Architectures: Continuous trajectory supervision generalizes improvements to Dream-Instruct-7B and blockwise D2F-LLaDA evaluations.The method is tested on an alternative pretrained backbone and on SVAMP with a blockwise diffusion model.
  • Inference Efficiency: EvoToken-DLM introduces negligible latency relative to standard MDLMs, with overhead primarily from element-wise token-embedding addition.The reported efficiency result characterizes the added computation as minimal.

6 Related Work

Related work develops diffusion language modeling, efficiency mechanisms, blockwise generation, and continuous or latent representations for reasoning. These directions address parallel decoding, iterative denoising cost, and information preservation during refinement.

  • Masked Diffusion Language Models: MDLMs adapt diffusion to discrete text generation, while LLaDA and Dream show that large-scale MDLMs can match autoregressive baselines in complex reasoning.The related-work discussion identifies masked diffusion as the dominant formulation and notes its reasoning capability.
  • KV-Cache Optimization: KV-cache research reduces redundant computation in bidirectional MDLMs through approximate caching, selective updates, and dynamic eviction.These approaches address the incompatibility between standard KV caching and bidirectional refinement.
  • Blockwise Diffusion Language Models: Blockwise diffusion combines autoregressive global ordering with intra-block diffusion to support KV-caching and reduce serial decoding bottlenecks.D2F additionally decodes future blocks from noisy intermediate states.
  • Reasoning in Continuous Space: Continuous-space reasoning methods use hidden states, projection modules, probability-weighted embeddings, or latent spaces to enhance chain-of-thought expressivity.The surveyed approaches replace or augment discrete token representations with continuous semantics.

7 Conclusion

EvoToken-DLM replaces rigid binary masks with evolving soft token distributions and trains this process with continuous trajectory supervision. Experiments report consistent gains over strong baselines alongside compatibility with KV-caching and blockwise architectures.

  • Conclusion: EvoToken-DLM replaces rigid binary masks with evolving soft token distributions for progressive decoding.The approach addresses irreversible discrete assignments in traditional MDLMs.
  • Conclusion: Continuous trajectory supervision aligns the training objective with iterative probabilistic refinement.The conclusion presents this supervision as the mechanism supporting progressive decoding.
  • Conclusion: Extensive experiments show consistent improvement over strong baselines while preserving compatibility with KV-caching and blockwise architectures.The reported conclusion covers both performance and integration across these inference settings.

Limitations

The paper’s main stated limitation concerns adapting EvoToken-DLM to autoregressive-initialized models. The mismatch between unidirectional pretraining and bidirectional iterative refinement increases training difficulty and slows convergence.

  • Limitations: Autoregressive-initialized backbones face increased training difficulty and slower convergence under EvoToken-DLM.The authors attribute this to the discrepancy between unidirectional AR pretraining and iterative bidirectional refinement.

B More Methodological Details

EvoToken-DLM uses progressive inference with soft token states and trains the model to follow this iterative trajectory. Evaluation spans arithmetic and reasoning benchmarks with specified parameter-efficient fine-tuning and simulation settings.

  • Progressive Inference: EvoToken-DLM’s progressive inference manages token states and continuous embeddings across multiple refinement passes within each block.Tokens begin masked; selected positions transition through soft states, while historical high-confidence predictions support final decoding.
  • Continuous Trajectory Supervision: Continuous trajectory supervision simulates inference during training by partitioning blocks into soft and masked subsets and refining them for Δτ steps.This trains parameters for the evolutionary path rather than only single-step recovery.
  • Evaluation Benchmarks: The evaluation covers Countdown, GSM8K, MATH500, SVAMP, and S1K, spanning arithmetic, logical, and reasoning difficulty.The supplied benchmark descriptions include combinatorial arithmetic, grade-school arithmetic, competition mathematics, linguistic robustness, and challenging reasoning problems.
  • Training Configuration: Training uses S1K for 10k steps with LoRA applied to query, key, and value projections.The LoRA configuration sets rank r = 128, alpha α = 256, and dropout 0.05.
  • Training Configuration: Continuous simulation uses Δτ = 4, block size 512, dynamically sampled transition-token counts, and α sampled from U(0.5, 1.0).The maximum sequence length is 1,024 tokens, and transition-token counts are sampled from {1, 2, 4, 8}.

C.3 Inference and Evaluation Setup

Inference is standardized across datasets with fixed temperature and seed, while α is selected by grid search. Generation length and NFEs vary by pretrained backbone.

  • Inference Setup: Inference uses temperature 0.5 and random seed 42 across datasets, with α selected from {0.5, 0.6, 0.7, 0.8, 0.9}.Performance is reported using the optimal α for each setting.
  • Evaluation Setup: Dream-Instruct-7B evaluations use generation length 256 with 128 NFEs, whereas D2F-LLaDA evaluations use maximum generation length 512.These settings are specified separately for the two pretrained models.

D.1 Paradigm Consistency: Why MDLMs Rapidly Adapt to EvoToken-DLM

MDLMs provide a compatible prior for EvoToken-DLM because they infer token distributions from context, whereas autoregressive backbones conflict with iterative full-context refinement. Experiments report generalizable gains and stable α sensitivity across benchmarks.

  • MDLM Paradigm: MDLM training masks tokens through a forward process and learns reverse denoising that iteratively predicts masked tokens as t decreases.At t = T the sequence is fully masked; the reverse process recovers the data distribution toward t = 0.
  • Paradigm Consistency: Despite supervision on mask tokens, MDLMs can infer distributions at all positions from surrounding context.Random token substitutions remain associated with ground-truth tokens among top-ranked predictions.
  • Paradigm Consistency: This contextual prediction capability lets EvoToken-DLM adapt from pretrained MDLMs with minimal supervised fine-tuning.The paper characterizes MDLM pretraining as a robust prior for real-time token refinement.
  • AR Adaptation Hurdles: Autoregressive backbones require substantial training resources because causal, unidirectional attention conflicts with iterative full-context token refinement.The authors therefore do not further explore AR-based variants.
  • Generalization and Robustness: EvoToken-DLM improves reasoning precision over binary masking on another pretrained model, and performance remains stable for α ∈[0.5, 0.9] across four datasets.The sensitivity analysis spans benchmarks from simple arithmetic to complex reasoning.

E.3 Additional Results with Different Block Sizes

EvoToken-DLM consistently outperforms the baselines across block sizes, tasks, and inference configurations. Additional visualizations show progressive refinement from uncertain intermediate representations to sharp, symbolically correct tokens.

  • Block-size robustness: Across block sizes {8, 16, 32, 64}, EvoToken-DLM consistently outperforms the original LLaDA-Instruct-8B and FT-baseline across tested datasets.Table S2 uses generation length 256 and NFE 128 after 10k-step fine-tuning with continuous trajectory supervision.
  • Block-size robustness: +10.55% to +13.67% absolute gains over the baseline are observed on Countdown across different block sizes.These results indicate improved reasoning performance across varying inference granularities.
  • Block-size robustness: EvoToken-DLM remains resilient to block-size changes, yielding stable and superior results across all tested discretization settings.The reported token evolution patterns are described as agnostic to specific block partitions.
  • Cross-task consistency: Performance superiority is maintained across symbolic reasoning and mathematical benchmarks, including Countdown, GSM8K, and MATH500.The results support consistency across diverse downstream tasks rather than specialization to one benchmark.
  • Cross-task consistency: The results indicate that EvoToken-DLM is not finely tuned to a specific inference setting but enhances the underlying diffusion generation process.This conclusion follows the reported stability across block sizes and tasks.
  • Qualitative refinement: Figures S3 and S4 visualize intermediate refinement states, with uncertain embeddings progressively becoming sharp, symbolically correct tokens as steps increase.The examples use block sizes 12 and 8 and showcase refinement for the first 16 output tokens.
Loading 2601.07351v2…