通过解耦扩散语言模型的隐式正则化,实现无需突变的快速泛化。
Tuning the Implicit Regularizer of Masked Diffusion Language Models: Enhancing Generalization via Insights from $k$-Parity
- 将扩散目标分解为信号与噪声两部分,揭示其学习机制
- 在80亿参数模型上提升5.8%精度,5000万参数模型困惑度显著降低
- 适合追求高效训练与强泛化的大型语言模型研究者
掩码扩散语言模型作为新兴生成范式,其泛化能力尚未得到充分研究。本文在k-奇偶性问题(计算k个相关位的异或和)背景下,理论分解掩码扩散(MD)目标为驱动特征学习的信号项和充当隐式正则化的噪声项。通过在纳米GPT上使用MD目标训练,发现该目标从根本上改变学习轨迹,实现快速且同步的泛化,避免了传统神经网络常见的‘领悟’现象(长期随机水平表现后突然泛化)。进一步基于理论洞察优化掩码概率分布,显著提升5000万参数模型的困惑度,并在从零预训练和监督微调中分别取得8.8%和5.8%的性能增益,验证了方法在大规模掩码扩散语言模型中的可扩展性与有效性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。