NGM：一种即插即用、无需训练的大语言模型记忆模块

# NGM：一种面向LLM的即插即用免训练记忆模块

来源：https://arxiv.org/html/2605.16893

屈宇闻* 南京大学 [email protected]  
董文辉* 南京大学 [email protected]  
司辰阳 南京大学 [email protected]  
单彩峰† 南京大学 [email protected]

###### 摘要

近期研究提出了条件记忆模块，将知识存储与神经计算解耦，从而实现更直接的知识访问。与依赖动态计算路径的MoE相比，显式查找提供了一种更高效的知识检索机制。然而，这些方法仍然依赖于学习得到的记忆嵌入，需要额外训练且限制了灵活性。为解决此问题，我们提出**N元记忆（N-gram Memory, NGM）**，一个免训练、即插即用的模块，由**因果N元编码器**和**余弦门控记忆注入器**组成。因果N元编码器直接对主干模型的预训练词元嵌入进行平均，构建N元表示，从而无需从头训练独立的N元嵌入。该设计既不需要额外的记忆表，也不需要检索流水线。余弦门控记忆注入器随后使用非参数化的余弦门控（带ReLU）将检索到的嵌入调制到上下文表示中。我们在0.6B到14B的Qwen3系列模型上，跨八个基准评估NGM。NGM将平均性能提升了0.5到1.2个点，在代码生成和知识密集型任务上尤其明显（例如，Qwen3-14B在LiveCodeBench上+3.0，在GPQA上+3.03）。此外，NGM在多模态基准中也提升了性能（例如，Qwen3-VL-2B在MMStar上+1.53）。代码可在 https://github.com/PioneerQyw/NGM 获取。

11脚注：同等贡献。
22脚注：通讯作者。

## 1 引言

