arXiv:2603.22241cs.CL2026-03

让语言模型学会逐步去噪,提升长文本理解与生成能力

MemDLM: Memory-Enhanced DLM Training

  • 用双层优化引入模拟去噪轨迹,构建可参数化的记忆通道
  • 训练损失更低,收敛更快,长文本表示能力显著增强
  • 适合需要精准长上下文理解的任务,如复杂检索

扩散语言模型(DLM)相比自回归模型具有全注意力并行解码和灵活生成的优势。但标准DLM训练采用静态单步掩码预测目标,未暴露模型于推理时的渐进去噪过程,且所有上下文信息仅靠令牌空间注意力维持,导致长序列下信息逐渐稀释。本文提出MemDLM(记忆增强型DLM),通过双层优化在训练中嵌入模拟去噪轨迹,内层更新一组快速权重形成参数化记忆,外层基于该记忆更新主模型。通过将部分记忆负担从令牌空间注意力转移至参数空间,MemDLM实现更快收敛、更强长上下文表征与更低训练损失,即使推理时丢弃快速权重亦有效。重新启用内层优化可实现提示级自适应,使参数化记忆成为应对挑战性‘大海捞针’任务的涌现式内存检索机制。

原文摘要 · Abstract (English)

Diffusion Language Models (DLMs) offer attractive advantages over Auto-Regressive (AR) models, such as full-attention parallel decoding and flexible generation. However, standard DLM training uses a static, single-step masked prediction objective that never exposes the model to the progressive denoising dynamics of inference, and forces all contextual information to be maintained purely through token-space attention, which becomes increasingly diluted as context length grows. We propose MemDLM (Memory-Enhanced DLM), which introduces a second memory channel by embedding a simulated denoising trajectory 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, while an outer loop updates the base model conditioned on this memory. By offloading part of the memorization burden from token-space attention to parameter space, MemDLM yields faster convergence, stronger long-context representations, and lower training loss, even when the fast weights are discarded at inference time. Re-enabling the inner loop at inference provides an additional prompt-specific adaptation effect, where the Parametric Memory acts as an emergent in-weight retrieval mechanism on challenging Needle-in-a-Haystack tasks. Code: https://github.com/JarvisPei/MemDLM.

扩散模型长文本生成记忆机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。