近年来，AI 模型性能不断提升——但规模也变得越来越大。神经网络权重占用的内存量持续增长，有些模型拥有 5000 亿个参数，甚至达到数万亿。在传统计算机架构上执行推理时，每次移动这些权重进行计算都会消耗时间和能量。模拟内存计算将内存与计算融为一体，消除了这一瓶颈，在节省时间和能量的同时仍能提供卓越的性能。

在三篇新论文中，IBM Research 的科学家展示了他们的研究成果：面向大模型的可扩展 3D 模拟内存架构硬件、用于紧凑边缘规模模型的相变存储器，以及加速 transformer 注意力机制的算法创新。

根据 IBM Research 团队的一项新研究，基于模拟内存计算的芯片特别适合运行前沿的混合专家（MoE）模型，在多项指标上优于 GPU（图形处理器）。他们的工作登上了《Nature Computational Science》期刊的封面，表明 MoE 网络每一层中的每个专家都可以映射到 3D 非易失性存储器的物理层上，这种独特的类脑芯片架构存在于基于 3D 的模拟内存计算芯片中。通过大量的数值模拟和基准测试，团队发现这种映射可以在运行 MoE 模型时实现卓越的吞吐量和能效。

连同 IBM Research 的另外两篇新论文，这展示了内存计算在为采用 transformer 架构的 AI 模型提供动力方面的前景，无论是边缘还是企业云应用。根据这些新论文，将这项实验性技术从实验室推向市场的时间已经临近。

专业知识的层叠

“将模拟内存计算带入第三维度，确保即使是大型 transformer 架构的模型参数也能完全存储在芯片上，”IBM Research 科学家 Julian Büchel 说道，他是这篇 MoE 论文的第一作者，该论文证明了将 MoE 的每个“专家”堆叠在 3D 模拟内存计算单元中是有益的。

在 MoE 模型中，神经网络的特定层可以被拆分为更小的层。每个更小的层被称为一个“专家”，指的是它专门处理数据的一个子集。当输入到来时，路由层决定将数据发送给哪个（或哪些）专家。当他们将两个标准 MoE 模型运行在性能模拟工具中时，模拟硬件的表现优于最先进的 GPU。

“这样，你可以更好地扩展神经网络，并以小得多的计算足迹部署一个大型且功能强大的神经网络，”IBM Research 科学家 Abu Sebastian 说道，他领导了这些新论文背后的团队。“可以想象，这也将推理所需的计算量降至最低。”Granite 1B 和 3B 使用这种模型架构来降低延迟。

在这项新研究中，他们使用模拟硬件将 MoE 网络的各层映射到模拟内存计算单元上，每个单元由多个垂直堆叠的层级组成。这些包含模型权重的层级可以单独访问。在论文中，团队将这些层比作一栋高层办公楼，不同楼层包含不同的专家，可以根据需要随时调用。

将专家层堆叠到不同的层级上很直观，但这一策略的结果才是关键。在他们的模拟中，与运行相同模型的商用 GPU 相比，3D 模拟内存计算架构在运行 MoE 模型时实现了更高的吞吐量、更高的面积效率和更高的能效。在能效方面的优势最为显著，因为 GPU 在内存和计算之间移动模型权重时牺牲了大量时间和能量——而这一问题在模拟内存计算架构中并不存在。

据参与这项工作的 IBM Research 科学家 Hsinyu (Sidney) Tsai 表示，这是朝着成熟 3D 模拟内存计算迈出的关键一步，最终可以加速云环境中的企业 AI 计算。

边缘端的 Transformer（深度学习中的注意力模型架构）

团队研究的第二篇论文是一项加速器架构研究，于12月在 IEEE 国际电子器件会议（IEEE International Electron Devices Meeting）的受邀演讲中展示。他们展示了在边缘应用中利用超低功耗设备进行 AI 推理的可行性，这些设备使用相变存储器（Phase-change memory, PCM）通过硫系玻璃的导电性来存储模型权重。当更多电压通过玻璃时，它从晶态转变为非晶态固体。这使其导电性降低，从而改变矩阵向量乘法运算的值。相变存储器（PCM）用于模拟存内计算——高效地将 Transformer 模型带到边缘端。在团队的分析中，他们的神经处理单元对 Transformer 模型展现出有竞争力的吞吐量，并具有预期的能效优势。

