Meta AI 可扩展内存层-尧图网络科技

摘要：本文围绕 Meta AI 提出的可扩展内存层展开长文拆解，从 Transformer 的稠密记忆瓶颈出发，系统介绍产品键查找机制、键值记忆块、稀疏激活训练、分布式工程实现、实验评估与产品级落地思路。文章同时给出可复现的简化实现和调参建议，帮助读者理解“把记忆从计算中解耦”这一下一代模型扩展范式。

一、大模型扩展的新瓶颈：为什么必须重新审视“记忆”

过去几年，大语言模型的进步很大程度上来自一条被反复验证的技术路线：扩大参数量、扩大训练数据、扩大训练算力。从早期的 GPT 系列到 LLaMA、ChatGPT 等主流系统，规模法则长期保持着较好的预测能力。只要模型参数、训练 token 数和算力预算同步增长，模型的困惑度、下游任务能力、推理表现通常会持续改善。这条路线简洁有力，也在工程上催生了大规模分布式训练、混合精度、数据流水线等一系列基础设施的成熟。

但当模型规模从几十亿参数推进到数百亿、数千亿参数之后，单纯“堆参数”开始出现明显问题。第一是训练成本增长过快，第二是推理延迟和显存占用难以承受，第三是很多新增参数并没有被高效利用。一个重要原因是：传统 Transformer 在记忆事实性知识时，主要把知识编码在前馈网络的稠密权重矩阵中。这种方式并不是专门为“存事实”设计的，而是前馈网络在非线性变换过程中顺带形成的副产品。结果就是，为了记住更多事实，前馈网络必须变得更宽；而前馈网络变宽后，每次前向传播又必须激活全部参数，导致计算效率下降。

换句话说，模型把“记忆”和“计算”耦合在了一起。推理时，即使当前 token 只需要回忆一个很小的事实，模型仍然要完成整个前馈网络的矩阵乘法。对于需要大规模世界知识、长尾实体、代码库、多语言事实等场景，这种耦合会越来越吃力。因此，研究者开始尝试把知识存储从稠密前馈网络中剥离出来，用专门的模块承担记忆任务。Meta AI 的可扩展内存层就是其中一条代表性路线。

Meta AI 相关工作的核心动机可以概括为：如果我们能够用一个显式、可扩展、稀疏激活的键值记忆模块，替换或增强传统前馈网络，是否可以在增加模型知识容量的同时，不让推理计算量成比例增长？如果可行，模型就不再需要把所有知识都塞进稠密矩阵，而是可以用“参数化数据库”的方式存储海量事实，并按需读取。这不仅可能改善事实记忆，也可能为大模型的规模扩展提供一条不同于堆叠前馈网络的第二增长曲线。

本文的目标是把这一思路讲透。文章不会默认读者已经读过原始论文，而是从 Transformer 的基础结构讲起，逐步引入产品键查找、键值聚合、稀疏梯度训练和分布式实现。为了让内容尽量接近工程实操，本文还会给出简化版的 PyTorch 实现、伪代码、参数选择建议和常见问题排查思路。

二、从稠密前馈到显式记忆：背景与直觉

2.1 前馈网络为什么能“记住”事实

标准 Transformer 层通常包含多头自注意力和前馈网络两个核心子层。自注意力负责建模 token 之间的依赖关系，前馈网络则对每个位置独立施加非线性变换。近年来大量分析和编辑实验表明，语言模型中的事实性知识有相当一部分被编码在前馈网络的参数中。例如，当模型被问到“法国的首都是哪里”时，能够输出“巴黎”，并不是模型在回答时临时连接了外部知识库，而是相关事实已经被隐式地存储在多层前馈网络的权重里。

从形式上看，前馈网络可以写成两层线性变换加一次非线性激活。第一层把输入向量映射到更高的隐藏维度，激活函数引入非线性，第二层再映射回原始维度。这个结构可以看作一种高维查找表：输入经过线性变换后进入一个很大的特征空间，不同的输入模式会激活不同的特征组合，从而实现不同信息的读取。问题在于，这个“查找表”是稠密的，输入通过第一层矩阵乘法时，所有隐藏单元都会参与计算，即使当前任务只需要其中极少数单元。

2.2 稠密记忆的三个代价

稠密前馈网络负责记忆时，主要存在三个代价。第一是计算代价：每次前向都要对完整矩阵做乘法，参数越多，计算越大。第二是干扰代价：许多事实被压缩到同一组参数中，不同事实之间容易相互干扰，尤其是低频事实可能被高频模式覆盖。第三是扩展代价：为了显著提高记忆容量，通常需要成倍增加前馈网络宽度，而宽度增加后的边际收益可能并不理想。

这些问题并不只是在理论层面存在。实际训练中，超大模型虽然能记住更多知识，但其中很多知识分布在极高维空间中，难以被精准定位和更新。对于需要持续更新知识、记住大量低频实体的应用，稠密记忆的成本与灵活性都显得不够友好。这也解释了为什么检索增强、外部记忆、混合专家等方案会不断出现。

2.3 混合专家模型：一条有益但仍不够彻底的路线

混合专家模型把前馈网络拆分为多个并行的专家，每个 token 根据路由选择其中的一小部分专家进行计算。这样做可以在参数量很大的情况下，保持每次前向的激活参数量相对较小，因此被称为条件计算。Mixtral、Gemma 等模型已经证明，混合专家可以在相同训练成本下取得更好的语言建模能力，并为大规模模型部署提供了更高效的选择。

