InfoDLM: an Information-Adaptive Framework for Discrete Diffusion Language Model Pretraining

Diffusion language models (DLMs) can match or surpass similarly sized autoregressive language models on language understanding and reasoning. However, their mask-and-denoise pretraining relies on heuristic random masking, which fails to target the most informative tokens. Consequently, the model spends significant computational effort on redundant or trivial tokens. To address this, we propose InfoDLM, an adaptive DLM pretraining framework that reformulates mask selection as an active, feedback-driven process. InfoDLM targets tokens that offer the highest measurable information gain during mask selection. Specifically, we: (1) introduce a Trainable Information-Gain (TIG) signal to quantify information gain of each masking configuration; (2) develop a feedback mechanism that adapts the masking policy to the model’s evolving state with a maturity indicator; and (3) jointly optimize the DLM and masking policy through an interleaved training flow with minimal computational overhead. Across reasoning-oriented benchmarks, InfoDLM achieves up to 13\% improvement in reasoning accuracy over a small variant of LLaDA under comparable pretraining budgets.

Large language models are trained on enormous amounts of text, but not every word is equally useful for learning. A word that is challenging for an early model may become obvious later, while another word may reveal exactly what the model still fails to understand. This means that the value of data depends on the current state of the model. We study this idea in diffusion language models, which learn by hiding parts of text and training the model to recover the missing words. In a diffusion language model, masking is more than a way of corrupting text before reconstruction, because it is the mechanism that turns raw text into a learning environment. The same sentence can teach the model different skills depending on which words are hidden, just as the same lesson can challenge different students in different ways. InfoDLM treats masking as an adaptive teaching process. As training continues, InfoDLM shifts what the model is asked to practice, focusing less on what the model already knows and more on what can still improve its reasoning. Our experiments show that this adaptive strategy improves language understanding and reasoning performance under comparable training settings. More broadly, our work suggests that improving language models is not only about using more data, but also about turning the same data into better learning challenges at the right time.