Analog In-Memory Computing Attention Mechanism for Fast and Energy-Efficient Large Language Models

📄 arXiv: 2409.19315v2 📥 PDF

作者: Nathan Leroux, Paul-Philipp Manea, Chirag Sudarshan, Jan Finkbeiner, Sebastian Siegel, John Paul Strachan, Emre Neftci

分类: cs.NE, cs.AI, cs.AR, cs.ET

发布日期: 2024-09-28 (更新: 2024-11-25)

备注: 25 pages, 6 figures, 1 table

💡 一句话要点

提出基于模拟存内计算的Attention机制，加速节能的大语言模型推理。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 存内计算 模拟电路 自注意力机制 大型语言模型 低功耗 Transformer 增益单元

📋 核心要点

  1. 现有GPU存储的Transformer模型在生成新token时，需要频繁地将投影数据从GPU加载到SRAM，造成显著的延迟和能量消耗。
  2. 论文提出一种基于增益单元的模拟存内计算架构，直接在存储器中进行自注意力计算，避免了数据搬运，从而降低延迟和功耗。
  3. 实验结果表明，该架构在注意力计算的延迟和能耗方面，相比GPU分别降低了高达两个和五个数量级，性能提升显著。

📝 摘要（中文）

Transformer网络，特别是其中的自注意力机制，是大型语言模型的核心。在生成式Transformer中，自注意力使用缓存来存储token的投影，避免在每个时间步重复计算。然而，存储在GPU中的投影必须加载到SRAM中才能进行每个新的生成步骤，这导致了延迟和能量瓶颈。本文提出了一种定制的自注意力存内计算架构，该架构基于新兴的基于电荷的存储器（称为增益单元），该存储器可以有效地写入以在序列生成期间存储新的token，并实现自注意力所需的并行模拟点积计算。然而，模拟增益单元电路引入了非理想性和约束，阻止了预训练模型的直接映射。为了规避这个问题，我们设计了一种初始化算法，该算法实现了与GPT-2相当的文本处理性能，而无需从头开始训练。我们的架构分别将注意力延迟和能耗降低了高达两个和五个数量级，标志着朝着超快速、低功耗生成式Transformer迈出了重要一步。

🔬 方法详解

问题定义：现有的大型语言模型，特别是生成式Transformer，在推理过程中需要频繁访问存储在GPU上的token投影。每次生成新的token，都需要将这些投影加载到SRAM中进行自注意力计算，这导致了显著的延迟和能量消耗，成为模型推理的瓶颈。

核心思路：论文的核心思路是利用新兴的增益单元存储器，构建一个模拟存内计算架构。通过将token投影直接存储在增益单元中，并在存储器内部进行自注意力计算所需的点积运算，避免了数据在GPU和SRAM之间的频繁搬运，从而降低了延迟和功耗。

技术框架：该架构主要包含以下几个关键模块：1) 增益单元存储器阵列，用于存储token投影；2) 模拟计算电路，用于在存储器内部进行点积运算；3) 初始化算法，用于将预训练模型的参数映射到增益单元存储器中。整体流程是：首先，使用初始化算法将预训练模型的参数映射到增益单元存储器中。然后，在推理过程中，新的token投影被写入增益单元。最后，通过模拟计算电路在存储器内部进行自注意力计算，得到最终的输出。

关键创新：最重要的技术创新点在于利用模拟存内计算架构，将计算过程集成到存储器内部，从而避免了数据搬运带来的延迟和功耗。与传统的基于GPU的计算架构相比，该架构能够实现更高的计算效率和更低的功耗。此外，针对模拟电路的非理想性，论文还设计了一种初始化算法，保证了模型的性能。

关键设计：论文设计了一种特殊的初始化算法，用于将预训练模型的参数映射到增益单元存储器中。由于模拟电路存在非理想性，直接映射预训练模型的参数会导致性能下降。因此，初始化算法需要对参数进行调整，以适应模拟电路的特性。具体的参数设置和网络结构细节在论文中没有详细描述，属于未知信息。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，基于模拟存内计算的自注意力架构在延迟和能耗方面相比GPU分别降低了高达两个和五个数量级。该架构在文本处理性能上与GPT-2相当，但无需从头开始训练，显著提升了效率。这些结果表明，该架构在加速和节能大型语言模型推理方面具有显著优势。

🎯 应用场景

该研究成果可应用于各种需要快速、低功耗的大型语言模型推理场景，例如移动设备、边缘计算设备和嵌入式系统。通过降低延迟和功耗，可以实现更高效的自然语言处理应用，例如智能助手、机器翻译和文本生成等。该技术还有潜力应用于其他类型的神经网络，加速其推理过程。

📄 摘要（原文）

Transformer networks, driven by self-attention, are central to Large Language Models. In generative Transformers, self-attention uses cache memory to store token projections, avoiding recomputation at each time step. However, GPU-stored projections must be loaded into SRAM for each new generation step, causing latency and energy bottlenecks. We present a custom self-attention in-memory computing architecture based on emerging charge-based memories called gain cells, which can be efficiently written to store new tokens during sequence generation and enable parallel analog dot-product computation required for self-attention. However, the analog gain cell circuits introduce non-idealities and constraints preventing the direct mapping of pre-trained models. To circumvent this problem, we design an initialization algorithm achieving text processing performance comparable to GPT-2 without training from scratch. Our architecture respectively reduces attention latency and energy consumption by up to two and five orders of magnitude compared to GPUs, marking a significant step toward ultra-fast, low-power generative Transformers.