内存层的工作原理以及它们如何增强LLM?

内存层的工作原理以及它们如何增强LLM

大语言模型（LLMs）在其参数中存储着海量的信息知识（主要以稠密层中线性矩阵变换的权重形式存在）。然而，随着参数规模的不断扩大，计算成本和能源消耗也与日俱增。那么，这些能否用简单且低成本的键值查找机制来替代呢？此前已有许多相关研究，但至今都未能达到当前AI架构的规模。

不过，Meta的研究人员终于取得突破，开发出了能够增强现有大语言模型能力的内存层。这些内存层取代了一个或多个Transformer层中的前馈网络（FFN），效果出奇地好！

内存层能将大语言模型的事实准确性提高100%以上，在编码性能和常识方面的提升，可与经过4倍以上计算量训练的传统大语言模型相媲美。这些配备内存层的大语言模型，还击败了在相似计算资源和参数规模下训练的混合专家（Mixture-of-experts ）大语言模型架构，在实际任务中的表现尤为突出。

接下来，我们将深入探讨内存层的工作原理，以及它们是如何为大语言模型赋能的。可以说，如果下一代AI架构不采用这些技术，将会错失良机。

什么是内存层？

内存层的工作方式与Transformer中的注意力机制有些相似。给定查询（Q）、键（K）和值（V），内存层会输出值（V）的加权和，其中权重是通过softmax函数，依据查询和键之间的相似度来确定的。

不过，内存层与传统注意力机制有两个关键区别。其一，在注意力机制中，键和值是针对每个查询动态计算的；而在内存层里，键和值是可训练的参数，能够持续学习并存储下来。其二，内存层中使用的键值对数量极为庞大（可达数百万个）。计算输出时，仅会采用最相似的前k个键及其对应的值，如此一来，在大规模数据下，查找和更新操作的计算效率得以显著提高。

内存层可以用以下公式来描述：

1. 首先，根据与查询的相似度，为选定的前k个键计算索引（I）。

2. q和K分别代表查询键和可训练键。接着，计算选定键的相似度分数（K(I)q），并使用Softmax进行归一化，从而得到权重（s）。

3. q和K（I）分别表示查询和选定的前k个键。最后，利用选定的前k个值的加权和来计算输出（y）。

4. s表示经过softmax归一化的权重，V（I）代表选定的前k个值。

与传统Transformer中的前馈层类似，每个令牌嵌入都会独立通过一个内存层。

如何大规模搜索相似的键？

为查询找到最相似的键，计算成本相当高。简单的最近邻搜索步骤如下：

1. 计算查询和键之间的相似度分数（比如余弦相似度）（时间复杂度：对于N个维度为n的键，时间复杂度为O(N ⋅ n) ）。
2. 按照相似度分数对键进行排序（时间复杂度：对于N个键，时间复杂度为O(N log(N)) ）。
3. 挑选出相似度分数最高的前k个键。
4. 使用前k个键计算最终输出。

对于维度为n的N个键，上述方法的内存成本为O(N ⋅ n) 。鉴于内存层中有数百万个键，这种方法显然不可行。近似最近邻（ANN）搜索在这里也不太适用，因为ANN需要计算搜索的静态索引，而内存层中的键是可训练的，在训练过程中会不断更新，这就意味着需要持续重新索引。

难道就没有其他办法了吗？当然有！这个方法借鉴了之前一篇研究论文中描述的可训练乘积量化键。

分割键

不是使用单个大的键矩阵（K），而是先将其分成两个较小的矩阵（K(1)和K(2)）。大矩阵的维度是N × n ，而小矩阵的维度是√N × n/2，这里N是键的数量，n是每个键/查询向量的维度。大矩阵实际上是这两个较小键矩阵的笛卡尔积，即K = K(1) X K(2) ，但这个大矩阵并不会被显式创建，从而节省了内存和计算资源。

分割查询

接下来，查询向量（Q）也会被分成两个较小的向量（Q(1)和Q(2)）。原始查询向量的维度是n，分割后得到的较小向量维度是n/2。这两个小向量会分别与对应的较小键矩阵进行交互。

寻找前k个相似键并计算相似度分数

对于Q(1)，在K(1)中找到前k个相似的键，其索引记为（I(1)），然后使用Softmax计算相似度分数（s(1)）。对Q(2)也执行相同的步骤。

寻找整体的前K个索引和分数

通过对索引和分数计算Argmax函数，来确定整体的前k个索引和分数。

