摘要本文围绕 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 来自码本 Aj 来自码本 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, biasFalse) 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 可扩展内存层所验证的产品键和稀疏激活路径正是这一进程中重要的一步。十七、总结Meta AI 可扩展内存层以键值记忆为核心通过产品键把大规模相似度检索转化为多个小规模码本查找再结合稀疏激活让模型可以拥有远超稠密前馈网络的知识容量同时避免推理算力随参数规模线性膨胀。它继承并升级了早期键值记忆网络的思路又克服了其在规模、训练效率与硬件可用性上的短板。与传统稠密前馈、混合专家、检索增强生成等方法相比内存层在记忆容量、长尾事实、推理成本和知识隔离性之间提供了新的折中。它并不是对所有模型都必要的万能组件但在大模型需要承载海量事实知识、追求更低推理成本的背景下具有明显的战略价值。当前内存层仍然面临训练稳定性、负载均衡、硬件支持、解释性和可删除性等挑战。但随着产品键查找、通信优化、量化压缩和持续学习技术的进步可扩展内存层有潜力成为下一代基础模型中负责长期知识存储的标准组件。对于希望深入理解大模型扩展前沿、或尝试构建知识密集型应用的读者这是一个值得长期跟踪和实践的方向。