Meta探索大模型记忆层，扩展至1280亿个参数，优于MoE

2025-01-04 发布浏览345次点赞0次收藏0次

预训练语言模型通常在其参数中编码大量信息，并且随着规模的增加，它们可以更准确地回忆和使用这些信息。对于主要将信息编码为线性矩阵变换权重的密集深度神经网络来说，参数大小的扩展直接与计算和能量需求的增加相关。语言模型需要学习的一个重要信息子集是简单关联。虽然前馈网络原则上（给定足够的规模）可以学习任何函数，但使用联想记忆（associative memory）会更高效。

记忆层（memory layers）使用可训练的键值查找机制向模型添加额外的参数，而不会增加 FLOP。从概念上讲，稀疏激活的记忆层补充了计算量大的密集前馈层，提供了廉价地存储和检索信息的专用容量。

最近，Meta 的一项新研究使记忆层超越了概念验证，证明了它们在大型语言模型（LLM）扩展中的实用性。

论文标题：Memory Layers at Scale
论文地址：https://arxiv.org/pdf/2412.09764
项目地址：https://github.com/facebookresearch/memory

在下游任务中，通过改进的记忆层增强的语言模型的性能优于计算预算两倍以上的密集模型，以及在计算和参数相当的专家混合（MoE）模型。

这项工作表明，当记忆层得到充分改进和扩展时，可以用于增强密集神经网络，从而带来巨大的性能提升。通过用记忆层替换一个或多个 transformer 层的前馈网络（FFN）来实现这一点（保持其他层不变）。这些优势在各种基本模型大小（从 1.34 亿到 80 亿参数）和内存容量（最多 1280 亿参数）中都是一致的。这意味着存储容量实现了两个数量级的飞跃。

记忆增强架构

可训练的记忆层类似于注意力机制。给定一个查询

，一组键

，以及值

。输出是值的软组合，根据 q 和相应键之间的相似性进行加权。

在使用时，记忆层与注意力层之间存在两个区别。

首先，记忆层中的键和值是可训练参数，而不是激活参数；
其次，记忆层在键和值的数量方面通常具有更大的规模，因此稀疏查询和更新是必需的。

该研究将键-值对的数量扩展到数百万。在这种情况下，只有 top-k 最相似的键和相应的值被输出。一个简单的记忆层可以用下面的等式来描述：

其中，I 是一组指标，

，输出

。

扩展记忆层

扩展记忆层时面临的一个瓶颈是「查询 - 键」检索机制。简单的最近邻搜索需要比较每一对查询 - 键，这对于大型记忆来说很快就变得不可行。虽然可以使用近似向量相似性技术，但当键正在不断训练并需要重新索引时，将它们整合起来是一个挑战。相反，本文采用了可训练的「product-quantized」键。

并行记忆。记忆层是记忆密集型的，主要是由于可训练参数和相关优化器状态的数量庞大导致的。该研究在多个 GPU 上并行化嵌入查找和聚合，记忆值在嵌入维度上进行分片。在每个步骤中，索引都从进程组中收集，每个 worker 进行查找，然后将嵌入的部分聚合到分片中。此后，每个 worker 收集与其自身索引部分相对应的部分嵌入。该过程如图 2 所示。