基于Transformer的大语言模型（LLMs）[39 (https://arxiv.org/html/2605.16893#bib.bib2)] 提供了强大的上下文建模和语义推理能力，然而语言建模融合了两种性质不同的需求：动态组合计算与局部、静态、刻板模式的复用[15 (https://arxiv.org/html/2605.16893#bib.bib36), 10 (https://arxiv.org/html/2605.16893#bib.bib37)]。命名实体、重复标识符、单位、术语和公式化短语往往不像是需要深度推理的问题，而更像是可以通过廉价查找恢复的模式[4 (https://arxiv.org/html/2605.16893#bib.bib38), 29 (https://arxiv.org/html/2605.16893#bib.bib39), 33 (https://arxiv.org/html/2605.16893#bib.bib40)]。然而，标准的Transformer缺乏针对这种局部词汇和符号依赖的原生知识查找原语，迫使LLMs在推理时通过注意力和前馈计算来重建它们[40 (https://arxiv.org/html/2605.16893#bib.bib35), 8 (https://arxiv.org/html/2605.16893#bib.bib1)]。

查找式记忆为将静态模式复用与动态Transformer计算分离提供了一种自然方式[24 (https://arxiv.org/html/2605.16893#bib.bib15), 40 (https://arxiv.org/html/2605.16893#bib.bib35)]。近期工作通过引入显式学习的记忆组件探索了这一方向：Engram[8 (https://arxiv.org/html/2605.16893#bib.bib1)] 使用学习的N元查找表和上下文相关的门控机制形式化**条件记忆**，而嵌入缩放方法通过额外的词元级或N元嵌入参数扩展容量[42 (https://arxiv.org/html/2605.16893#bib.bib41), 38 (https://arxiv.org/html/2605.16893#bib.bib43), 28 (https://arxiv.org/html/2605.16893#bib.bib42), 12 (https://arxiv.org/html/2605.16893#bib.bib44)]。这些方法表明，局部查找是改进语言模型的一个有用维度，但它们是通过额外的可训练参数、专用训练以及在某些情况下专门的存储或检索基础设施来获得这种好处的。

这引出了我们的核心研究问题：

> *已经训练好的LLMs能否在不重新训练或添加学习的记忆表的情况下恢复有用的局部记忆益处？*

典型的查找式记忆流水线首先构建训练的N元嵌入，检索稀疏的相关记忆条目，然后通过上下文感知的门控将检索到的记忆与隐藏状态融合。该流水线的第一个障碍是需要训练一个独立的N元嵌入空间。我们转而探究是否可以复用主干模型已经训练好的词元嵌入：通过在局部因果窗口内平均预训练词元嵌入，我们无需引入任何新的记忆表即可获得N元特征。

只有当聚合后的N元特征与模型的隐藏状态保持兼容时，这种简单构造才有用。如图1 (https://arxiv.org/html/2605.16893#S1.F1) 所示，N元嵌入与Qwen3-8B隐藏状态的对齐程度，在各个深度上都显著强于位置打乱的N元对照和随机词元对照。在两个默认注入层，实际平均余弦相似度分别为0.312和0.137，而打乱对照为0.172和0.084，随机对照为0.014和0.008。这表明非参数化聚合的N元嵌入可以通过免训练的余弦门控直接与隐藏状态融合。

参见图说明Figure 1: Qwen3-8B模型中隐藏状态与聚合N元嵌入的对齐程度。

受此观点启发，我们提出**NGM（N元记忆）**，一个免训练、即插即用的模块，将局部N元信号注入冻结的仅解码器LLMs。关键思想是将预训练嵌入空间不仅视为输入接口，也视为可复用局部记忆的轻量级来源：如果邻近词元形成稳定的词汇、符号或短语级模式，它们的聚合嵌入可能提供有用的线索，解码器可以复用而无需完全通过更深的Transformer计算重建。如图2 (https://arxiv.org/html/2605.16893#S3.F2) 所示，NGM通过两个非参数化组件实现这一思想：**因果N元编码器**和**余弦门控记忆注入器**。给定输入序列，因果N元编码器通过聚合主干模型预训练词元嵌入在局部尾随窗口内构建因果多尺度N元表示，从而无需学习单独的记忆条目即可捕获不同粒度的局部模式。余弦门控记忆注入器随后使用ReLU过滤的余弦门控比较这些源于输入的N元表示与解码器隐藏状态，并通过带缩放残差连接写入得到的记忆更新，从而仅注入正对齐的局部记忆信号。该设计从实践和分析角度都有意义。实践中，它可以附加到已经训练好的LLMs上，无需额外参数、外部知识源或检索基础设施。从分析角度，它提供了一种受控方式，用于测试预训练嵌入空间是否已包含可利用的局部记忆结构，从而改进生成。

我们在从0.6B到14B的Qwen3模型上跨八个基准进行评估，涵盖数学、代码、知识和对齐。在所有测试规模上，NGM一致地将平均分数提升了+0.5到+1.2个点，在代码生成和几个知识密集型基准上观察到最显著的提升，例如Qwen3-14B在LiveCodeBench上+3.0，在GPQA上+3.03。此外，我们将方法扩展到多模态任务。在Qwen3-VL-2B上的结果表明，仅将NGM应用于语言解码器就能改进所有报告的基准，展示了我们方法的一定程度的通用性。

## 2 相关工作

#### 条件记忆与嵌入缩放。

经典的N元模型通过固定阶马尔可夫假设捕获短程统计[25 (https://arxiv.org/html/2605.16893#bib.bib31), 7 (https://arxiv.org/html/2605.16893#bib.bib32)]，而局部词汇模式具有强预测结构这一见解在神经时代仍然相关[32 (https://arxiv.org/html/2605.16893#bib.bib33), 2 (https://arxiv.org/html/2605.16893#bib.bib34)]。混合专家（MoE）模型通过条件计算扩展容量[36 (https://arxiv.org/html/2605.16893#bib.bib3), 16 (https://arxiv.org/html/2605.16893#bib.bib4)]；条件记忆则探索了基于查找的互补稀疏性轴。最近，一系列工作重振了这一直觉，称为**嵌入缩放**，将N元或词元级嵌入表视为用于扩展模型容量的专用参数轴。SCONE[42 (https://arxiv.org/html/2605.16893#bib.bib41)] 训练一个辅助Transformer来生成上下文相关的N元嵌入，但依赖于引入额外训练FLOPs的辅助编码模型；L3[38 (https://arxiv.org/html/2605.16893#bib.bib43)] 通过静态路由将分词器嵌入表泛化到解码器层，但仍需要学习的逐层聚合矩阵和CPU卸载存储；LongCat-Flash-Lite[28 (https://arxiv.org/html/2605.16893#bib.bib42)] 将基于哈希的N元嵌入扩展到超过300亿参数，要求大规模分布式训练和哈希表基础设施；MeKi[12 (https://arxiv.org/html/2605.16893#bib.bib44)] 注入重新参数化为静态查找表的词元级记忆专家，仍需要专用训练阶段来学习记忆库。关系最密切的是Engram[8 (https://arxiv.org/html/2605.16893#bib.bib1)]，它通过哈希N元查找、上下文感知门控和稀疏性分配框架形式化了**条件记忆**，并通过算法-系统协同设计扩展到270亿参数以进行深层注入。另一个相关线通过非参数化数据存储或对隐藏状态和外部语料库的检索来增强语言模型[24 (https://arxiv.org/html/2605.16893#bib.bib15), 40 (https://arxiv.org/html/2605.16893#bib.bib35), 19 (https://arxiv.org/html/2605.16893#bib.bib17), 3 (https://arxiv.org/html/2605.16893#bib.bib16)]；相比之下，NGM直接复用主干嵌入矩阵，不构建数据存储或检索索引。所有这些方法都有一个共同要求：训练专用的嵌入参数，并在大多数情况下需要专门的存储和检索基础设施。NGM在更严格的约束下重新审视了相同的直觉——它在推理时直接从主干现有的词元嵌入中构建因果多尺度N元表示，不需要额外训练、外部记忆表或专门基础设施。

#### 残差流对齐。

关于Transformer可解释性的工作已将残差流描述为后续计算的共享线性工作空间[14 (https://arxiv.org/html/2605.16893#bib.bib45)]。Logit Lens[34 (https://arxiv.org/html/2605.16893#bib.bib46)] 及后续探针[11 (https://arxiv.org/html/2605.16893#bib.bib48)] 表明，中间隐藏状态仍可通过*反嵌入*矩阵（语言建模头）部分投影到词汇空间。对于具有绑定嵌入的模型，这直接意味着与输入嵌入层的对齐；对于具有非绑定嵌入的模型（包括这里使用的Qwen3系列），则意味着间接对齐。因此，我们的**残差对齐**论点采用了一个更弱的、经验性前提：隐藏状态与*输入*嵌入空间保持足够的几何兼容性，使得余弦相似度可以用作有效的免训练门控信号。我们在§4.5 (https://arxiv.org/html/2605.16893#S4.SS5) 中验证了这一前提，其中隐藏状态与源于输入的N元嵌入之间的余弦相似度显著超过了打乱对照和随机对照。

## 3 方法

参见图说明Figure 2: NGM概览。因果N元编码器从主干模型的词元嵌入构建多尺度N元表示；余弦门控记忆注入器将它们与解码器隐藏状态进行评分，并将聚合的残差注入选定层。如图2 (https://arxiv.org/html/2605.16893#S3.F2) 所示，NGM是一个免训练的记忆模块，它直接从主干模型的词元嵌入矩阵中推导局部记忆信号，并通过非参数化余弦门控注入冻结的解码器表示。该模块包含两个组件。**因果N元编码器**仅使用预训练词元嵌入从输入序列构建多尺度局部记忆向量，而**余弦门控记忆注入器**测量它们与解码器隐藏状态的相似性，并通过残差连接将得到的记忆更新集成到主干中。因此，编码器指定了哪些局部信息可用作记忆，而注入器决定了这些信息何时应影响解码器。算法1 (https://arxiv.org/html/2605.16893#alg1) 总结了整个过程。在我们考虑的推理设置中，所有主干参数保持冻结，唯一的额外计算由当前输入序列和一小部分预定义的N元大小引起。

### 3.1 因果N元编码器

NGM的第一个组件是因果N元编码器，它仅使用主干模型的词元嵌入矩阵将输入前缀转换为多尺度局部记忆向量。令输入词元ID为X={x1,...,xT} 且主干词元嵌入矩阵为 E∈RV×d。词元嵌入为 et=E[xt]∈Rd。对于每个 n∈N（例如，{2,3}），我们首先用 (n−1) 个零向量左填充嵌入序列，形成 e~t：

e~t={0if2-n≤t≤0,etif1≤t≤T。(1)

然后，我们通过在一个包含 n 个词元的尾随窗口上进行平均池化，定义位置 t 处的**因果**N元表示：

gt,n=1n∑k=0n-1e~t-k。(2)

这使用了一个**嵌入袋（bag-of-embeddings）**近似：算术平均可以在不同粒度下捕获局部模式，而无需学习单独的记忆条目。得到的表示是有意在窗口内对顺序不敏感的，并且不旨在恢复完整的短语语义；相反，它提供了一个简单的局部摘要，可以在没有额外参数的情况下计算。左填充保持输出长度不变并确保因果性（位置 t 仅依赖于 ≤t 的词元）。对于多个窗口大小，我们将每个大小的向量堆叠成一个矩阵：

Gt=[gt,n]n∈N∈R|N|×d,(3)

然后由注入器消费（§3.2 (https://arxiv.org/html/2605.16893#S3.SS2.SSS0.Px1)）。实现中，左填充和因果平均池化通过F.pad后接核大小n、步长1的1D平均池化实现，完全可并行化。

算法1 NGM：N元记忆（*因果N元编码器* + *余弦门控记忆注入器*）
1: 词元ID X={x1,...,xT}; 隐藏状态 Hl={htl}