不过，混合专家并不是真正意义上的显式记忆。其每个专家仍然是一个完整的前馈网络，内部仍然是稠密矩阵乘法。专家记住事实的方式仍然是把知识编码在权重中，激活某个专家时，该专家内部的全部参数都会参与计算。专家数量通常也只有几个到几十个，很难扩展到百万级离散知识单元。因此，混合专家虽然在计算上做了一定程度的稀疏化，但没有从存储结构上解决“事实应该被单独存放、按需读取”的问题。

2.4 键值记忆的长期脉络

键值记忆的历史可以追溯到早期的记忆网络、神经图灵机等模型。这些工作的共同思想是：维护一组可读写的记忆槽位，每个槽位包含一个键和一个值；查询向量通过相似度与键匹配，再读取对应值并加权汇总。由于键值对相对独立，这种结构天然适合表示离散知识。

早期记忆网络的主要困难是规模与训练效率。如果键的数量很大，每次查询都需要在全量键上计算相似度并排序，计算开销很高；同时，记忆寻址过程往往不稳定，训练起来比较困难。Meta AI 可扩展内存层的关键创新，是使用“产品键”这种结构化查找方式，把一个超大的相似度搜索问题拆解为多个小规模查找，配合稀疏激活训练，从而把键值数量扩展到数百万甚至数十亿级别，并保持 GPU 上的高效并行。

三、可扩展内存层的整体设计

3.1 内存层不是独立模型，而是 Transformer 中的一层

Meta AI 的可扩展内存层并不是一个单独运行的记忆库，而是嵌入 Transformer 内部的一种参数化子层。它可以替换部分前馈网络，也可以与前馈网络并存。一般来说，可以在标准 Transformer 的若干层之后插入内存层，内存层接收前一层的隐藏状态作为输入，对每个 token 独立执行键查找和值聚合，输出一个与输入维度相同的向量。该向量通过残差连接加入主线，再经过层归一化后继续向下传递。

这种设计的一个重要好处是兼容性强。因为输入输出维度相同，内存层可以出现在任意位置，并且不需要改变 Transformer 的其余部分。研究人员可以逐步把标准前馈层替换为内存层，也可以在低层保留前馈网络，在中高层插入内存层，从而在表示能力和记忆容量之间取得平衡。

3.2 内存层的组成要素

一个典型的内存层包含以下要素：查询投影、键集合、值集合和聚合逻辑。设模型隐藏维度为 d，输入向量为 x，查询投影通常是一个可训练的线性变换，把 x 映射为查询向量 q。键集合由大量键向量组成，值集合与键集合一一对应。对于某个 token，内存层通过 q 在键集合中找到最相关的一批键，取出这些键对应的值向量，再根据相似度得分对值向量做加权求和，得到最终输出。

用伪代码可以表示如下：

q = W_q * x top_indices = top_k_similarity(q, keys) selected_values = values[top_indices] scores = dot(q, keys[top_indices]) weights = softmax(scores) output = sum(weights * selected_values) return x + output

其中top_k_similarity并非在全量键上直接扫描，而是通过产品键机制高效实现。这是可扩展内存层与传统键值记忆网络最核心的区别之一。

3.3 稀疏激活：参数规模与计算规模解耦

内存层最吸引人的特性是稀疏激活。假设内存层总共有上千万甚至上亿个键值对，但对于每个 token 的查询，模型只激活几十个候选键值对。参数总量尽管巨大，单次前向的实际计算量却只与候选数量、隐含维度以及子键空间大小有关，与总键值对数量没有直接线性关系。

这意味着，我们可以通过增加键值对总数来扩大模型“背下来多少东西”，而不必同步增加每一步推理的计算量。这与稠密前馈网络形成鲜明对比：稠密模型增加事实记忆通常意味着增加前馈宽度，而宽度增加会直接增加矩阵乘法的计算量。内存层则把参数规模和推理算力解耦，为超大规模模型提供了一条新的扩展方式。

3.4 与注意力机制的关系

注意力机制和内存层都涉及“查询—键—值”结构，但二者作用范围不同。自注意力是在当前序列内部查找 token 之间的依赖，键和值来自输入序列本身；内存层则是在可训练的参数化记忆库中查找知识，键和值是可以跨样本共享的长期参数。自注意力的计算复杂度与序列长度相关，而内存层的计算复杂度与序列中的 token 数量呈线性关系，与序列长度平方无关。

正是因为二者不冲突，Meta AI 的方案可以把内存层直接叠加在注意力层和前馈层构成的骨干网络上。注意力继续负责上下文的局部与全局依赖，前馈网络继续负责非线性变换，内存层专门负责大规模事实记忆。三者各司其职，共同构成更完整的模型表达。

四、产品键：如何让百万级查找变得可训练

4.1 朴素键值查找的问题

如果直接维护 K 个高维键向量，并对每个查询执行全量相似度计算，时间复杂度至少与 K 成正比。当 K 达到百万、千万甚至更大时，这一步就会成为瓶颈。虽然可以使用近似最近邻索引，但近似最近邻通常依赖聚类、树结构或哈希，其训练过程与深度学习优化器的兼容性并不理想，也难以在高吞吐 GPU 集群上实现稳定的端到端反向传播。

产品键的出发点，是避免显式地对全部 K 个键做最近邻搜索，而是把键集合构造成多个子键空间的笛卡尔积。通过分别进行小规模查找，再组合候选，就可以覆盖巨大的产品键空间。

4.2 产品键的定义

产品键的想法可以用一个简单例子说明。假设我们希望拥有 1,000,000 个键。直接存储 1,000,000 个键向量成本很高。若把键拆成两组子键，第一组包含 1,000 个子键，第二组包含 1,000 个子键，那么二者的笛卡尔积正好可以表示 1,000 × 1,000 = 1,000,000 个产品键。也就是说，我们只需要维护 2,000 个基础子键向量，就可以表达 1,000,000 个组合键。

