Language Model Memory and Memory Models for Language¶

会议: NeurIPS 2026
arXiv: 2602.13466
OpenReview: https://openreview.net/forum?id=SCvNm4wa7n
代码: https://github.com/blbadger/memorymodels
领域: LLM 其他
关键词: 记忆模型, 语言模型可逆性, 信息保留量, 编码器-解码器架构, 联合目标训练

一句话总结¶

针对自回归语言模型隐层向量对输入信息保留极少且难以支持任意信息检索的痛点，本文提出了可并行化的编码器-解码器记忆模型，结合因果语言建模与复制保留联合训练目标，使模型在大幅降低推理计算复杂度的同时保留高保真记忆。

研究背景与动机¶

大语言模型通常依靠模型权重存储训练知识，并依赖隐层激活值捕获当前输入 token 序列的信息。为了大幅降低自回归解码阶段的计算开销、首 token 延迟（TTFT）以及 KV 缓存显存占用，学术界尝试使用紧凑的记忆向量（Memory Embeddings）来替代长文本序列。然而现有的记忆增强模型（如 RMT、Transformer-XL 变体）多局限于单次查表或特定关联检索，在通用的广泛语言理解任务上表现显著落后于全上下文 Transformer。

这一困境的核心矛盾在于：现代复杂语言任务本质上要求对输入上下文具备「任意信息的精确检索能力」，而标准的自回归因果语言模型在数学上是一个典型的多对一（Many-to-one）非单射映射函数。模型在生成单个下一 token 时通常仅需参考极少数的前序上下文信息，导致因果预训练过程并无驱动力在最后的隐藏状态中保留完整的输入信息；经验测定表明因果 LLM 的记忆保留度与未训练模型相差无几。相反，面向输入重构训练的自编码器（Autoencoder）却能够将输入信息近乎无损地压缩进极少量的隐层向量中。

本文的切入角度是打破因果生成与记忆保留的隔阂：既然因果目标难以自然孕育信息完备的向量记忆，就必须在架构上解耦记忆编码与生成解码，并通过专门的显式信息保留目标与课程学习引导解码器掌握记忆提取。核心 idea：设计一种可并行的分块编码器-解码器记忆模型架构，并提出结合因果预测与文本复制（Copy Task）的联合优化目标，利用高质量自编码器特征与课程训练使解码器学会从高压缩记忆向量中精确解码任意信息。

方法详解¶

整体框架¶

本文提出的记忆模型将长输入序列划分为 $s$ 个长度为 $n/s$ 的子块（Chunks）。编码器在各个 chunk 上完全并行执行，分别将各个子块压缩为单个隐层记忆向量；解码器则接收这 $s$ 个记忆向量作为前置上下文，并结合后续 token 序列进行因果生成。与传统的自回归循环记忆机制相比，这种并行编码极大缩短了训练与推理前向耗时。

在理想的分块设置下，针对长度为 $n$ 的序列，当取分块数 $s = n^{2/3}$ 时，编码器与解码器的总计算复杂度可从全上下文的 $\mathcal{O}(n^2)$ 降低至 $\mathcal{O}(n^{4/3})$；在完全并行的硬件前向下，编码阶段的有效时序复杂度进一步下降至 $\mathcal{O}(n^{2/3})$，显存带宽占用也成倍降低。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["长输入文本序列<br/>拆分为 s 个子块 (Chunk)"] --> B["并行子块编码器<br/>多分块独立并行前向压缩"]
    B --> C["信息保留度量与展开投影<br/>滑动窗口线性映射 P 展开特征维度"]
    C --> D["因果解码器<br/>固定位置填充拼接记忆向量"]
    D --> E["双任务联合目标优化<br/>因果语言建模 + 显式文本复制"]
    E --> F["输出：低延迟生成与任意上下文检索"]

关键设计¶

1. 基于可逆性与熵比率的信息保留度量

为了定量评测模型隐层向量中到底沉淀了多少输入信息，本文建立了一套独立于词表大小的模型可逆性度量方法。给定被测模型作为编码器 $\mathcal{E}$，将其隐层输出向量经由滑动窗口线性投影 $P: \mathbb{R}^d \to \mathbb{R}^{n \times d}$ 展开后送入结构匹配的解码器 $\mathcal{D}$，优化重构损失。本文定义了归一化熵比率 $H_r$ 与精确 token 匹配准确率 $h(x,y)$ 两个核心量化指标： $$H_r = 1 - \frac{H(p, q)}{H(p_0, q)} = 1 - \frac{-\sum_{x} q(x) \log p(x)}{-\sum_x q_0(x) \log p(x)}$$ 其中分母 $H(p_0, q)$ 取均匀分布 $\mathcal{U}(|t|)$（最大熵无信息基线）与目标 token 的平均交叉熵，对于 8k 词表该常数约为 $9.03$。若 $H_r \to 1$ 且准确率接近 $100\%$，表明向量几乎完美保留了全部输入信息；测试证实主流因果 LLM 的 $H_r$ 仅为 $0.3\sim 0.4$，而自编码器可达 $0.95$ 以上。

