arXiv:2411.12992cs.CL2024-11NeurIPS被引 7

用内存查表替代全连接层,大幅降低Transformer计算量

MemoryFormer: Minimize Transformer Computation by Removing Fully-Connected Layers

  • 用内存中的离散向量表替代全连接层权重,通过哈希动态检索
  • 在多个基准上实现显著更低的浮点运算量(FLOPs),性能接近原模型
  • 适合追求高效推理、资源受限场景下的大模型部署

为降低大语言模型的计算复杂度,现有工作多聚焦于改进注意力机制如线性注意力或Flash-Attention。然而,模型规模与计算开销仍在持续增长以追求更高性能。本文提出MemoryFormer,一种新型Transformer架构,从新视角显著降低计算复杂度(FLOPs)。其核心是移除几乎全部非必要计算,仅保留多头注意力所需操作。为此,我们引入一种替代方案:用存储大量离散向量的内存查找表替代全连接层中的线性投影权重矩阵,并通过哈希算法根据输入嵌入动态检索相关向量子集。这些检索出的向量组合后形成输出嵌入,近似完成全连接层的矩阵乘法结果。相比矩阵乘法,从内存中检索数据块的计算成本极低。我们从零训练MemoryFormer,且在多个基准上进行充分实验,验证了该方法的有效性。

原文摘要 · Abstract (English)

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.

Transformer优化内存计算低延迟推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。