从高维空间的角度看，产品键相当于把一个完整键向量拆分为多个低维片段，不同片段分别来自不同的码本。查询向量也按相同方式拆分。每个片段在其对应码本中独立查找最相似的若干子键，再把选中的子键索引组合起来，得到候选产品键。

4.3 两段产品键的查找过程

以最常见的两段产品键为例。设查询向量 q 被拆分为 q1 和 q2，键码本 A 包含 K1 个子键，键码本 B 包含 K2 个子键。完整产品键空间共有 K1 × K2 个键。前向查找过程如下：

用 q1 与码本 A 中的所有子键计算相似度，选择 top-k1 个子键；
用 q2 与码本 B 中的所有子键计算相似度，选择 top-k2 个子键；
将 A 中选中的 k1 个子键索引与 B 中选中的 k2 个子键索引做笛卡尔积，得到 k1 × k2 个产品键候选；
对于这些候选产品键，构造完整键向量，并计算查询与它们的得分；
根据得分取 softmax，对候选所对应的值向量进行加权求和。

假设 k1 和 k2 都是 4，那么每个 token 最终只需要处理 4 × 4 = 16 个候选键值对。即使完整键空间高达百万级，前向计算的瓶颈也主要由码本 A、码本 B 的内积搜索和少量候选聚合构成，而不是由百万级键空间直接决定。

4.4 为什么产品键高效

产品键高效的关键在于“加法级查找覆盖乘法级空间”。以两段结构为例，完整键空间是 K1 × K2，但基础码本的查找成本只与 K1 + K2 相关。若 K1 和 K2 都取 1000，那么完整空间规模是百万级，基础查找规模只有两千级。即使在每一段取前若干相似子键，计算成本也远远低于直接搜索百万键。

更重要的是，每个码本内的相似度计算就是矩阵乘法，可以在 GPU 上充分发挥并行能力。q1 与码本 A 的内积、q2 与码本 B 的内积都可以作为批量矩阵乘法完成。top-k 选择虽然不像矩阵乘法那样规则，但码本规模通常控制在几十万以内，因此取 top-k 的开销仍然可控。

4.5 多段扩展与索引表示

两段产品键是最常见的配置，但产品键并不限于两段。理论上可以把键拆分为更多段，进一步降低基础码本规模，或覆盖更大的总键空间。不过，段数越多，候选笛卡尔积也越复杂，索引管理、值矩阵映射和通信成本都会上升。实践中，两段结构通常已经能在容量、精度和工程复杂度之间取得良好平衡。

每个产品键可以用一个组合索引表示，例如(i, j)，其中 i 来自码本 A，j 来自码本 B。值矩阵则按照产品键索引排列。只要知道组合索引，就可以通过简单偏移计算获得对应值向量的行号。这种规则性使得值矩阵可以按行分片存储，也方便分布式场景中的通信和 gather 操作。

五、键值存储、得分与聚合细节

5.1 键向量的组织

内存层中的键并非完整维度为 d 的向量，而是根据产品键拆分为多段。每一段是维度为 d/g 的向量，g 为分段数。两段产品键中，键向量的前一半来自码本 A，后一半来自码本 B。查询向量 q 也按同样方式拆分成 q1 和 q2。这种拆分不仅带来了查找效率，也降低了每个基础码本的内存占用。

码本 A 和码本 B 的规模可以不同。例如，可以让第一段有 32768 个子键，第二段有 32768 个子键，则总产品键空间约为 10 亿级别。也可以让第一段较小、第二段较大，以调整候选分布和计算负担。这种灵活性为不同规模和不同任务提供了可调空间。

5.2 相似度得分与温度控制

在码本内查找时，通常使用查询子向量与所有子键的点积作为相似度分数。为了避免点积值过大导致 softmax 过度集中，通常需要引入温度缩放。常见做法是将点积结果除以根号维度或一个可学习、可调节的温度系数。温度较高时，候选分布更平滑；温度较低时，模型更倾向于选择得分最高的少量键。

在训练初期，查询投影尚未学好，点积分布可能不稳定。此时适当增大温度或缩放下标，可以防止 top-k 候选过于集中、导致部分键永远得不到梯度。随着训练推进，可以逐步恢复标准缩放，让模型学会更尖锐地寻址。

5.3 值的读取与加权输出

得到候选产品键后，需要从值矩阵中读取对应值向量。值矩阵通常形状为[K1 * K2, d]，因此是一个非常大的可训练参数。前向过程中不会读取整个值矩阵，而是只读取候选索引对应的若干行。这个操作可以使用连续索引的 gather，也可以通过偏移计算实现。

读取完候选值后，对候选得分套用 softmax，得到每个候选值的权重，再把候选值按权重求和。由于候选数量很小，softmax 和加权求和的计算开销几乎可以忽略。最终输出向量维度与输入一致，可以进入残差连接。

5.4 数值稳定性

在大规模键值查找中，数值稳定性非常重要。如果长期不约束键向量的模长，点积可能随训练逐渐变大，导致 softmax 饱和。常见做法包括对键和查询进行归一化，或在点积前对键向量做 rescale。也可以限制键向量的初始化方差，并在训练中监控点积得分的均值与方差。

此外，如果需要使用嵌入归一化，应确保归一化操作在分布式 shard 内和跨设备之间保持一致。数值问题一旦出现，往往会先表现为 loss 尖峰、梯度爆炸或候选分布过于集中，因此在训练内存层时应重点监控这些信号。

六、内存层的稀疏训练与优化

6.1 稀疏激活动态下的梯度流动

内存层的一个训练特点是：每个训练步骤只有少量键值对参与了前向和反向传播，因此只有这些被激活的参数会累积梯度。这为内存层带来了稀疏更新特性，但也带来新的挑战。对于极少数被频繁命中的热点键值对，它们的优化器状态会更新得非常频繁；而对于大量几乎不被命中的键值对，梯度可能长期为零，参数也不再变化。

