Mask Is What DLLM Needs: A Masked Data Training Paradigm for Diffusion LLMs

Linrui Ma, Yufei Cui, Kai Han, Yunhe Wang
3/16/2026
cs.LG

No confident code implementation found for this paper.

Abstract

Discrete diffusion models offer global context awareness and flexible parallel generation. However, uniform random noise schedulers in standard DLLM training overlook the highly non-uniform information density inherent in real-world sequences. This wastes optimization resources on low-density structural glues while leaving high-density logical pivot points severely under-optimized. To address this, we propose an Information Density Driven Smart Noise Scheduler. By extracting information-dense hubs and applying Complementary Priority Masking, our method decouples a single training instance into mutually reinforcing reasoning and syntax samples, forcing the model to master both logical deduction and foundational sequence structure. Experiments demonstrate that our approach improves average accuracy by ~4\% across four Code and Math reasoning benchmarks, significantly outperforming uniform baselines. Mechanistic analyses further reveal that probabilistic priority masking effectively mitigates contextual collapse during block diffusion training. Overall, this density-aware strategy efficiently unlocks the reasoning potential of diffusion language models at minimal annotation cost, emerging as a promising new masked data training paradigm for Diffusion LLMs. Our processed dataset can be found at https://huggingface.co/datasets/malr07/opc-sft-stage2-dense-extracted.

View on arXivView PDF

Code Implementations

No confident code match yet

We couldn't find an author-owned or strongly-evidenced community implementation for this paper. Any repos shown below are weak matches — verify before relying on them.

No code implementations found yet.

Know of an implementation? Let us know in the comments below!

Similar Papers

Cite this paper

@article{ma2026mask,
  title  = {Mask Is What DLLM Needs: A Masked Data Training Paradigm for Diffusion LLMs},
  author = {Linrui Ma and Yufei Cui and Kai Han and Yunhe Wang},
  year   = {2026},
  eprint = {2603.15803},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2603.15803v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2603.15803.svg)](https://codeofpaper.com/paper/2603.15803)

Discussion