点击下方
卡片
，关注
「3DCV」
公众号
选择
星标
，干货第一时间送达
添加小助理：cv3d001，备注：方向+学校/公司+昵称，拉你入群。文末附3D视觉行业细分群。
扫描下方二维码，加入
「3D视觉从入门到精通」知识星球
(
点开有惊喜
)
，星球内凝聚了众多3D视觉实战问题，以及各个模块的学习资料：
近20门独家秘制视频课程
、
最新顶会论文
、计算机视觉书籍
、
优质3D视觉算法源码
等。想要入门3D视觉、做项目、搞科研，欢迎扫码加入！
作者丨
王云鹤
来源丨https://zhuanlan.zhihu.com/p/9264019510
编辑丨极市平台
本文
介绍了MemoryFormer，一种新型的Transformer模型，它通过使用存储空间来替代传统的全连接层，显著降低了推理时的计算复杂度。MemoryFormer利用哈希算法和局部敏感哈希索引方法，在保持模型性能的同时，减少了模型的计算量，为大模型的高效推理提供了新的解决方案。
论文链接：
https://arxiv.org/abs/2411.12992
引言
在大模型快速发展的今天，深度学习模型的规模不断增大。然而，随着模型规模的提升，对于计算资源的消耗也随之剧增，这使得模型的部署和推理面临巨大挑战。近期，一篇来自北京大学和华为的研究人员共同合作的论文提出了一种创新性的方法——
MemoryFormer
，通过巧妙地利用存储空间来替代传统的密集计算的全连接层，显著降低了推理时的计算复杂度（FLOPs）。这一思路为大模型的高效推理提供了全新的解决方案。
用存储来替代全连接层
传统的Transformer模型中，计算资源的开销不仅来自于多头自注意力模块，还来自于模型内部的全连接层。针对这一特点，MemoryFormer提出了一种“记忆层”（Memory Layer）来代替这些计算量庞大的全连接层，从而极大地减少了模型的计算复杂度。具体来说，Memory Layer通过哈希算法来对输入特征进行离散化，并在内存中构建查找表，将原本需要的矩阵乘法替换为在计算上更轻量的哈希查找。这种设计使得MemoryFormer能够在推理时大幅降低所需的计算量，并且几乎不影响模型的性能。
MemoryFormer不仅简单地引入哈希查找表，还设计了一种基于局部敏感哈希（LSH）的极简索引方法，使得哈希表中的每一个存储的向量都可以在训练中根据梯度动态更新。这种设计确保了哈希表中存储的特征能够不断适应输入数据，并在推理阶段根据输入特征的相似性高效检索出近似的输出结果，实现全连接层所需的特征变换功能。此外，MemoryFormer通过多表分块和向量分段的方式来控制哈希表的存储规模，使得内存需求不会因哈希表的引入而暴增。
Memory Layer结构示意图
方法：可学习的局部敏感哈希
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的概率是均等的。
MemoryFormer Block示意图
在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维度一致）：
消融实验