为了让所有键值对都有机会被训练，需要在训练早期保证查询投影能够相对均匀地覆盖码本空间，并在损失函数中加入负载均衡约束。理想情况下，每个训练批次的 token 所命中的产品键应该尽量分散，这样值矩阵的各行才能保持有意义的更新频率。

6.2 初始化与缩放

查询投影、键码本和值矩阵都需要仔细初始化。查询投影通常使用小方差随机初始化，使得初始查询子向量与码本子键的点积不会过大。键码本也可以使用小方差高斯分布初始化。值矩阵的初始化通常更敏感，因为值向量的输出会直接进入残差流，如果初始值太大，可能破坏前几层已经学到的表示。

一种稳妥的做法是：将查询投影初始化为接近正交的小矩阵；键码本和值矩阵使用低标准差初始化；在点积计算中加入固定的缩放因子。还可以在内存层输出上增加一个可学习或固定的缩放，使初始内存层几乎相当于恒等映射，从而降低训练初期的优化难度。

6.3 独立优化器状态与学习率

由于键值参数数量庞大，为所有参数维护完整的 Adam 一阶矩和二阶矩成本很高。工程上通常对稀疏参数使用稀疏优化器状态，只为实际获得梯度的参数分配状态；或者使用 Adafactor 等低内存优化器。也有团队对查询投影、键码本和值矩阵使用不同学习率，因为它们在训练动态中的角色不同。

通常，键码本和值矩阵的学习率可以略低于主干网络，以避免记忆层对知识表示造成过大扰动。查询投影则需要更高的学习率和更长的 warmup，因为它决定了查找质量。具体配置需要根据模型规模和任务调整，但“分离学习率 + 稀疏优化器状态”几乎是内存层训练的工程共识。

6.4 负载均衡策略

如果没有任何约束，某些产品键可能被大量 token 同时选中，形成热点；另一些键可能从未被使用，成为“死参数”。这不仅降低参数利用率，还会造成分布式训练中的设备负载不均。为解决这一问题，可以在损失中加入辅助项，鼓励每个码本的各子键命中次数趋于均匀。

例如，可以统计当前训练批次内每个码本子键被选中的频率，并计算该频率与均匀分布的偏差。将其作为附加损失，与语言建模损失加权相加。也可以使用简单有效的机制：若某个子键被选次数超过容量上限，则将其从候选中排除或降低权重。实践表明，负载均衡损失对最终效果和训练稳定性都有明显帮助。

七、与主流扩展范式的横向对比

7.1 与稠密前馈网络对比

稠密前馈网络训练稳定、实现简单，但每次前向都会激活全部参数。对于纯粹的语言建模任务，稠密模型在中小规模下仍然表现优秀。但随着对事实记忆容量要求提高，稠密模型需要不断加宽前馈层，计算成本快速上升。内存层则允许总参数量远大于稠密基线，同时保持较低的激活参数量。

不过，内存层并不会在所有场景都优于稠密前馈。对于短上下文、非事实密集型、或模型规模本来就不大的任务，稠密前馈仍然可能是更简单、更稳定的选择。内存层的价值主要在大规模、强记忆需求、对推理成本敏感的设定中体现。

7.2 与混合专家模型对比

混合专家和内存层都有稀疏激活、大规模参数、按需计算等特点，但机制差异很大。混合专家的专家是完整前馈网络，路由作用在选择专家这一层级；内存层的存储单元是键值对，寻址作用在更细粒度的知识单元层级。混合专家的每个专家内部仍然是稠密矩阵乘法，而内存层通过对键值对的相似度查找完成读取。

从容量角度看，内存层更容易扩展到百万级或更高级别的独立知识单元；混合专家扩大容量通常需要增加专家数量，而专家数量过大时会带来路由不稳定和负载均衡困难。从适用角度看，两者甚至可以结合使用：前馈层采用混合专家提供稀疏推理，中间插入内存层提供显式知识存储，从而同时获得条件计算和可扩展记忆的好处。

7.3 与检索增强生成对比

检索增强生成通过外部文档库检索相关内容，再把检索结果拼接到上下文窗口中。它的优势是知识可更新、证据可解释，并且不需要重新训练模型。但每次推理需要维护外部索引并执行检索，检索到的文本还会占用上下文窗口，进一步增加注意力计算和 KV Cache 成本。

可扩展内存层将知识直接编码在模型参数中，不需要外部数据库，不增加序列长度，也不依赖独立的检索服务。它的知识是端到端训练得到的，更适合存储需要频繁访问、难以用文本形式表达、或要求低延迟访问的内部知识。二者并不是竞争关系，很多场景下可以互补：用外部检索处理高频变化的文档知识，用内存层承载模型内部的长期背景知识和结构化事实。

7.4 与早期可微记忆网络对比

神经图灵机、内存网络等早期模型提出了软寻址和外部记忆的构想，但由于训练困难、容量有限、缺乏高效硬件实现，未能在产业级模型中广泛落地。可扩展内存层继承了键值寻址的核心思想，并通过产品键、稀疏激活、专用通信和负载均衡，解决了规模化训练与推理问题。

可以说，早期的可微记忆网络证明了“显式记忆”概念可行，而 Meta AI 等工作把它提升到了可以在现代大规模语言模型中实际使用的工程级别。这种从概念到系统的跨越，是内存层真正引人关注的原因。

八、实验评估：语言建模、事实任务与扩展性

8.1 语言建模困惑度

