学会记忆：基于混合邻居归纳记忆的半参数模型可扩展持续学习

# 学习记忆：基于混合邻居归纳记忆的半参数模型可扩展持续学习  
来源：https://arxiv.org/html/2303.01421  

彭广跃♡，葛涛♠，罗文♡，李伟♡，王厚峰♡†  
♡多媒体信息处理国家重点实验室，北京大学计算机学院  
♠微软  
{agy,wanghf}@pku.edu.cn, [email protected]  
[email protected], [email protected]  

###### 摘要  
半参数语言模型在各种自然语言处理任务中展现出潜力。然而，它们将非参数记忆用作静态存储，缺乏学习能力且与参数模型的内部信息流脱节，从而限制了可扩展性和效率。基于近期对语言模型的可解释性理论，我们将由 kkNN-LM 表示的非参数记忆重新概念化为一种可学习的混合邻居归纳记忆（MoNIM），该记忆协同了注意力头的归纳能力与前馈网络的记忆强度。通过集成到模型的信息流中，MoNIM 在 Transformer 架构中充当类似 FFN 的旁路层，从而能够有效学习新知识。大量实验表明，MoNIM 在数据维度和模型维度上均是一种保持性强且可扩展的持续学习者，提升了半参数语言模型的可扩展性和持续学习性能。¹代码公开于 https://github.com/viniferagy/MoNIM。  

## 学习记忆：基于混合邻居归纳记忆的半参数模型可扩展持续学习  

彭广跃♡††本文完成于作者在微软亚洲研究院实习期间，葛涛♠††通讯作者，罗文♡，李伟♡，王厚峰♡†  
♡多媒体信息处理国家重点实验室，北京大学计算机学院  
♠微软  
{agy,wanghf}@pku.edu.cn, [email protected]  
[email protected], [email protected]  

## 1 引言  

半参数语言模型因其照相式记忆能力和领域内准确性而受到越来越多的关注。通过将参数化神经模型与可扩展的非参数记忆相结合，它们擅长利用记忆辅助预测。然而，当前的半参数语言模型由于记忆管理策略低效，不适合我们快速变化的世界。它们通常将所有训练数据记录在静态记忆中，然后依赖额外模块或可调超参数搜索有用信息。缺乏学习能力阻碍了记忆压缩——这是语言模型高效学习的关键过程——并使记忆组件与模型的信息流分离。因此，随着数据或模型维度的增加，这些模型的内存使用和搜索时间呈线性增长。这种低效在大语言模型中尤其不切实际，因为它们处理海量训练数据。  

本文探讨关于半参数语言模型的两个问题：为什么它们强大但低效，以及我们能否构建高效的记忆策略？受解释大语言模型学习能力的研究启发，我们提出非参数记忆，特别是 kkNN-LM，本身就具有类似于多头自注意力层中归纳头的能力。完美的记忆带来完美的局部下一词归纳和良好性能，但全局视野的缺失阻碍了记忆的高效推理。  

参见图注。  
参见图注。  
图 1：MoNIM 的学习机制。(a) 归纳头（上）与 MoNIM（下）的类比。线条深浅表示分配的注意力分数。归纳头在上下文中同化相关信息，而 MoNIM 则收集学习历史中记忆的相似样本。(b) MoNIM 作为可学习模块。MoNIM 与最终 FFN 层共享相同的输入、工作流和输出。FFN 使用参数化键 n^L 匹配查询 x̃^L，以促进编码在参数化值 o^L（绿色）中的已学概念，而 MoNIM 则促进嵌入在邻近记忆值中的概念混合 c（带条纹的绿色）。MoNIM 的学习过程由为模型损失函数优化的压缩器 M̂_Θ(x) 控制。  

为了发挥局部优势同时避免全局弱点，我们基于前馈网络中概念提升的组件，构建了一个可学习的混合邻居归纳记忆（MoNIM）。作为类似 FFN 的旁路层，MoNIM 能够以逐渐减小的内存占用选择并吸收知识，这表明它与参数模型中的梯度下降一致——新数据的影响随着数据量增大或模型规模扩大而减弱。因此，MoNIM 的记忆随模型能力的增强而呈次线性增长，成为一个无需训练的可扩展持续学习者。在持续学习场景中，MoNIM 的性能与原始模型相当，而内存空间仅消耗一半。我们的贡献总结如下：  

- • 我们引入了混合邻居归纳记忆（MoNIM），一种在半参数持续学习设置中充当类似 FFN 旁路层的可学习记忆。  
- • MoNIM 通过整合多头自注意力的归纳能力与前馈网络的记忆功能来实现其卓越性能。  
- • 在语言建模和下游任务上的大量实验表明，MoNIM 能有效压缩已见信息，并且在数据和模型维度上均可扩展，因此适用于使用半参数语言模型在流式数据上进行持续学习。  

## 2 MoNIM：混合邻居归纳记忆  

### 2.1 预备知识：kkNN-LM  

形式上，我们用 Θ = (θ, M) 表示一个半参数语言模型，其中 θ 代表参数化语言模型，M 代表非参数记忆。作为代表，kkNN-LM 通过利用 M 中 k 个最近邻的信息来增强 θ 的预测。给定一个左侧上下文 x = (x₁, ..., x_t)，kkNN-LM 使用在 FFN 层 l ∈ {1, ..., L} 之前最终位置的隐藏状态作为上下文表示 x̃^l ∈ R^d，并如下计算其下一个词 y 的概率：  