为什么这种方法效果这么好呢？这是因为该方法不是将查询与全部N个键进行比较，而是与两个小得多的集合进行比较，对于维度为n的N个键，时间和空间复杂度从O(N ⋅ n) 降低到了O(√N ⋅ n) 。

这些操作在GPU上是如何执行的？

内存层包含数百万个可训练参数（键和值）。为了在这些参数嵌入上进行大规模运算，首先会沿着嵌入维度将它们分割成多个分片，并分布到多个GPU上。每个GPU负责管理和处理分配给自己的那部分内存参数分片。通过一个进程组来协调多个GPU之间的操作。

在查询操作时，会先识别出相关索引，并将其分发到各个GPU上。然后，每个GPU在自己的分片内查找与索引对应的相关嵌入。最后，各个GPU将部分结果进行共享和汇总，从而计算出最终输出。

加快GPU操作

可以使用PyTorch的EmbeddingBag函数来计算内存层中前k个嵌入的加权和。然而，其默认实现受到GPU内存带宽的限制，只能达到不到400GB/s的内存带宽，这与现代GPU的潜力相差甚远。

因此，研究人员为该操作的前向和反向传递实现了更高效的自定义CUDA内核。这些内核能够实现3TB/s的内存带宽，接近NVIDIA H100 GPU 3.35TB/s的理论最大值，使得嵌入袋操作的端到端速度比PyTorch的默认EmbeddingBag函数快6倍。

通过引入基于SiLU非线性的输入相关门控机制，内存层的训练性能得到了进一步提升。输出公式变为：

silu(x) = x ∗ σ(x)，其中σ(x)是sigmoid函数；⊙表示逐元素乘法操作；x是内存层的输入；y是经过基于x的门控机制缩放后的内存层输出；W(1)和W(2)是可训练的权重矩阵。

训练小型基础模型和大型内存层时，有时会出现不稳定的情况。为确保稳定性，研究人员采用了QK归一化方法。该方法在计算查询（Q）和键（K）向量的点积之前，先对它们进行归一化处理。

内存层会取代哪些前馈层？

在深度神经网络中，较低层学习基本特征，较高层则学习复杂模式。因此，在多个层中添加内存层效果最佳。

为避免增加大语言模型的总参数数量，所有层会共享一个内存池。这样一来，多个层都可以访问相同的内存，从而提高了架构的效率。研究还发现，在多层（最多3层）中使用内存层能够提升模型性能，但如果用内存层替换过多的密集前馈网络（FFN）层，反而会导致性能下降。这表明，稀疏内存层和密集前馈层都有着重要作用，二者结合使用时效果最佳。

内存层增强的大语言模型表现如何？

研究人员使用Llama系列模型（Llama2和Llama3）进行实验，将其中一个或多个前馈层（FFN）替换为共享内存层。普通内存模型只有一个内存层，而“Memory +”模型则有三个内存层，并采用了SwiLU（SiLU的缩放加权版本）非线性激活函数。

![](https://files.mdnice.com/user/13819/64c4c395-08bd-4f6d-9923-8...

.png)

其中β是可学习参数，σ(x)是sigmoid函数

除了默认的Llama模型，研究人员还选用了一个混合专家（MoE）模型（使用专家选择路由进行训练）和一个配置相当的PEER模型进行对比。

下面来看看它们的性能表现：

1. 在问答（QA）任务上，内存模型的表现优于同规模的密集模型，达到了参数数量翻倍的密集模型的性能水平。
2. “Memory +”模型的表现更出色，其性能与经过两到四倍计算量训练的密集模型相当。值得注意的是，在参数数量相同的情况下，PEER模型的表现与内存模型相近，但落后于“Memory +”模型。此外，MoE模型的表现则大幅落后于内存增强模型。

3. 在固定基础模型的情况下，进一步扩大内存参数的规模，内存模型在事实性问答任务上的性能有显著提升。拥有6400万个键的13亿参数内存模型，尽管训练时使用的令牌数量只有Llama2 70亿参数模型的一半，计算量仅为其1/10，但性能却与之相当。
   
4. 在80亿参数规模下，内存模型在一般科学、世界知识和编码基准测试中，显著优于密集基线模型。经过1万亿个令牌训练后，“Memory +”模型的性能接近经过15万亿个令牌（数据量多15倍）训练的Llama3.1 80亿参数模型。

本文由mdnice多平台发布