2. 并行化分块编码与固定槽位记忆注入

针对循环记忆模型无法并行推理的弊端，本文设计了分块独立编码机制。输入序列被切分为 $s$ 个互不重叠的 chunk，并行通过自编码器架构编码为 $s$ 个记忆向量。在送入解码器时，本文对比了变长紧凑拼接与固定槽位填充（Fixed Embedding Padding）两种策略，证实固定槽位对齐能让解码器中的自注意力头形成稳定的位置偏置，更便于跨 chunk 定位信息。此外，当编码器被预先训练并冻结时，不仅能消除反向传播显存开销，还能使多轮批处理吞吐量大幅提升。

3. 因果建模与文本复制的联合目标函数

直接用因果目标训练带记忆的解码器时，解码器倾向于走捷径（Shortcut），即完全忽略复杂的记忆向量，仅依赖自身局部因果上下文，导致记忆模块退化为噪声。为了强制解码器学会检索与解码记忆向量，本文引入了全序列复制任务（Copy Task）与无可见前缀的空白复制任务（Blank Copy Task），并构建联合损失函数： $$\mathbb{L}' = \mathbb{L}(O(X, \theta), X_s) + \mathbb{L}(O(X_c, \theta), X_{sc})$$ 式中 $X_s$ 表示标准下一个词预测的目标位移输入，$X_c$ 表示拼接了特殊记忆分隔符的前半段复制序列，$X_{sc}$ 为对应的复制重构监督。在推理时无需分隔符即可无缝回归标准生成，但此时注意力权重已学会根据需求主动从记忆槽位提取细粒度知识。

4. 冻结高保真编码器的课程学习解码

实验发现，即使为模型提供由自编码器生成的完美记忆向量，若从一开始就直接进行因果生成训练，解码器由于初始梯度剧烈震荡仍会忽略记忆特征。为此本文提出了两阶段课程学习（Curriculum Learning）：第一阶段执行 Blank Copy 训练，将所有前序 token 掩码，强迫解码器仅依靠记忆向量重建原序列；第二阶段再放开 token 上下文并接入因果生成联合目标。该策略成功使 Llama 3.2 等预训练模型顺利适配记忆向量输入，打破了局部最优陷阱。

损失函数 / 训练策略¶

模型预训练基于 FineWeb-edu 子集，使用 AdamW 优化器，设置 500 step 线性预热，Transformer 峰值学习率设为 $2 \times 10^{-4}$，Mixer 设为 $5 \times 10^{-4}$，随后进行线性衰减。常规训练全局 batch size 为 128（$n_{ctx}=512$），长上下文 copy 任务设置 batch size 为 64（$n_{ctx}=1024$）。当引入外部预训练大模型（如 Llama 3.2 1B）作为解码器时，需要使用更小的学习率以避免反向传播梯度发散。

实验关键数据¶

主实验¶

论文首先对比了不同模型架构与预训练目标下，单个隐藏层向量重构输入文本的信息保留度（上下文长度 $n_{ctx}=512$）。如表所示，传统因果语言模型与检索模型的信息保留极为有限，而自编码器能够形成高保真记忆。

模型类型与架构	损失值 (Loss)	熵比率 ($H_r$)	恢复准确率 (Accuracy)
Autoencoder (Mixer, $n_{ctx}=256$)	0.010	0.999	99.60%
Autoencoder (Transformer, $n_{ctx}=256$)	0.090	0.990	95.40%
Autoencoder (Mixer, $n_{ctx}=512$)	0.292	0.968	91.84%
Autoencoder (Transformer, $n_{ctx}=512$)	2.730	0.698	43.30%
Causal LM (Mixer, $n_{ctx}=512$)	5.815	0.356	5.48%
Causal LM (Transformer, $n_{ctx}=512$)	6.214	0.310	5.01%
Causal $\to$ Retrieval (Mixer, $n_{ctx}=512$)	5.594	0.381	5.45%
Causal $\to$ Retrieval (Transformer, $n_{ctx}=512$)	6.380	0.290	5.29%
Untrained Baseline (Transformer, $n_{ctx}=512$)	6.643	0.264	4.67%

在下游长程与信息检索基准测试中，基于冻结编码器训练的记忆 Transformer 在联合训练目标（Combined）下，相比纯因果训练（Causal）表现出显著优势：

评估基准类别	具体任务 / 数据集	纯因果目标 (Causal)	联合优化目标 (Combined)	提升幅度
信息密集型评测	LongBench - Code Completion	13.17 ± 0.31	16.48 ± 0.33	+3.31
信息密集型评测	LongBench - Few-Shot Learning	0.80 ± 0.11	1.57 ± 0.12	+0.77
信息密集型评测	LongBench - Single Doc QA	0.81 ± 0.07	1.07 ± 0.08	+0.26
信息密集型评测	LongBench - Summarization	2.89 ± 0.11	3.41 ± 0.12	+0.52
信息密集型评测	SWDE (Contains)	0.00	0.81	+0.81
常规语言理解	ARC-Easy (Accuracy %)	44.20 ± 1.02	43.73 ± 1.02	-0.47
常规语言理解	WikiText (Bits Per Byte)	1.6752	1.6850	+0.0098
常规语言理解	HellaSwag (Accuracy %)	27.71 ± 0.45	27.52 ± 0.45	-0.19