"在边缘设备中，你面临能源约束、成本约束、面积约束和灵活性约束，"参与该论文的 IBM Research 科学家 Irem Boybat 表示。"因此我们提出了这种架构，以满足边缘 AI 的这些需求。"她和同事概述了一种神经处理单元，混合了基于 PCM 的模拟加速器和数字加速器节点，二者协同工作以处理不同级别的精度。

Boybat 表示，由于这种灵活的架构，多种神经网络可以在这些设备上运行。在本文中，她和同事探索了一种为移动设备定制的 Transformer 模型 MobileBERT。根据他们自己的吞吐量基准测试，团队提出的神经处理单元表现优于市场上现有的低成本加速器，并且在 MobileBERT 推理基准测试中接近某些高端智能手机的性能。

Sebastian 表示，这项工作代表着向未来迈出了一步——届时模拟存内计算设备可以低成本大规模生产，将 AI 模型的所有模型权重存储在芯片上。此类设备可以构成微控制器的基础，帮助边缘应用的 AI 推理，例如自动驾驶汽车的摄像头和汽车传感器。

模拟中的 Transformer

最后但同样重要的是，研究人员概述了 Transformer 架构在模拟存内计算芯片上的首次部署，涵盖了涉及静态模型权重的每一次矩阵向量乘法运算。与所有运算均以浮点执行的场景相比，它在一个名为 Long Range Arena 的基准测试中精度差距在 2% 以内，该基准测试评估长序列上的准确性。研究结果发表在《Nature Machine Intelligence》期刊上。

从更宏观的角度来看，这些实验表明，可以用模拟存内计算来加速注意力机制——这是 Transformer 的主要瓶颈，IBM Research 科学家 Manuel Le Gallo-Bourdeau 表示。"Transformer 中的注意力计算必须完成，而这并非可以在模拟中直接加速的事情，"他补充道。障碍在于注意力机制中需要计算的值。它们是动态变化的，这将需要不断重新编程模拟设备——从能源和耐久性角度来看这是一个不切实际的目标。

为了克服这一障碍，他们使用了一种称为核近似（kernel approximation）的数学技术，在他们的实验模拟芯片上执行非线性函数。Sebastian 表示，这一进展很重要，因为此前人们认为这种电路架构只能处理线性函数。该芯片采用受大脑启发的设计，将模型权重存储在相变存储器设备中，以交叉阵列方式排列，类似于 MoE 实验中模拟的系统。

"注意力计算是一种非线性函数，对任何 AI 加速器来说都是一种非常棘手的数学运算，尤其是对模拟存内计算加速器，"Sebastian 表示。"但这证明了我们可以通过使用这个技巧来实现它，并且我们还可以提高整体系统效率。"

这个技巧——核近似——通过使用随机采样的向量将输入投影到更高维空间，然后在所得的高维空间中计算点积，从而绕过了对非线性函数的需求。核近似是一种通用技术，可以应用于多种场景，不仅限于使用模拟存内计算的系统，但它恰好在这方面效果非常好。

"这些论文为未来现代 AI 工作负载可以在云端和边缘端同时运行提供了关键突破，"IBM 院士 Vijay Narayanan 评论道。

注释

• 注释 1：相变存储器（PCM）设备通过硫系玻璃的导电性来存储模型权重。当更多电压通过玻璃时，它会从晶态重新排列为非晶态固体。这使其导电性降低，从而改变矩阵向量乘法运算的值。↩︎

转载声明：本文为转载整理，原文版权归原作者及原站点所有。

原文来源：IBM Research Blog

原创地址：https://research.ibm.com/blog/how-can-analog-in-memory-computing-power-transformer-models