NeurIPS’24｜推理计算量减小10倍!MemoryFormer：华为提出存储代替计算的Transformer新架构

论文链接：https://arxiv.org/abs/2411.12992

引言

在大模型快速发展的今天，深度学习模型的规模不断增大。然而，随着模型规模的提升，对于计算资源的消耗也随之剧增，这使得模型的部署和推理面临巨大挑战。近期，一篇来自北京大学和华为的研究人员共同合作的论文提出了一种创新性的方法——MemoryFormer，通过巧妙地利用存储空间来替代传统的密集计算的全连接层，显著降低了推理时的计算复杂度（FLOPs）。这一思路为大模型的高效推理提供了全新的解决方案。

用存储来替代全连接层

传统的Transformer模型中，计算资源的开销不仅来自于多头自注意力模块，还来自于模型内部的全连接层。针对这一特点，MemoryFormer提出了一种“记忆层”（Memory Layer）来代替这些计算量庞大的全连接层，从而极大地减少了模型的计算复杂度。具体来说，Memory Layer通过哈希算法来对输入特征进行离散化，并在内存中构建查找表，将原本需要的矩阵乘法替换为在计算上更轻量的哈希查找。这种设计使得MemoryFormer能够在推理时大幅降低所需的计算量，并且几乎不影响模型的性能。

MemoryFormer不仅简单地引入哈希查找表，还设计了一种基于局部敏感哈希（LSH）的极简索引方法，使得哈希表中的每一个存储的向量都可以在训练中根据梯度动态更新。这种设计确保了哈希表中存储的特征能够不断适应输入数据，并在推理阶段根据输入特征的相似性高效检索出近似的输出结果，实现全连接层所需的特征变换功能。此外，MemoryFormer通过多表分块和向量分段的方式来控制哈希表的存储规模，使得内存需求不会因哈希表的引入而暴增。

方法：可学习的局部敏感哈希

Memory Layer的具体实现方式如下：在一层Memory Layer内设置K个哈希表 ，每个哈希表中存储 个d维向量(τ=d/k)，对于一个输入的d维token embedding向量x，首先对其进行分段操作切分为k段τ维向量：

接下来对每段子向量 进行局部敏感哈希操作，分别求得k个哈希表的对应索引，使用的极简哈希函数为：

其中sign()函数对每一位只取符号位，integer()函数首先将输入中的-1变为0再将二值化向量转换为对应的十进制数。然后将K个哈希表的检索结果 进行加权聚合得到该Memory Layer层的输出：

其中 为各个哈希表的归一化权重，表征了embedding切分后的第k段子向量与第k个哈希表索引结果的契合程度，其计算方式为

其中

sim函数表示 向量与第 个hash bucket的相关程度，如下图中的二维坐标系示意图所示，sim函数既考虑了 与sign 之间的夹角、又考虑了 本身的幅值，在对第k个哈希表内所有 个位置进行类似softmax形式的归一化后，我们即可得到每个哈希表 的检索结果在聚合时的加权系数：

函数将整数i（0≤i＜ ）转换为对应的符号向量{-1, +1}∈ 。

考虑数学形式上的2个等价变化：

可进一步简化为

Memory Layer最终的输出为

Memory Layer对于相似的输入 和 具有相似的输出，且通过设置多个哈希表与加权聚合的方式丰富了单层的特征表达能力。MemoryFormer将Baseline Transformer中的q_proj、k_proj、k_proj等Linear Layer全部替换为Memory Layer。对于原模型中的FFN模块，我们用2个连续的且容量scale up的Memory Layer进行替代，来模拟原FFN中多个Linear Layer“先升维后降维”的特点。因此，MemoryFormer中每个Block的计算过程为：

我们移除了FFN模块中的非线性函数，取而代之的是每个Memory Layer之前都有一个Norm Layer以保证输入的向量在分布上是zero-centered的，也就是说在sign()函数的结果中出现-1和+1的概率是均等的。

在Baseline Transform中，每个block的计算量为：

在MemoryForm中，每个block的计算量为：

实验：

论文基于Pythia训练框架，在多个公开的自然语言处理基准测试中验证了MemoryFormer的性能，包括PIQA、WinoGrande、WSC、ARC-E、ARC-C、LogiQA等任务。我们在每种setting下都保持2种模型的层数和token维度是一致的、且保持Attention部分的计算量一致。实验结果表明，在多种模型尺寸下，MemoryFormer几乎消除了注意力模块之外的计算量，大幅减少了推理时模型的总计算量，且总体性能不逊于baseline Transformer模型，例如MemoryFormer将Pythia-410M模型的Attention之外per-block计算量从25.8GFLOPs降低至1.6GFLOPs，展示出更优的计算效率。

我们还将MemoryFormer与其他基于“局部自注意力”方案的Efficient Transformer进行了性能于计算量方面的对比（保持层数于embedding维度一致）：

消融实验

不同分表设置下的存储占用和模型性能：

增加FFN模块中哈希表的容量，存储消耗与模型性能收益之间的关系：

是否移除了FFN模块中的非线性函数：

哈希过程可视化

我们还对不同Memory Layer（Q、K、V、FFN）中的多个哈希表中每个位置的命中情况进行了可视化，结果表明在推理过程中，大部分哈希表中的不同位置都会以相似的概率被hash到，少数哈希表中会出现个别位置被大量token embedding都hash到的情况。

总结：

MemoryFormer为实现高效的Transformer模型提供了全新的思路，用存储取代计算、哈希表替代全连接层的设计，减少了模型推理时的计算复杂度，为大模型的高效部署、尤其是边缘设备和CPU设备的部署、带来了新的契机，为大模型的多方面可持续发展提供了新的可能性。

公众号后台回复“数据集”获取100+深度学习各方向资源整理