消融实验¶

论文评估了在全参数可训（All Trainable）与冻结编码器（Frozen Encoder）条件下，不同训练目标对因果预测（CLM Eval）和长文本复制（Copy Eval）性能的消融影响（模型配置为 $d_m=512, n_l=16, n_{ctx}=256, s=4$）：

编码器设定	训练目标配置	CLM $H_r$	CLM Accuracy	Copy $H_r$	Copy Accuracy
All Trainable	纯因果目标 (CLM)	0.684	43.15%	0.688	43.65%
All Trainable	联合目标 (Combined)	0.675	42.07%	0.947	86.10%
All Trainable	联合目标-空白复制 (Combined Blank)	0.669	41.50%	0.756	50.60%
Frozen Encoder	纯因果目标 (CLM)	0.677	42.32%	0.677	42.63%
Frozen Encoder	联合目标 (Combined)	0.664	40.80%	0.927	82.70%
Frozen Encoder	课程学习+联合目标 (Curriculum Blank)	0.660	40.40%	0.759	52.60%

关键发现¶

因果目标天然抑制信息保留：纯因果预训练的向量在复制任务上的准确率仅为 42%~43%，这与没有输入记忆时单凭局部上下文推测下一个 token 的 baseline 完全一致，说明模型在没有额外约束时几乎不将非当前急需的历史信息压缩进隐层。
联合目标带来跨越式记忆提取能力：引入 Copy 联合目标后，Copy 准确率从 43.65% 跃升至 86.10%，且在常规因果语言生成上仅损失约 1% 的精度，证实模型能同时兼顾局部因果连贯性与全局历史记忆检索。
模态流形冲突（Manifold Mismatch）：Masked Mixer 虽在编码阶段具有极高的压缩能效，但其生成的隐层流形与 Transformer 解码器高度不兼容，强行跨架构拼接反而会降低训练效率；同构架构配对表现最稳定。

亮点与洞察¶

首次从函数映射的非单射性（Many-to-one）角度严格解释了传统自回归因果语言模型难以胜任记忆压缩的本质机理，打破了「增大模型尺寸和训练算力即可自然涌现记忆回溯」的盲目预期。
提出了理论完备且不受词表规模影响的熵比率归一化指标 $H_r$，为不同架构（Transformer vs. Mixer）与不同规模模型的信息保留能力提供了统一的可度量衡基准。
巧妙利用两阶段课程学习（Blank Copy $\to$ Copy），成功引导解码器克服局部极小值捷径，掌握了从高度抽象且冻结的向量记忆中解压缩长序列的技术路径。

局限与展望¶

基准性能有待进一步放缩：当前主要实验基于中小规模自研模型（512 维）与 1B 级模型（Llama 3.2、Qwen 2.5/3 0.6B），超大参数规模（如 70B+）下多层注意力与记忆交互的动力学行为仍待深入探索。
自然任务中的主动检索机制不足：目前联合目标中依赖合成的复制标记与分隔符显式切换模式，未来应探索如何在无需人造分隔符的前提下，让模型根据自然语言指令自主决定何时深层提取记忆向量中的细粒度细节。

评分¶

新颖性: ⭐⭐⭐⭐ [从信息论与映射可逆性角度切入解释记忆失效，提出行之有效的混合目标与课程训练]
实验充分度: ⭐⭐⭐⭐ [涵盖了不同架构、不同规模模型、自编码重构、复制评测及 LongBench 等丰富实验]
写作质量: ⭐⭐⭐⭐⭐ [逻辑链条极为清晰，理论推导、问题剖析与解决手段层层递进]
价值: ⭐⭐⭐⭐ [对超长文本压缩、端侧高效推理和下一代记忆增强大模型架构设计具有极高参考价值]

模型类型与架构	损失值 (Loss)	熵比率 (\(H_r\))	恢复准确率 (Accuracy)
Autoencoder (Mixer, \(n_{ctx}=256\))	0.010	0.999	99.60%
Autoencoder (Transformer, \(n_{ctx}=256\))	0.090	0.990	95.40%
Autoencoder (Mixer, \(n_{ctx}=512\))	0.292	0.968	91.84%
Autoencoder (Transformer, \(n_{ctx}=512\))	2.730	0.698	43.30%
Causal LM (Mixer, \(n_{ctx}=512\))	5.815	0.356	5.48%
Causal LM (Transformer, \(n_{ctx}=512\))	6.214	0.310	5.01%
Causal \(\to\) Retrieval (Mixer, \(n_{ctx}=512\))	5.594	0.381	5.45%
Causal \(\to\) Retrieval (Transformer, \(n_{ctx}=512\))	6.380	0.290	5.29%
Untrained Baseline (Transformer, \(n_{ctx}=512\))	6.643	0.264	4.67%