MemDLM: Memory-Enhanced DLM Training¶

会议: NeurIPS 2026
arXiv: 2603.22241
代码: https://github.com/Hao-Zhe-Shi/MemDLM
领域: 预训练
关键词: 扩散语言模型, 双层优化, 参数化记忆, 快速权重, 长上下文检索

一句话总结¶

本文提出 MemDLM 框架，通过双层优化在训练内循环中将模拟的逐步去噪轨迹写入快速权重构成参数化记忆，卸载注意力机制的局部记忆负担，在仅训练模式下大幅提升扩散语言模型的长上下文表征与检索能力，且可在推理时按需激活内循环实现基于权重的快速自适应检索。

研究背景与动机¶

离散扩散语言模型（DLMs）凭借双向全注意力并行解码与灵活的局部编辑能力，正成为自回归语言模型之外最具潜力的生成架构。然而，现有的掩码扩散语言模型（MDLMs）在训练范式与长上下文建模机制上面临两个根深蒂固的瓶颈。首先，标准 MDLM 采用静态单步掩码预测目标，在均匀采样的噪声时间步下仅根据真实上下文单步预测被掩码 token，使得模型在整个训练期从未经历推理阶段自上而下逐步去噪、动态依赖自身历史预测的迭代过程。这种暴露偏差造成训练目标与逐步推理动力学严重脱节，模型缺乏在带噪局部轨迹中持续校准和纠错的能力。

更棘手的矛盾在于长上下文下的注意力稀释（Attention Dilution）。标准的密集 Softmax 注意力要求将概率质量分配到上下文全序列的所有 token 上，随着序列长度显著增加，单个关键 token 所分配到的注意力权重急剧衰减，导致模型在中长距离信息查找与聚合时发生灾难性遗忘。由于现存标准 MDLM 强迫模型完全依赖 token 空间的注意力通路来维持所有任务相关的上下文线索，缺乏独立的记忆承载途径，只要序列一长或干扰项增多，模型就极难在“大海捞针”中锁定关键事实。

为了从根本上缓解这一矛盾，本文提出不能只在 token 空间被动分配注意力，而应开辟参数空间的第二记忆通道。如果能在训练期模拟逐步去噪轨迹，并将局部轨迹信息直接内化到轻量参数中，就能有效分担注意力机制的记忆压力。核心 idea：通过双层优化（Bi-level Optimization）将逐步去噪轨迹嵌入训练流程，内循环利用两阶段局部梯度更新快速权重作为参数化记忆，外循环在此记忆条件下更新基座模型，实现注意力与参数记忆的双通道协同。

方法详解¶

整体框架¶

MemDLM 的核心设计是将单一样本的训练过程重构为一个双层优化问题。系统将网络参数显式拆分为两部分：长期保留的基座模型参数 \(\theta\)（通过 LoRA 适配），以及每个 batch 临时动态生成的快速权重（Fast Weights）\(\phi\)（轻量级参数化记忆模块）。在一次前向更新中，模型首先在内循环中沿“预锚点 \(\to\) 锚点 \(\to\) 干净目标”的局部轨迹进行逐步去噪模拟，执行两步梯度更新，将局部上下文与过渡特征内化为参数化状态 \(\phi_2\)；随后在外循环中，基座模型在保持 \(\phi_2\) 条件化的前提下计算真实的锚点预测损失，并通过一阶梯度反向传播更新 \(\theta\)。

在多阶段去噪与参数协同流向中，整体架构包含输入加噪、锚点一致轨迹构建、两阶段快速权重累积、外循环记忆条件化预测以及基座参数更新。各阶段数据流向如下所示：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列与锚点构建<br/>干净序列 x0 采样时间步 t 生成锚点 xt"] --> B["阶段1：预锚点对齐<br/>xt 进一步加噪得 xt_pre，计算损失并更新 phi0 -> phi1"]
    B --> C["阶段2：锚点到目标演进<br/>基于 xt 预测目标 x0，计算损失并更新 phi1 -> phi2"]
    C --> D["阶段3：外循环记忆条件化预测<br/>基于固定的快速权重 phi2 计算锚点 xt 损失"]
    D --> E["阶段4：一阶梯度更新基座模型<br/>冻结内循环二阶导数，仅更新基座参数 theta"]

