aiqbench analysis
This study demonstrates that optimizing the implicit regularizer in Masked Diffusion Language Models can eliminate the "grokking" phenomenon, enabling rapid generalization on tasks like the $k$-parity problem. By theoretically decomposing the training objective into signal and noise regimes, the authors achieve significant perplexity improvements across models ranging from 50M to 8B parameters. The findings offer a scalable framework for enhancing pre-training and fine-tuning efficiency without relying on auto-regressive architectures.
Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from $k$-Parity
arXiv:2601.22450v2 Announce Type: replace Abstract: 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.
Abstract: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.
Machine Learning (cs.LG); Artificial Intelligence (cs.AI)
arXiv:2601.22450 [cs.LG]
(or arXiv:2601.22450v2 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2601.22450 Focus to learn more arXiv-issued DOI via DataCite
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 278 words
Related coverage
Improving CoT Monitorability of Evaluation Awareness via Verbalization Training
LessWrong · Oct 2, 2026
AirLLM: Diffusion Policy-based Adaptive LoRA for Remote Fine-Tuning of LLM over the Air
arXiv · Aug 28, 2026
A Continuous-Time Reinforcement Learning Framework for Fine-Tuning Discrete Diffusion Models
arXiv · Sep 30, 2026
Masked Diffusion Language Agents for Tool-Integrated Chemical Reasoning
arXiv · Oct 12, 2026