Meta AI 的“记忆层”技术详解

Meta AI 近期发表的一篇论文 "Memory Layers at Scale" 为我们提供了一个新的思路：为模型配备可学习的“记忆”模块，使其能够像人脑一样，选择性地存储和检索信息，从而在不显著增加计算成本的前提下，大幅提升模型的性能，尤其是在处理事实性知识方面。

https://ai.meta.com/research/publications/memory-layers-at-scale/

本文将解读这项有趣的技术。

一、核心思想：为 LLM 打造专属“记忆库”

想象一下，我们的大脑并不是存储所有的信息，而是更倾向于记住关键的事实和联系。当我们需要回答问题时，大脑会迅速检索相关的信息。Meta AI 的“记忆层”正是受到了这种机制的启发。

传统 LLM 的知识存储主要依赖于模型自身的参数。参数越多，理论上能记住的信息也越多。但这种方式效率低下，因为所有的信息都以分布式的方式存储在巨大的参数矩阵中，检索过程也比较“笨重”。

“记忆层”的核心思想是引入额外的、可训练的“记忆槽”，用于存储模型的“知识”。 这些“记忆槽”以 键（Keys）- 值（Values） 对的形式存在，类似于一个可学习的数据库。

• 键（Keys）： 可以理解为知识的索引或标签。

• 值（Values）： 存储着实际的信息或知识表示。

二、记忆层如何工作？

当模型需要利用“记忆”时，会经历以下步骤：

生成查询（Query）： 模型会根据当前的输入生成一个查询向量，用于在“记忆库”中查找相关信息。

查找相似的键（Key Lookup）： 查询向量会与记忆层中所有的键（Keys）进行比较，计算它们之间的相似度。这里需要强调的是，记忆层与传统注意力机制的一个关键区别在于，这里的键（Keys）和值（Values）都是模型的可训练参数，而不是像注意力机制中那样，动态地来自输入序列的隐藏层状态（activations）。 我们将在下文详细解释这个区别。

选择最相关的“记忆”（Top-k Selection）： 根据相似度得分，选择最相关的若干个键（Top-k）。

检索对应的值（Value Retrieval）： 获取与选定的键（Keys）相对应的值（Values）。

融合检索到的信息： 将检索到的值（Values）与模型的原始表示进行融合，从而增强模型的能力。

可以用以下公式来概括记忆层的基本操作：

I = SelectTopkIndices(Kq)
s = Softmax(K_I q)
y = sum(s_i * V_I_i)

• q：查询向量（Query）

• K：所有键（Keys）的集合

• V：所有值（Values）的集合

• SelectTopkIndices(Kq)：选择与查询 q 最相似的 k 个键的索引。

• K_I 和 V_I：分别是被选中的 k 个键和值。

• s：对相似度进行 Softmax 归一化后的权重。

• y：加权平均后的输出，代表从记忆中检索到的信息。

三、记忆层与注意力机制的本质区别：键和值的来源

记忆层与注意力机制最关键的区别之一在于键（Keys）和值（Values）的来源。理解这一点至关重要。

1. 注意力机制 (Attention Mechanism)：动态生成的“关注点”

在标准的注意力机制中，键（Keys）、值（Values）和查询（Query）通常都来自输入序列的隐藏层状态（activations）。

• 什么是隐藏层状态？ 在深度学习模型中，每一层都会对输入数据进行一系列的变换，并输出一个中间表示，这个中间表示就称为隐藏层状态或激活值。你可以把它理解为模型在处理输入过程中提取到的特征。

• 注意力机制如何生成 Key 和 Value？ 当一个输入序列进入注意力层时，该层会将输入序列中每个位置的隐藏层状态通过不同的线性变换（由可训练的权重矩阵定义）映射成 Query、Key 和 Value。 这意味着，注意力机制的 Key 和 Value 是根据当前的输入动态生成的，它们并不存在于模型之外。

• 注意力机制的作用： 注意力机制的作用是让模型能够关注输入序列中不同的部分，并根据不同部分的关联性赋予不同的权重。它更像是模型在运行时动态地决定要“关注”哪些输入信息。

2. 记忆层 (Memory Layers)：可学习的“知识条目”

与注意力机制不同，记忆层中的键（Keys）和值（Values）是模型独立的可训练参数。

可训练参数的意义： 这意味着在模型的训练过程中，Key 和 Value 的值会不断地被调整和优化，以存储模型需要记住的特定信息或知识。

记忆层的 Key 和 Value 是什么？ 可以将记忆层的 Key 和 Value 看作是模型内部存储的一组固定的“记忆条目”。每个 Key 代表一个可识别的“知识点”或“概念”，而 Value 则存储着与该 Key 相关的信息的表示。