关键设计¶

1. 锚点一致轨迹模拟：在内循环中精确注入渐进去噪动态

标准单步训练无法让模型感知多步去噪的演进过程。为了在内循环中高效重建具有物理意义的推理轨迹，MemDLM 没有随意挑选离散时间步，而是设计了锚点一致的局部轨迹（Anchor-Consistent Local Trajectory）。内循环从零初始化的快速权重 \(\phi_0 = \mathbf{0}\) 出发，分为两个紧密耦合的阶段展开：第一阶段为“预锚点对齐（Pre-Anchor Alignment）”，在锚点状态 \(x_t\) 的基础上按比例 \(s_{\text{pre}}\) 进一步掩码得到更嘈杂的预锚点状态 \(x_{t_{\text{pre}}}\)（其中 \(t_{\text{pre}} > t\)），模型学习从 \(x_{t_{\text{pre}}}\) 向锚点状态 \(x_t\) 甚至干净目标恢复，并通过梯度反向传播将初始状态更新为 \(\phi_1\)；第二阶段为“锚点到目标演进（Anchor-to-Target）”，模型接收真实的锚点状态 \(x_t\)，在 \(\phi_1\) 的加持下预测最终干净序列 \(x_0\)，再次执行梯度更新得到最终的参数化记忆 \(\phi_2\)。通过这种双阶段递进式学习，\(\phi_2\) 真实刻画了局部轨迹在锚点附近的过渡动态，使扩散模型在训练步内就学会了跨时步协同。

2. 快速权重局部化约束：将参数记忆写入深层 FFN 模块

盲目增加快速权重的参数容量反而会导致内循环过拟合，削弱外循环基座的泛化泛用性。MemDLM 深入探讨了参数记忆写入的最佳结构位置与容量边界。实验表明，全参数更新或涉及全注意力层的快速权重往往因为表征过于发散而损害下游任务。因此，MemDLM 采用极其高效的结构约束方案：将快速权重 \(\phi\) 限制为仅作用于 Transformer 顶层最后 10% 层的门控前馈网络（FFN）模块中的轻量 LoRA 适配器（秩 \(r=32, \alpha=64\)），并通过单轮 SGD 配合局部参数级梯度裁剪（Gradient Clipping = 1.0）进行更新。这种设计使得参数化记忆充当了顶层语义特征的局部键值关联缓存，既将上下文信息强制固化到局域参数中、分担注意力的记忆负荷，又最大限度地降低了计算开销，保证每步训练耗时处于可控范围。

3. 外循环条件化与一阶更新：实现注意力与参数记忆的分流优化

当快速权重 \(\phi_2\) 吸收了批次样本特定的局部轨迹经验后，外循环需要在该记忆的先验条件下训练基座模型 \(\theta\)。外循环采用与标准 MDLM 形式一致的目标函数，但核心差别在于网络前向传播受 \(\phi_2\) 显式调制。在优化反向传播时，如果计算完整的二阶 Hessian 矩阵，显存和计算复杂度将呈指数级爆炸，完全无法工程落地。MemDLM 采用一阶近似（First-Order Approximation），在计算外循环损失对基座参数 \(\theta\) 的梯度时，阻断内循环更新梯度对 \(\theta\) 的反向图回传，将 \(\phi_2\) 视为当前样本的静态条件特征：

\[\theta \leftarrow \theta - \beta \nabla_{\theta} \left( \omega(t) \sum_{i \in \mathcal{M}_t} -\log p_{\theta, \phi_2}(x_0^i \mid x_t) \right)\]

