MemoryFormer: Minimize Transformer Computation by Removing Fully-Connected Layers
作者: Ning Ding, Yehui Tang, Haochen Qin, Zhenli Zhou, Chao Xu, Lin Li, Kai Han, Heng Liao, Yunhe Wang
分类: cs.CL
发布日期: 2024-11-20
备注: NeurIPS2024
💡 一句话要点
MemoryFormer：通过移除全连接层最小化Transformer计算量
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer 计算复杂度 内存查找表 哈希算法 大型语言模型 模型优化 高效计算
📋 核心要点
- 现有Transformer模型计算复杂度高，模型规模持续扩大，效率提升面临挑战。
- MemoryFormer通过构建内存查找表，利用哈希算法动态检索向量子集，替代全连接层的矩阵乘法。
- 实验表明，MemoryFormer在多种基准测试中表现出有效性，验证了其降低计算复杂度的能力。
📝 摘要（中文）
为了降低大型语言模型的计算复杂度，人们在改进Transformer模型的效率方面做出了巨大努力，例如线性注意力机制和FlashAttention。然而，为了追求更高的性能，模型规模和相应的计算复杂度也在不断扩大。本文提出了一种新的Transformer架构MemoryFormer，它从一个新的角度显著降低了计算复杂度（FLOPs）。我们几乎消除了Transformer模型的所有计算，只保留了多头注意力操作所需的必要计算。这通过使用一种替代的特征转换方法来代替全连接层的线性投影来实现。具体来说，我们首先构建一组内存中的查找表，这些表存储了大量的离散向量，以代替线性投影中使用的权重矩阵。然后，我们使用哈希算法根据输入嵌入动态地检索相关的向量子集。检索到的向量组合在一起将形成输出嵌入，从而提供对全连接层中矩阵乘法运算结果的估计。与进行矩阵乘法相比，从内存中检索数据块是一种计算量小得多的操作。我们从头开始训练MemoryFormer，并在各种基准上进行了广泛的实验，以证明所提出模型的有效性。
🔬 方法详解
问题定义：现有Transformer模型为了追求更高的性能，不断扩大模型规模，导致计算复杂度显著增加。全连接层是Transformer模型中计算量巨大的模块之一，如何降低全连接层的计算量是提升Transformer效率的关键问题。
核心思路：MemoryFormer的核心思路是使用内存查找表来近似全连接层的矩阵乘法操作。通过预先存储大量的离散向量，并使用哈希算法动态检索相关的向量子集，从而避免了耗时的矩阵乘法运算。这种方法将计算密集型的矩阵乘法转换为内存访问操作，显著降低了计算复杂度。
技术框架：MemoryFormer的整体架构与标准的Transformer类似，主要区别在于使用Memory Lookup模块替代了全连接层。该模块包含以下几个步骤：1) 构建内存查找表，存储大量的离散向量；2) 使用哈希算法根据输入嵌入检索相关的向量子集；3) 将检索到的向量组合成输出嵌入。其余部分，如多头注意力机制，保持不变。
关键创新：MemoryFormer最重要的创新点在于使用内存查找表和哈希算法来近似全连接层的矩阵乘法操作。与传统的矩阵乘法相比，内存访问操作的计算量要小得多，从而显著降低了模型的计算复杂度。这种方法避免了直接进行矩阵乘法，而是通过检索和组合预先存储的向量来完成特征转换。
关键设计：内存查找表的大小是一个关键参数，它决定了模型的容量和计算复杂度之间的平衡。哈希算法的选择也会影响检索的准确性和效率。此外，如何将检索到的向量组合成输出嵌入也是一个重要的设计细节。论文中可能使用了加权平均或其他组合方法，具体细节未知。
🖼️ 关键图片
📊 实验亮点
论文通过在各种基准测试上进行实验，验证了MemoryFormer的有效性。具体的性能数据和对比基线未知，但摘要中提到MemoryFormer能够显著降低计算复杂度（FLOPs），并且是从头开始训练的，这表明该方法具有一定的通用性和实用性。实验结果表明，MemoryFormer在保持性能的同时，能够显著降低计算成本。
🎯 应用场景
MemoryFormer具有广泛的应用前景，可以应用于各种需要高效Transformer模型的场景，例如移动设备上的自然语言处理、边缘计算设备上的图像识别等。通过降低计算复杂度，MemoryFormer使得在资源受限的环境中部署大型语言模型成为可能，从而推动人工智能技术的普及和应用。
📄 摘要（原文）
In order to reduce the computational complexity of large language models, great efforts have been made to to improve the efficiency of transformer models such as linear attention and flash-attention. However, the model size and corresponding computational complexity are constantly scaled up in pursuit of higher performance. In this work, we present MemoryFormer, a novel transformer architecture which significantly reduces the computational complexity (FLOPs) from a new perspective. We eliminate nearly all the computations of the transformer model except for the necessary computation required by the multi-head attention operation. This is made possible by utilizing an alternative method for feature transformation to replace the linear projection of fully-connected layers. Specifically, we first construct a group of in-memory lookup tables that store a large amount of discrete vectors to replace the weight matrix used in linear projection. We then use a hash algorithm to retrieve a correlated subset of vectors dynamically based on the input embedding. The retrieved vectors combined together will form the output embedding, which provides an estimation of the result of matrix multiplication operation in a fully-connected layer. Compared to conducting matrix multiplication, retrieving data blocks from memory is a much cheaper operation which requires little computations. We train MemoryFormer from scratch and conduct extensive experiments on various benchmarks to demonstrate the effectiveness of the proposed model.