语言建模困惑度是评价模型预测下一 token 能力的核心指标。Meta AI 的研究将包含内存层的模型与参数量相近或计算量相近的稠密 Transformer 进行对比。结果显示，在相同训练 token 数下，内存层模型通常能取得更低困惑度。这种优势在需要大量事实记忆、长尾分布明显的语料上尤其突出。

更关键的是，内存层模型的总参数量往往远大于稠密基线，但激活参数量和推理计算量并不成比例增加。如果以“相同训练预算下的困惑度”或“相同推理成本下的困惑度”来衡量，内存层模型的优势会更加明显。这也印证了产品键稀疏查找在扩展效率上的价值。

8.2 事实性问答与长尾实体

语言建模困惑度有时不能完整反映事实记忆能力，因此事实问答任务被用来进一步检验。测试中包含大量关于人物、地点、时间、组织等实体的问题。与传统稠密模型相比，内存层模型在低频实体和长尾事实上的回答准确率更高。

这是因为传统前馈网络把事实压缩进稠密参数，不同事实之间会形成竞争。高频知识更容易被保留，低频知识可能被覆盖。内存层则以键值对形式提供相对独立的存储位置，只要查询向量能够在码本中命中正确键，就能读取对应值。事实之间具有更强的隔离性，有利于长尾记忆。

8.3 推理任务的稳定表现

有人担忧大规模显式记忆会占用模型容量，从而损害推理能力。但从实验结果看，内存层模型在数学、代码、常识推理等任务上并未出现明显退化，甚至在需要结合大量背景知识的推理任务上表现更好。原因在于，记忆层解决的是“知识获取”问题，而注意力层和前馈网络仍然承担逻辑推理与上下文综合。

当模型能更可靠地访问事实时，后续推理模块就不必在模糊的稠密记忆中反复猜测，可以把更多计算用于结构化推理。因此，内存层并不是替代推理，而是为推理提供了更稳定的知识输入。

8.4 消融研究要点

消融实验通常围绕以下几个维度展开：内存层插入位置与频率、产品键码本规模、每段 top-k 数量、键值维度、负载均衡强度、学习率设置等。多项消融得到一些比较一致的结论：

内存层不宜完全取代所有前馈层，每隔若干层插入一个内存层通常效果更好；
增大产品键总空间几乎总能带来语言建模收益，且收益未出现明显饱和；
每段 top-k 数量需要与总候选数量平衡，过大增加计算量，过小则可能找不到正确键；
负载均衡损失对参数利用率和最终效果具有显著影响；
键值维度的提升有收益，但存在边际递减，需要结合计算成本选择。

这些发现为实际部署中的超参数选择提供了比较明确的方向。

九、工程实现：分片、内核与通信

9.1 大规模参数分片

内存层的键码本和值矩阵规模非常大，通常无法完整放在单张 GPU 上，需要采用模型并行分片。键码本可以按子键空间切分，值矩阵可以按产品键索引行切分。前向传播时，每个设备只负责自己分片内的码本内积、top-k 搜索和值读取。随后，各设备之间交换候选索引和部分加权结果，完成全局聚合。

这种分片方式的优点是通信量较低。由于每个 token 只产生少量候选索引，设备之间的通信集中在小的索引张量和小维度输出向量上，远小于标准稠密模型中的全量激活通信。因此，内存层可以在较大规模的 GPU 集群上训练，并保持良好的扩展性。

9.2 矩阵乘法与 top-k 算子

产品键查找包含三个关键算子：查询子向量与码本的内积、top-k 索引选择、候选值 gather。内积可以用标准矩阵乘法库实现，如 cuBLAS 或 PyTorch 底层算子。top-k 还可以通过对码本得分进行部分排序实现，但需要避免在反向传播时出现不稳定的梯度。当前常见做法是自定义 CUDA kernel，只对得分矩阵中的选中位置回传梯度。

值 gather 的难点在于不同 token 命中的索引不同，属于稀疏、不规则访问。通常会将候选索引展平为连续索引，再使用 index_select 或 scatter/gather 操作。也可以提前将值矩阵按映射表重新排列，使特定码本组合对应的值行连续存放，从而减少随机读取。

9.3 负载与热点处理

分布式训练中的负载不均主要来自某些码本子键或值分片被频繁命中。某些设备可能需要进行远多于其他设备的 top-k、gather 和梯度计算，形成木桶效应。解决方案包括：在损失中加入负载均衡惩罚、限制每个子键的容量、在路由后重新分布 token 到不同设备等。

Meta AI 的实现通常结合辅助损失与容量因子。辅助损失让命中分布尽量均匀，容量因子则提供硬性上限，避免偶发的极高峰值。监控每个设备的 token 数量、每个子键的命中次数以及梯度密度，是判断负载均衡是否有效的直接方法。定性地看，负载均衡做得好的内存层，各 GPU 之间的计算时间应当接近。

9.4 推理优化与长期存储

推理阶段，内存层的优势在于可以预先存储好键值参数，随模型一起加载到显存或高速内存中。由于查询只激活少量候选，推理延迟不会随着键值总规模的增大而线性上升。对于超大内存层，可以将不常使用的值分片放在 CPU 内存或 NVMe 上，通过小索引访问并按需加载，但需要平衡不同设备之间传输延迟。

此外，内存层可以支持后续微调和增量更新。如果希望注入新知识，可以冻结原有参数，只训练一组新增的键值单元；也可以通过蒸馏或正则化将新事实写入未被充分利用的键值区域。与外部知识库相比，参数化记忆的修改需要更谨慎，但通过合理设计仍然可以实现可控更新。

十、简化实现：用 PyTorch 理解产品键内存层

10.1 目标是理解，而不是追求完整复现

下面给出一段简化实现，用于帮助理解产品键内存层的前向逻辑。实际产品级场景还需要处理分布式分片、通信、负载均衡损失、量化、稀疏优化器状态等，不建议直接拿这段代码训练巨大模型。但它足以体现前文讨论的核心组件：查询拆分、码本检索、组合索引、值聚合。