这种分流解耦带来两个极其关键的优势：首先，由于局部样本的记忆信息已被快速权重分流吸收，外循环梯度 \(\nabla_{\theta}\) 不再需要强制基座模型通过全局注意力的细微权值去死记硬背每个局部 token，从而促使基座模型学到更纯粹、鲁棒的通用降噪表征；其次，即便在推理阶段将快速权重完全丢弃（Train-Only 模式），基座模型依然继承了在良好优化曲面上获得的平坦损失特征，展现出远超传统基线的长上下文处理能力。而在推理阶段按需开启内循环时，\(\phi_2\) 则自然转变为基于权重的检索器（In-Weight Retrieval），实现无需微调的瞬时提词自适应。

损失函数 / 训练策略¶

系统采用两阶段双层优化目标，内循环学习率为 \(\eta = 0.1\)，外循环优化器采用 AdamW，学习率为 \(\beta = 2 \times 10^{-5}\)，配合 10% 预热阶段的余弦衰减调度。

内循环两阶段快速权重累积过程形式化为：

\[\mathcal{L}_{\text{inner}}^{(1)} = \sum_{i \in \mathcal{M}_{t_{\text{pre}}}} -\log p_{\theta, \phi_0}(x_t^i \mid x_{t_{\text{pre}}}), \quad \phi_1 = \phi_0 - \eta \nabla_{\phi} \mathcal{L}_{\text{inner}}^{(1)}\]

\[\mathcal{L}_{\text{inner}}^{(2)} = \sum_{i \in \mathcal{M}_t} -\log p_{\theta, \phi_1}(x_0^i \mid x_t), \quad \phi_2 = \phi_1 - \eta \nabla_{\phi} \mathcal{L}_{\text{inner}}^{(2)}\]

外循环对基座模型参数 \(\theta\) 的整体优化目标为：

\[\mathcal{L}_{\text{MemDLM}}(\theta) = \mathbb{E}_{t \sim \mathcal{U}(0,1), x_0} \left[ \omega(t) \sum_{i \in \mathcal{M}_t} -\log p_{\theta, \phi_2}(x_0^i \mid x_t) \right]\]

其中时间步加权函数设为 \(\omega(t) = 1/t\)。训练在 LongAlpaca 指令微调集上进行，采用非对称掩码策略（Prompt 严格保留不掩码，仅对 Response 部分加噪计算损失）。

实验关键数据¶

主实验¶

主实验在 LLaDA-MoE-7B-A1B-Base 和 LLaDA2.1-mini 两大主流离散扩散基座上展开，评测基准涵盖长上下文“大海捞针”与变量追踪基准 RULER（多值检索 MV、变量追踪 VT、高频词抽取 CWE）以及复杂多跳推理长文本基准 BABILong。评估了标准 MDLM、仅在训练期使用快速权重的 MemDLM (Train-Only) 以及推理期同时开启快速权重的 MemDLM (Train & Inference)。

基座模型	方法配置	RULER-MV (4K / 8K)	RULER-VT (4K / 8K)	RULER-CWE (4K / 8K)	BABILong (4K / 8K)
LLaDA-MoE	Standard MDLM	87.90 / 92.50	99.40 / 78.84	69.10 / 55.18	45.60 / 41.80
LLaDA-MoE	MemDLM (Train-Only)	95.85 / 96.95	99.52 / 94.84	75.72 / 58.20	50.60 / 45.60
LLaDA-MoE	MemDLM (Train & Inference)	96.80 / 97.00	99.60 / 95.80	83.70 / 61.96	51.20 / 47.00
LLaDA2.1	Standard MDLM	98.65 / 98.00	59.32 / 43.72	72.34 / 67.08	48.80 / 47.40
LLaDA2.1	MemDLM (Train-Only)	99.65 / 99.00	88.56 / 59.15	78.96 / 68.02	58.50 / 55.00
LLaDA2.1	MemDLM (Train & Inference)	100.00 / 99.90	90.16 / 60.72	80.96 / 69.12	61.20 / 57.00

在跨长度外推测试中（LLaDA-MoE 在 16K 和 32K 极限序列），MemDLM 在 RULER-VT 16K 下取得 56.84%，显著优于标准基线的 52.56%；而在 LongBench 全面评测中，MemDLM 在 TriviaQA（87.77 vs 55.29）和 Passage Retrieval-zh（87.38 vs 74.50）上均取得巨大突破。