P(y|x; Θ) = f( P(y|x; θ)⏟模型 , P(y|x̃^l; M)⏟记忆 , λ )   (1)  

其中 f 表示插值函数，通过 λ 权衡模型和记忆的预测。kkNN-LM 通过训练集 D 构建 M 作为键值查找表，训练集中每个标记 x_t（作为值）及其上下文 <t 的表示 x̃^l（作为键）作为一个条目。在推断时，给定查询 x̃^l，检索 N(x̃) ⊂ M 中的 k 个最近邻；根据这些邻居在训练集中出现的频率计算概率 P(y|x̃^l; M)。  

### 2.2 从归纳头到 kkNN 记忆  

最近的研究表明，归纳头可能构成大语言模型中上下文学习的基本能力。注意力头表现出两种典型特性：前缀匹配——关注具有相似上下文的标记；以及复制——增加与被关注标记对应输出的对数几率。图 1(a) 展示了归纳头与 kkNN 记忆之间的推理机制。它们的相似性显而易见：归纳头同化分散在上下文中的相关信息，而 kkNN 记忆则从 N(x̃) 收集历史中记忆的相似信息。记忆将前缀编码为键向量，用其匹配查询，并将记忆的值复制到待预测位置——这与归纳头的功能一致。从这个角度看，kkNN 记忆是一个归纳缓冲区，将所有训练数据视为邻居，因此得名混合邻居归纳记忆（MoNIM）。  

然而，归纳头作为早期层的基本组件，在模型中传递上下文信息以实现更复杂的全局推理路径；而 MoNIM 作为高度局部的信息源，仅提供记忆的标签。其归纳最相似标记的能力也暗示其无法进行高效预测。为了充分利用 MoNIM 的局部精确性，我们着手开发一种新的 MoNIM 概念架构。  

### 2.3 MoNIM 用于局部概念混合提升  

根据先前的工作，FFN 层充当键值记忆，值向量可以投影到词汇空间以表示可理解的概念，例如食物或电影角色。一个 FFN 更新 F(x̃) 因此可以被解释为对目标标记含义的连续概念提升：  

F^l(x̃^l) = Σ_{i=1}^{d_F} f(x̃^l · n_i^l) o_i^l = Σ_{i=1}^{d_F} m_i^l o_i^l  

n_i^l, o_i^l 是 W_in^T, W_out ∈ R^{d_F × d} 的第 i 列，m_i^l 表示 o_i^l 的权重，模型在其中存储了对概念的理解。类似地，对于 MHSA 层之前的隐藏状态 x̂^l，给定长度为 T 的上下文的注意力模式 a^l ∈ R^T 以及相应的 v_j^l（W_V ∈ R^{T × d_A} 的第 j 列），MHSA 更新 A(x̂^l) 是输出矩阵 W_O ∈ R^{d_A × d} 向量的线性组合：  

A^l(x̂^l) = Σ_{j=1}^{d_A} (a^l · v_j^l) o_j^l = Σ_{j=1}^{d_A} m_j^l o_j^l  

由于 Geva 等人已证明层归一化几乎是线性的且不影响 MHSA 和 FFN 输出的线性性质，我们断言在解嵌入矩阵 E ∈ R^{|V| × d} 之前的最终预测可以分解为从 MHSA 和 FFN 层获得信息的加权和：  

x̃^l = Σ_{i=1}^{l·(d_F+d_A)} m̃_i^{<l} o_i^{<l}  
ỹ = softmax(E x̃^L)  

ỹ 的预测由概念 o 的混合决定。最大 log P(x_t|x_<t; Θ) 比较预测与真实标记的差距，如果差距较大，则 δ_ada < δ，表明我们应积极记忆该样本，因为其与排名第一的预测存在较大差距。AMT 在实践中简单有效。它允许我们跳过许多排名靠前的样本，显著减少内存大小而生成质量损失极小；此外，它缓解了对排名靠前样本的过拟合，起到类似标签平滑的作用，避免过度自信的预测。  

由于 MoNIM 更新总是提升概念而非消除或走捷径，如果提取的邻居过于稀疏散乱以至于没有可靠的概念可提升，则其权重 λ 应相应降低。受先前研究启发，我们训练一个简单的校准器，使用三类特征告知 MoNIM 的可靠性：参数化语言模型的分布信息、训练数据的词汇信息以及 MoNIM 的密度信息。按照 He 等人的方法，我们使用一个 4 层 MLP 网络，在验证集的一小部分上优化。⁴⁴⁴ 校准器的详细实现见附录 A.3，消融研究见附录 C.3。  

### 3.3 MoNIM 的可扩展性  

与空间效率低下的 kkNN-LM 不同，MoNIM 的学习能力使其能够在整个学习过程中压缩和减少内存需求。我们在持续学习设置中设计了实验，表明 MoNIM 在更新时能持续压缩。进一步，我们揭示了压缩的特征与模型参数更新的特征一致，从而赋予 MoNIM 可扩展性：(i) 数据可扩展性：在参数模型中，随着学习的进行，数据对参数更新的影响逐渐减小；MoNIM 的内存增长也遵循类似趋势。