记忆层的作用： 记忆层的作用是提供一个可学习的外部知识库。模型可以通过比较当前的 Query 和记忆库中的 Keys，来检索相关的 Values，从而利用预先存储的知识来增强自身的理解和生成能力。

用一个形象的比喻：

• 注意力机制： 就像你在阅读一篇文章时，你的目光（Query）会在不同的词语（Key）之间移动，并根据词语之间的关系（相似度）来理解句子的含义（Value）。你的关注点是动态的，取决于你正在阅读的内容。

• 记忆层： 就像你的大脑中存储的词汇表和知识库。 当你阅读到某个词语（Query）时，你会查找大脑中是否存在相关的知识条目（Key），并提取出相关的解释和信息（Value）。 这些知识是预先存在的，并且可以通过学习不断扩充。比如阅读到中文的成语时，自然会产生对应的联想。

四、如何应对规模挑战？记忆层的关键技术

要让记忆层真正发挥作用，就必须能够扩展到足够大的规模，存储足够多的“知识”。 这篇论文提出了几项关键技术来解决扩展性问题：

1、 乘积键查找（Product-key lookup）：化繁为简的搜索策略

当记忆层的规模非常大时，每次查询都需要将 Query 与所有的 Key 进行比较，这会带来巨大的计算开销。

乘积键查找的核心思想是利用 分而治之 的策略，将一个高维度的搜索问题分解为两个低维度的搜索问题，从而减少需要比较的次数。

• 拆分“半键”： 将每个完整的键拆分成两个更小的“半键”。例如，一个 128 维的键可以拆分成两个 64 维的半键。

• 分步查找：

  1. 将查询向量也拆分成两个对应的部分。

  2. 先将查询的第一部分与所有键的第一个“半键”进行比较，找到最相似的若干个“半键”。

  3. 再将查询的第二部分与所有键的第二个“半键”进行比较，找到最相似的若干个“半键”。

• 关键在于，我们并不需要将所有可能的“半键”组合都与查询进行比较。 通过一些高效的索引和搜索策略，我们可以利用在前面两步中获得的相似度信息，优先考虑那些更有可能组成最相似完整键的“半键”组合。 这就像是在拼图游戏时，我们先找到颜色和形状最接近的碎片，再尝试将它们拼合在一起，而不是盲目地尝试所有可能的组合。

• 减少比较次数的原因： 虽然看起来需要进行两次比较，但比较半键的计算成本远低于比较完整键。 更重要的是，通过高效的索引和剪枝策略，我们可以避免对大量的“不可能”组合进行比较，从而大幅减少实际需要计算相似度的次数。

2、并行化记忆层（Parallel memory）：多卡协作，分而治之

记忆层的一个主要瓶颈是内存访问。为了解决这个问题，论文采用了并行化的策略，将记忆层的键（Keys）和值（Values）分布在多个 GPU 上进行存储和计算。

每个 GPU 只负责存储和处理一部分的“记忆”，从而降低了单个 GPU 的内存压力。

3、共享内存（Shared memory）：知识共享，效率提升

为了更有效地利用参数，论文提出在多个记忆层之间共享同一个键值对池。这意味着不同的记忆层可以访问和利用相同的“知识”，而不需要各自存储一份，从而提高了参数的利用率。

五、实验结果：记忆层的威力初显

实验结果表明，使用记忆层增强的语言模型在各种任务上都取得了显著的性能提升，尤其是在需要事实性知识的任务上。

与相同计算预算的密集模型相比，记忆模型表现更优，甚至可以匹敌参数量更大的模型。这证明了记忆层在提升模型效率方面的潜力。

六、记忆层的局限性与未来展望：能记住整个互联网吗？

有朋友可能会问：记忆层这么厉害，那能不能把整个互联网的知识都存进去呢？

答案是：目前还不行。 虽然记忆层可以有效地利用参数来存储知识，但它仍然受限于模型的整体参数规模。将整个互联网的海量信息都存储在模型内部是不现实的。

但这并不妨碍记忆层成为一个极具潜力的研究方向。 未来，我们可以探索以下方向：

• 更高效的知识编码方式： 研究如何将更复杂的知识结构编码到键值对中。

• 动态更新的记忆层： 让记忆层能够不断学习和更新新的知识。

• 与外部知识库结合： 将记忆层与外部知识库相结合，实现更大规模的知识存储和检索。

七、总结

Meta AI 的 "Memory Layers at Scale" 论文为我们展示了一种提升 LLM 效率的新思路。通过引入可学习的“记忆层”，模型可以在不显著增加计算成本的前提下，大幅提升性能，尤其是在处理事实性知识方面。虽然目前的技术还无法将整个互联网的知识存储在模型内部，但记忆层的出现为我们打开了一扇新的大门，预示着未来 LLM 发展的更多可能性。