import torch import torch.nn as nn import torch.nn.functional as F class ProductKeyMemory(nn.Module): def __init__(self, dim, num_keys_a, num_keys_b, top_k): super().__init__() assert dim % 2 == 0 self.dim = dim self.half_dim = dim // 2 self.num_keys_a = num_keys_a self.num_keys_b = num_keys_b self.top_k = top_k self.query_proj = nn.Linear(dim, dim, bias=False) self.keys_a = nn.Parameter(torch.empty(num_keys_a, self.half_dim).normal_(0, 0.02)) self.keys_b = nn.Parameter(torch.empty(num_keys_b, self.half_dim).normal_(0, 0.02)) self.values = nn.Parameter( torch.empty(num_keys_a * num_keys_b, dim).normal_(0, 0.02) ) self.temperature = self.half_dim ** 0.5 def forward(self, x): q = self.query_proj(x) q1, q2 = q[..., : self.half_dim], q[..., self.half_dim :] score_a = q1 @ self.keys_a.t() / self.temperature score_b = q2 @ self.keys_b.t() / self.temperature top_a = score_a.topk(self.top_k, dim=-1).indices top_b = score_b.topk(self.top_k, dim=-1).indices flat_indices = [] for i in range(self.top_k): for j in range(self.top_k): flat_indices.append(top_a[..., i] * self.num_keys_b + top_b[..., j]) indices = torch.stack(flat_indices, dim=-1) selected_values = self.values.index_select(0, indices.reshape(-1)) selected_values = selected_values.reshape(*indices.shape, self.dim) selected_keys_a = self.keys_a[top_a] selected_keys_b = self.keys_b[top_b] key_scores_a = q1.unsqueeze(-2) @ selected_keys_a.transpose(-1, -2) key_scores_b = q2.unsqueeze(-2) @ selected_keys_b.transpose(-1, -2) key_scores_a = key_scores_a.squeeze(-2) key_scores_b = key_scores_b.squeeze(-2) combined_scores = key_scores_a[..., :, None] + key_scores_b[..., None, :] combined_scores = combined_scores.reshape(*indices.shape) weight = torch.softmax(combined_scores, dim=-1) output = (weight.unsqueeze(-1) * selected_values).sum(dim=-2) return x + output

这段代码中，查询向量被拆成两半，分别在keys_a和keys_b两个码本中取 top-k，组合索引后读取对应值。得分使用两段子键得分的相加结果，最终用 softmax 加权聚合。该实现展示了产品键的两阶段结构，但还没有加入负载均衡损失、分布式分片和稀疏梯度优化。

10.2 实际实现还要考虑什么

上述简化实现存在一些重要省略。首先，索引组合使用 Python 嵌套循环，仅适合代码理解，不适合大规模训练。真正的系统会使用批量化索引运算或专门内核，避免在设备端串行展开。其次，index_select在候选数量较大时可能不够高效，实际系统通常会预排序值矩阵，或使用定制化的 gather kernel。

实际训练还需要加入负载均衡损失。可以统计每个码本子键在当前 batch 中被选中的次数，并让这个分布尽量接近均匀分布。也可以在topk之前对得分施加噪声，帮助未命中键获得学习机会。此外，键码本和值矩阵通常需要分片到多张 GPU，此时查询投影可以复制到所有设备，码本内积和值读取则按分片执行，最后通过 all-reduce 或 reduce-scatter 完成输出聚合。

十一、训练内存层模型的实用建议

11.1 先小规模验证，再逐步扩展

内存层的超参数较多，直接在大规模模型上尝试风险很大。建议先用小模型、小码本和较短训练步数验证前向和反向的数值稳定性。可以设置一个极小的 memory layer，替换主干网络中的一层前馈网络，确认损失曲线不出现异常尖峰。之后再逐步扩大码本规模、插入更多内存层，并观察困惑度与吞吐变化。

小规模验证还有助于发现实现错误。例如，乘积键索引计算错误、值矩阵行映射错位、top-k 梯度传播缺失等，都可能在小规模下表现为训练不收敛或结果异常。尽早定位这些问题，比在大模型中途调试要经济得多。

11.2 监控训练信号

内存层训练中应重点监控几类信号：模型 loss 与梯度范数、码本命中分布、候选 softmax 熵、每个码本子键的命中频率、各 GPU 的计算时间。命中分布应该保持相对均匀但允许一定偏斜；候选 softmax 熵通常不应长期接近零；各类稀疏参数的更新频率也不宜过分集中于极少部分。

如果发现某几个子键每天都占据绝大多数命中，说明负载均衡不足，模型可能在走“退化为小码本”的捷径。通过降低温度、添加负载均衡损失或扩大码本，可以缓解这种现象。

11.3 数据与任务的匹配

内存层最适合需要大量事实记忆、长尾实体或结构化知识的数据。若训练数据主要是短文本、日常对话或低知识密度内容，内存层带来的收益可能有限。对预训练数据中的领域分布做统计，评估事实密度和长尾程度，有助于决定是否值得使用内存层模型。

对于代码数据，产品键同样可以记住函数签名、库 API、常见模式等知识。对于多语言语料，内存层可以承载跨语言事实并降低高频语言对低频语言的干扰。总之，数据中可被“长期参数化记忆”的结构化信息越多，内存层的价值越大。

11.4 调参顺序与建议范围

在调参时，建议按以下顺序进行：先固定键值维度、温度和初始化，调整码本规模；再确定每段 top-k；随后调节负载均衡损失系数；最后分别调整查询投影和稀疏参数的学习率。不要同时改变多个关键超参，否则难以判断效果来自哪个变量。

