Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from $k$-Parity

Masked Diffusion Language Models have recently emerged as a powerful generative paradigm, yet their generalization properties remain understudied compared to their auto-regressive counterparts. In this work, we investigate these properties within the setting of the $k$-parity problem (computing the XOR sum of $k$ relevant bits), where neural networks typically exhibit grokking—a prolonged plateau of chance-level performance followed by sudden generalization. We theoretically decompose the Masked Diffusion (MD) objective into a Signal regime which drives feature learning, and a Noise regime which serves as an implicit regularizer. By training nanoGPT using MD objective on the $k$-parity problem, we demonstrate that MD objective fundamentally alters the learning landscape, enabling rapid and simultaneous generalization without experiencing grokking. Furthermore, we leverage our theoretical insights to optimize the distribution of the mask probability in the MD objective. Our method significantly improves perplexity for 50M-parameter models and achieves superior results across both pre-training from scratch and supervised fine-tuning. Specifically, we observe performance gains peaking at $8.8$% and $5.8$%, respectively, on 8B-parameter models, confirming the scalability and effectiveness of our framework in large-scale masked diffusion language model regimes.

Most AI language tools you've heard of, like ChatGPT, write by guessing one word at a time, always predicting what comes next. A newer approach called masked diffusion works differently: rather than going word by word, it learns by filling in words that have been blanked out anywhere in a sentence—a bit like solving a crossword. It's a promising idea, but we don't yet have a good understanding of how it actually learns. To find out, we trained models on a simple, well-controlled test problem and watched the learning unfold. Vanilla methods show a strange pattern called "grokking": the model looks stuck and clueless for a long stretch, then suddenly "gets it" all at once. We found that masked diffusion avoids this entirely. Instead of the long, frustrating plateau, it learns quickly and steadily. To understand why, we separated the method's learning process into two parts working together. One part teaches the model the real patterns in the data; the other quietly stops it from over-committing too early. That insight pointed us to a single setting worth tuning—how often words get blanked out during training—and showed us how to set it well. When we applied this change to real, large-scale language models, it paid off. Models trained our way produced and predicted text more accurately, and the benefit held up at scale: our largest models improved by up to roughly 6–9%. In short, by understanding how these models learn, we found a simple adjustment that makes them learn better.