消融实验¶

针对内循环结构设计的消融分析以 LLaDA-MoE 在 BABILong-1K 任务上的准确率及训练开销为验证标准：

配置	核心改动说明	训练墙上时间 (h)	训练开销对比	BABILong-1K 准确率
Full MemDLM (Default)	2步双阶段更新，后10%层FFN	0.97	2.1×	0.684
w/o Pre-Anchor Stage	仅保留锚点到目标单步更新	0.76	1.6×	0.646
w/o Anchor-to-Target	仅保留预锚点对齐单步更新	0.81	1.7×	0.620
Multi-Step Pre-Anchor (4-step)	增加预锚点迭代至4步	1.46	3.1×	0.590
Full-Parameter Inner Loop	快速权重扩展为全模型参数	1.68	3.6×	0.602
Broad Scope (FFN 50% layers)	快速权重扩展到后50%层FFN	1.31	2.8×	0.574
FFN + Attention (10% layers)	快速权重同时作用于FFN与注意力	0.99	2.1×	0.648
Standard MDLM Baseline	无内循环，标准单步静态掩码训练	0.47	1.0×	0.538

关键发现¶

Train-Only 模式贡献了绝大部分增益：实验最突出的发现是，在推理期完全丢弃快速权重、没有任何额外计算开销的 Train-Only 配置下，模型在 RULER-VT 8K 上就从 78.84% 跃升至 94.84%。这强力证明了双层优化的主要红利来自于训练期对基座模型表征能力的重构，而非单纯依赖测试时外挂适配器。
参数记忆结构“少即是多”：内循环过度优化与过大容量不仅增加显存负担，还会损害下游泛化。将快速权重限制在顶层 10% 的 FFN 表现最佳（0.684），全参数适配反而降至 0.602，表明深层前馈网络更适合作紧凑的局部键值信息存储体。
锚点一致轨迹对齐不可或缺：如果将内循环替换为与外循环锚点脱钩的任意掩码轨迹，优化损失虽然下降，但 BABILong 得分暴跌至 0.604；双阶段中缺少任意一环都会造成显著掉点（0.646 和 0.620），验证了“预锚点对齐 \(\to\) 锚点到目标”双步演进的合理性。

亮点与洞察¶

参数空间记忆分流机制：巧妙地将神经系统“快速权重”经典理论引入现代扩散语言模型，通过构建双层优化解耦全局通用注意力与局部瞬时记忆，直击长上下文注意力稀释的痛点。
高效的一阶近似工程闭环：通过阻断内循环的高阶梯度计算，避免了灾难性的 Hessian 矩阵运算，使得双层优化仅引入约 2.1× 的训练时间开销，具备极强的工程落地可行性。
训练表征强化与推理检索自适应的统一：既能在零推理延迟代价下提升基座泛化性能，又能根据需要平滑切换到推理自适应模式，将快速权重转化为动态检索器。

局限与展望¶

训练与推理延时开销：内循环引入了约 2.1× 的训练墙上时钟延迟；若在推理时开启内循环，也会引入额外的 Prompt 编码等待时间，在超低时延服务场景中受限。
解码阶段自适应尚未覆盖：当前推理期快速权重自适应仅在 Prompt 阶段执行单次更新，在自回归或块级逐步生成解码过程中尚未实现动态步进更新。
验证范围有待拓宽：目前实验主要在预训练完成后的指令微调阶段及特定大小（7B/16B）的 DLM 上验证，未来需进一步探索其在大规模无监督预训练阶段的应用潜力。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将双层优化与快速权重参数化记忆引入扩散语言模型，理论视角独到。
实验充分度: ⭐⭐⭐⭐⭐ 在两大多模/MoE扩散基座上完整覆盖检索、长上下文、外推与详尽消融。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻，数学建模严谨，图表分析透彻。
价值: ⭐⭐⭐⭐⭐ 为解决长文本上下文稀释和离散扩散模型训练偏差提供了极具启发性的全新范式。