作为经验参考，很多实现会选择两段产品键，每段 top-k 取 4 到 16，候选总数控制在几十到几百之间。查询和键点积温度通常使用根号维度缩放。码本规模可以根据目标参数总量反推。例如，若希望产品键空间达到千万级，可以考虑两个约 3 千到 4 千规模的码本；若希望达到亿级，则可使用两个约 1 万规模的码本。具体数值必须根据硬件与任务调整。

十二、内存层的可解释性与安全挑战

12.1 可解释性的机会

与大模型的稠密参数相比，键值记忆层在可解释性上有天然优势。研究者在分析模型行为时，可以记录某个查询命中了哪些产品键、读取了哪些值、候选得分如何分布。通过观察高命中键的 k 近邻样本，可以大致判断该键编码了哪一类模式。这种可追踪性为模型审计提供了一种稠密网络难以提供的入口。

不过，产品键中的基础码本并不直接与自然语言概念对应。一个子键是组合空间中的一个坐标，只有通过后续分析才能解释它代表的隐含模式。把内存层与字典学习、探针分类器、对比样本等方法结合，是未来提升可解释性的可行方向。

12.2 隐私与参数化记忆的删除难题

如果模型在预训练或微调中记住了个人信息、敏感数据或有害内容，这些信息可能以键值形式长期保存。与外部知识库可以删除文档不同，参数化记忆的删除并不直观。我们很难精确知道某条事实存储在哪些键值对上，直接删除整个值矩阵又会影响模型能力。

为了解决这一问题，可以研究记忆定位与机器遗忘技术。在内存层模型中，键值结构为定位和编辑提供了更清晰的假设：可以先找到与目标事实相关的查询模式，再定位其命中键和读写路径，进而对特定值向量进行修正或掩蔽。相比稠密模型，这种思路在原理上更有希望实现精确遗忘，但目前仍处于探索阶段。

12.3 偏见与有害知识的扩散

内存层的高容量意味着它可能比小型稠密模型记住更多偏见和有害内容。如果训练数据中存在偏见，内存层会更完整地保留这些信息，并在合适的查询下被激活。因此，在预训练数据过滤、安全对齐和推理侧防护方面，使用内存层时不能因为架构变化而放松标准。

也可以利用稀疏寻址特点设计防线。例如，在推理时对敏感查询的命中键进行检测，若进入高风险区域则拒绝响应或改写生成。由于内存层候选数量少，这种检查比扫描稠密模型要便宜得多。这为未来内容安全提供了新的实现可能。

12.4 稳定性与不可控寻址

虽然产品键查找比早期的软寻址记忆网络稳定，但当码本规模非常大、负载不均衡或训练超参不适当时，模型仍可能进入不可控寻址状态。若通过代码审查或监控无法快速恢复，应及时降低训练规模、增大温度、加强负载均衡损失或暂停更新键码本。

参数化记忆层的训练稳定性与模型安全是同一问题的两面：一个不稳定、不可预测的寻址机制也会带来难以预测的内容风险。工程化落地时，应把训练监控、候选正则、安全过滤和模型审计统筹考虑。

十三、不同应用场景中的可扩展内存层

13.1 大规模多语言模型

多语言模型需要同时记忆多种语言中的词汇、语法、实体和事实。传统稠密模型容易在不同语言之间形成知识干扰，低资源语言的事实更容易被高资源语言覆盖。内存层可以提供更大的显式存储空间，让不同语言的事实以更独立的方式保存。

在设计多语言内存层时，可以探讨共享码本与语言专属码本。共享码本有利于跨语言迁移，语言专属码本则有助于保护低资源语言的特殊知识。通过扩大总键值规模，模型可以覆盖更多语言的长尾实体，并在翻译、跨语言问答等任务中受益。

13.2 垂直领域大模型

医疗、法律、金融、工业等领域对知识准确性要求极高，且知识更新频繁。内存层可以用键值结构存储领域实体、关系、规范与案例。相比每次推理都检索外部文档，参数化内存层可以提供更低延迟、更稳定的知识访问。

垂直领域部署中，可以结合“基础内存 + 领域内存”的方式。基础模型已经掌握了通用知识，新增领域知识以独立值分片或增量训练注入，从而降低全量重训成本。对知识版本的管理也可以借助键值分片与元数据实现，逐步形成可治理的模型知识库。

13.3 个性化助手与长期用户记忆

个性化助手需要记住用户偏好、长期目标、历史交互等大量信息。把全部上下文塞进提示词会快速耗尽上下文窗口，也增加注意力计算量。内存层可以把用户相关信息编码到键值对中，在需要时按相关性读取，在不需要时保持静默。

这种方式还能增强隐私隔离：不同用户的信息可以映射到不同键区域，并在访问时施加权限控制。与完全依赖外部数据库的方案相比，内存层能更紧密地融入模型生成过程；与完全依赖上下文窗口相比，它又能支持更长的记忆周期并节省推理成本。

13.4 知识图谱与科学计算

知识图谱中的实体、关系和属性天然适合用键值形式表达。内存层可以被看作一个软性、端到端训练的知识图谱。模型可以将实体作为键、属性和关系向量作为值，在回答复杂科学问题时快速调用背景知识。与传统符号知识图谱相比，这种表示可以学习连续语义，并直接在生成过程中使用。

在科学发现、代码库理解、数学定理检索等任务中，内存层有望作为模型内部的知识基础设施。外部知识库仍然适合用于可追溯证据和高频更新，而内存层适合承载模型已经掌握、需要快速调用的结构化背景知识。

十四、从论文到产品：落地路线图

14.1 明确目标：是记忆密集型任务还是普通任务

