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

• URL: http://arxiv.org/abs/2601.22450v1
• Date: Fri, 30 Jan 2026 01:36:00 GMT
• Title: Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from $k$-Parity
• Abstract summary: Masked Diffusion Language Models have recently emerged as a powerful generative paradigm.<n>In this work, we investigate these properties within the setting of the $k$-parity problem.<n>By training nanoGPT using MD objective on the $k$-parity problem, we demonstrate that MD objective fundamentally alters the learning landscape.
• Score: 29.467658072159114
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion for High-Resolution Image Synthesis [93.7120829382642]
  Nemotron-Labs-Diffusion-Image is a state-of-the-art masked discrete diffusion model for text-to-image synthesis.<n>Standard MDMs lack self-correcting capability because discrete tokens cannot be modified once they are unmasked.<n>Nemotron-Labs-Diffusion-Image incorporates a token-editing mechanism that enables the model to dynamically revise already-unmasked tokens during inference.
  arXiv  Detail & Related papers  (2026-06-29T05:48:41Z)
• VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editing [20.078996384901966]
  Masked Diffusion Model (MDM-VGB) is a discrete diffusion sampler that augments unmasking generation with principled reward-guided remasking.<n>We prove that MDM-VGB is robust to process-verifier noise and achieves quadratic complexity, while popular test-times such as best-of-$N$ can incur exponential complexity due to error accumulation.
  arXiv  Detail & Related papers  (2026-06-26T17:47:09Z)
• Mixture-of-Experts Models in Vision: Routing, Optimization, and Generalization [0.0]
  We study MoE behavior in an image classification setting, focusing on predictive performance, expert utilization, and generalization.<n>We compare dense, SoftMoE, and SparseMoE classifier heads on the CIFAR10 dataset under comparable model capacity.<n>Both MoE variants achieve slightly higher validation accuracy than the dense baseline while maintaining balanced expert utilization through regularization.<n>We find that SoftMoE exhibits higher sharpness by these metrics, while Dense and SparseMoE lie in a similar curvature regime, despite all models achieving comparable generalization performance.
  arXiv  Detail & Related papers  (2026-01-21T14:22:25Z)
• Towards Efficient General Feature Prediction in Masked Skeleton Modeling [59.46799426434277]
  We propose a novel General Feature Prediction framework (GFP) for efficient mask skeleton modeling.<n>Our key innovation is replacing conventional low-level reconstruction with high-level feature prediction that spans from local motion patterns to global semantic representations.
  arXiv  Detail & Related papers  (2025-09-03T18:05:02Z)
• Unveiling the Potential of Diffusion Large Language Model in Controllable Generation [36.05635830551406]
  Even state-of-the-art autoregressive Large Language Models (LLMs) exhibit unreliability when required to generate structured output.<n>Inspired by the current new diffusion-based large language models (dLLM), we realize that the architectural difference may be the key to unlock next-level controllable generation.<n>We propose Self-adaptive Scaffolding, a novel framework that enables dLLM to stably generate reliable structured outputs.
  arXiv  Detail & Related papers  (2025-07-06T18:41:34Z)
• Supervised Optimism Correction: Be Confident When LLMs Are Sure [91.7459076316849]
  We establish a novel theoretical connection between supervised fine-tuning and offline reinforcement learning.<n>We show that the widely used beam search method suffers from unacceptable over-optimism.<n>We propose Supervised Optimism Correction, which introduces a simple yet effective auxiliary loss for token-level $Q$-value estimations.
  arXiv  Detail & Related papers  (2025-04-10T07:50:03Z)
• Boosting Large Language Models with Mask Fine-Tuning [60.56962908455601]
  We introduce Mask Fine-Tuning (MFT) to show that properly breaking the integrity of the model can surprisingly lead to improved performance.<n>Experiments show that MFT gains a consistent performance boost across various domains and backbones.
  arXiv  Detail & Related papers  (2025-03-27T20:17:57Z)
• Steering Masked Discrete Diffusion Models via Discrete Denoising Posterior Prediction [88.65168366064061]
  We introduce Discrete Denoising Posterior Prediction (DDPP), a novel framework that casts the task of steering pre-trained MDMs as a problem of probabilistic inference. Our framework leads to a family of three novel objectives that are all simulation-free, and thus scalable. We substantiate our designs via wet-lab validation, where we observe transient expression of reward-optimized protein sequences.
  arXiv  Detail & Related papers  (2024-10-10T17:18:30Z)
• Model Inversion Attacks Through Target-Specific Conditional Diffusion Models [54.69008212790426]
  Model inversion attacks (MIAs) aim to reconstruct private images from a target classifier's training set, thereby raising privacy concerns in AI applications. Previous GAN-based MIAs tend to suffer from inferior generative fidelity due to GAN's inherent flaws and biased optimization within latent space. We propose Diffusion-based Model Inversion (Diff-MI) attacks to alleviate these issues.
  arXiv  Detail & Related papers  (2024-07-16T06:38:49Z)
• Feature Attenuation of Defective Representation Can Resolve Incomplete Masking on Anomaly Detection [1.0358639819750703]
  In unsupervised anomaly detection (UAD) research, it is necessary to develop a computationally efficient and scalable solution. We revisit the reconstruction-by-inpainting approach and rethink to improve it by analyzing strengths and weaknesses. We propose Feature Attenuation of Defective Representation (FADeR) that only employs two layers which attenuates feature information of anomaly reconstruction.
  arXiv  Detail & Related papers  (2024-07-05T15:44:53Z)
• Cauchy-Schwarz Regularized Autoencoder [68.80569889599434]
  Variational autoencoders (VAE) are a powerful and widely-used class of generative models. We introduce a new constrained objective based on the Cauchy-Schwarz divergence, which can be computed analytically for GMMs. Our objective improves upon variational auto-encoding models in density estimation, unsupervised clustering, semi-supervised learning, and face analysis.
  arXiv  Detail & Related papers  (2021-01-06T17:36:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.