Diffusion Language Models (DLMs) offer attractive advantages over Auto-Regressive (AR) models, such as full-attention parallel decoding and flexible generation. However, they suffer from a notable train-inference mismatch: DLMs are trained with a static, single-step masked prediction objective, but deployed through a multi-step progressive denoising trajectory. We propose MemDLM (Memory-Enhanced DLM), which narrows this gap by embedding a simulated denoising process into training via Bi-level Optimization. An inner loop updates a set of fast weights, forming a Parametric Memory that captures the local trajectory experience of each sample, while an outer loop updates the base model conditioned on this memory. By offloading memorization pressure from token representations to parameters, MemDLM yields faster convergence and lower training loss. Moreover, the inner loop can be re-enabled at inference time as an adaptation step, yielding additional gains on long-context understanding. We find that, when activated at inference time, this Parametric Memory acts as an emergent in-weight retrieval mechanism, helping MemDLM further reduce token-level attention bottlenecks on challenging Needle-in-a-Haystack retrieval tasks. Code: https://github.com/JarvisPei/MemDLM.


翻译:扩散语言模型相比自回归模型具有显著优势,例如全注意力并行解码与灵活生成。然而,这类模型存在训练-推理不匹配的显著缺陷:扩散语言模型通过静态单步掩码预测目标进行训练,但在实际部署中需要执行多步渐进式去噪轨迹。我们提出MemDLM(记忆增强型扩散语言模型),通过双层优化将模拟去噪过程嵌入训练过程来缩小这一差距。内循环更新一组快速权重,形成捕获每个样本局部轨迹经验的参数化记忆,外循环则基于该记忆条件更新基础模型。通过将记忆压力从词元表征转移到参数化载体,MemDLM实现了更快的收敛速度与更低的训练损失。此外,在推理阶段可重新启用内循环执行自适应步骤,从而在长上下文理解任务中取得额外性能提升。我们发现,当该参数化记忆在推理阶段激活时,会展现出新兴的权值内检索机制,助力MemDLM在具有挑战性的"大海捞针"检索任务中进一步突破词元级注意力瓶颈。开源代码:https://github.com/JarvisPei/MemDLM。

0
下载
关闭预览

相关内容

扩散语言模型综述
专知会员服务
19+阅读 · 2025年8月15日
Llama-3-SynE:实现有效且高效的大语言模型持续预训练
专知会员服务
36+阅读 · 2024年7月30日
「知识增强预训练语言模型」最新研究综述
专知
18+阅读 · 2022年11月18日
自然语言处理中的语言模型预训练方法
PaperWeekly
14+阅读 · 2018年10月21日
国家自然科学基金
6+阅读 · 2017年12月31日
国家自然科学基金
5+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
7+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
11+阅读 · 2012年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
VIP会员
最新内容
学习数据的几何:形状空间分析数学综述
专知会员服务
8+阅读 · 6月17日
定向能反无人机系统最新发展动态
专知会员服务
10+阅读 · 6月17日
从燃煤战舰到算法战争:水面指挥的永恒要求
专知会员服务
6+阅读 · 6月17日
相关VIP内容
扩散语言模型综述
专知会员服务
19+阅读 · 2025年8月15日
Llama-3-SynE:实现有效且高效的大语言模型持续预训练
专知会员服务
36+阅读 · 2024年7月30日
相关基金
国家自然科学基金
6+阅读 · 2017年12月31日
国家自然科学基金
5+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2015年12月31日
国家自然科学基金
7+阅读 · 2015年12月31日
国家自然科学基金
1+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
11+阅读 · 2012年12月31日
国家自然科学基金
18+阅读 · 2012年12月31日
Top
微信扫码咨询专知VIP会员