在决定采用内存层前，先要明确业务任务是否属于记忆密集型。事实问答、知识型助手、长尾实体识别、代码补全、多语言客服等场景更可能受益；情感分析、短文本分类、对话意图判断等任务则未必需要超大规模参数化记忆。可以先评估现有模型在事实覆盖率、长尾准确率和上下文窗口上的瓶颈，再决定是否引入内存层。

14.2 渐进式引入，降低风险

生产中不建议一次性把整个稠密模型替换为内存层模型。更稳妥的路径是保留基础模型权重，只在少数层加入内存层或替换少量前馈层，先在离线评测中验证收益。确认困惑度、事实 QA 和推理能力没有退化后，再逐步扩大内存层的插入频率和码本规模。

渐进式引入还可以帮助评估资源和成本。通过对比每个阶段的实际吞吐、显存占用和训练时间，可以判断内存层方案在自身硬件条件下的真实性价比。理论上的稀疏优势并不总是能直接转化为端到端加速，实际硬件对稀疏算子的支持程度非常关键。

14.3 建立评测与回滚机制

内存层模型上线后，需要一套覆盖事实记忆、推理能力和安全性的评测集。重点监控长尾实体准确率、已知安全问题的回答、候选分布和命中热点。任何训练策略调整都应通过离线评测与 A/B 实验，避免模型能力发生不可预见的退化。

同时应建立模型版本管理和回滚机制。因为参数化记忆的更新和遗忘尚未完全成熟，新版本模型可能出现旧知识遗忘、新知识抢占等问题。保留每轮训练前的检查点，并记录键值存储和索引映射的变化，是后续排查问题的重要前提。

14.4 构建从数据到知识注入的闭环

内存层模型要长期保持价值，需要建立知识更新闭环。可以从用户反馈、专业文档和结构化数据库中持续提取高频且稳定的知识，再通过微调或增量训练写入内存层。对于时效性强、变化快的内容，仍建议使用检索增强，而不是强行写进参数。只有把参数化记忆与外部知识库合理分工，才能兼顾准确性和时效性。

十五、常见问题与排查

15.1 训练 loss 突然飙升

如果训练 loss 突然飙升，应首先检查查询投影与键码本的点积是否失控。可以查看当前候选 softmax 熵，若熵很低且得分极大，说明寻址过度尖锐。常见处理是增大温度缩放、约束键向量模长、降低内存层输出缩放，或降低键码本和值矩阵的学习率。

也可以回退到上一个稳定检查点，调整初始化后重新训练。梯度范数监控同样重要：如果内存层相关参数的梯度显著高于主干网络，应使用梯度裁剪，并降低对应参数组的学习率。

15.2 部分键被频繁命中，其他键长期失效

这说明负载均衡不足。适当增大负载均衡损失系数，或对命中过多的子键施加容量限制。也可以增大码本规模，给查询提供更多可选键。训练早期可对查询投影使用较强 warmup，让寻址分布逐渐扩大。若问题持续，则应检查数据分布中是否存在极端重复的 token 模式。

15.3 困惑度下降但下游事实任务没有提升

这种不一致可能来自数据分布和评价目标的不匹配。语言建模在常见词上取得的提升不一定反映长尾事实。应使用专门的事实问答集和低频实体测试集进行评测，检查真实的知识覆盖率。同时可以在训练数据中提升领域文档和事实密度，使内存层真正学习到需要被调用的知识。

15.4 推理加速不达预期

内存层虽然在理论计算量上稀疏，但实际加速受 top-k 内核、gather 效率和硬件支持影响。可以先对内存层各算子做单独基准测试，定位瓶颈究竟在码本内积、top-k 还是值 gather。若能使用融合 kernel、减少中间显存分配、避免 Python 循环和重复 reshape，往往能显著提升实际性能。分布式部署中还应检查通信是否为瓶颈。

十六、未来方向与开放问题

16.1 动态更新与持续学习

当前的大模型记忆主要依赖离线训练，模型上线后难以可靠地新增或遗忘知识。内存层的键值结构为持续学习提供了新的可能。未来方向包括：为新知识动态分配未使用键值、对旧知识进行局部修改、在保持模型整体稳定的前提下增量训练。如何保证新增知识不破坏已有知识，是持续学习仍需攻克的问题。

16.2 硬件与稀疏计算生态

内存层的普及还依赖稀疏计算生态。GPU、TPU 等硬件需要为大规模稀疏 gather、不规则 top-k 和低比特稀疏参数提供更高效的原语。软件框架也需要完善稀疏优化器状态、混合分片和自动并行。随着硬件和框架的演进，内存层的实际性价比有望进一步提升。

16.3 压缩与量化

内存层的键值参数数量庞大，即使稀疏激活，存储成本仍然不可忽略。低比特量化、乘积量化和结构化剪枝可以降低内存占用。键码本本身具有结构化特征，非常适合向量量化；值矩阵也可以按分片做低比特压缩。压缩不应对检索精度造成过大影响，因此需要在码本规模、量化误差和命中质量之间做权衡。

16.4 统一的记忆与推理框架

更长远地看，未来的基础模型可能需要统一管理多种记忆形式：上下文记忆、工作记忆、外部检索记忆和参数化长期记忆。内存层可以作为参数化长期记忆的核心组件，与外部知识库、缓存和上下文共同构成分层记忆体系。研究的方向包括：如何决定哪些知识该写进参数，哪些该留在外部文档；如何在不同记忆系统之间路由；如何在有限成本下保证知识新鲜度与准确性。

这种统一框架一旦成熟，模型将从“只能利用上下文和静态权重”演进为“能够按需调用内部长期记忆与外部知识”的系统。Meta AI 可扩展内存层所验证的产品键和稀疏激活路径，正是这一进程中重要的一步。