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

TL;DR

This paper introduces a density-aware masked data training paradigm for diffusion language models, improving reasoning accuracy by focusing on information-dense sequence parts and mitigating contextual collapse.

Contribution

It proposes an information density driven smart noise scheduler with complementary priority masking, enhancing diffusion LLM training efficiency and reasoning capabilities.

Findings

Achieves ~4% accuracy improvement on reasoning benchmarks

Effectively mitigates contextual collapse during training

Demonstrates the effectiveness of density-aware training strategies

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.