ConSmax: Hardware-Friendly Alternative Softmax with Learnable Parameters
作者: Shiwei Liu, Guanchen Tao, Yifei Zou, Derek Chow, Zichen Fan, Kauna Lei, Bangfei Pan, Dennis Sylvester, Gregory Kielian, Mehdi Saligane
分类: cs.AR, cs.AI, cs.LG
发布日期: 2024-01-31 (更新: 2024-11-15)
期刊: International Conference on Computer-Aided Design 2024
🔗 代码/项目: GITHUB
💡 一句话要点
提出ConSmax以解决LLM推理中的Softmax效率问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 自注意力机制 Softmax替代 硬件友好设计 功耗优化 并行计算 机器学习 自然语言处理
📋 核心要点
- 现有的LLM推理方法在使用Softmax时面临低算术强度和处理并行性不足的问题，尤其是在长上下文情况下。
- 本文提出的ConSmax通过可微分的归一化参数消除了Softmax中的最大值搜索和分母求和，提升了并行处理能力。
- 实验结果显示，ConSmax在16nm技术下实现了3.35倍的功耗节省和2.75倍的面积节省，同时保持了与基线模型相当的准确性。
📝 摘要（中文）
自注意力机制使基于变换器的大型语言模型（LLMs）在性能上优于卷积和递归神经网络。然而，由于Softmax在自注意力中的广泛使用，实现实时LLM推理仍然具有挑战性。为此，本文提出了Constant Softmax（ConSmax），这是一种软件-硬件协同设计，作为Softmax的高效替代方案。ConSmax利用可微分的归一化参数，消除了Softmax中最大值搜索和分母求和的需求，从而实现了广泛的并行化。实验结果表明，ConSmax在16nm FinFET技术下以1250MHz工作频率实现了0.2mW的极低功耗和0.0008mm²的面积，且在保持与GPT-2模型和WikiText103数据集相当的准确度的同时，显著节省了功耗和面积。
🔬 方法详解
问题定义：本文旨在解决在大型语言模型推理中，Softmax的低算术强度和处理并行性不足的问题。现有方法在长上下文处理时效率低下，限制了实时推理的实现。
核心思路：ConSmax的核心思路是通过引入可微分的归一化参数，消除Softmax中对最大值搜索和分母求和的需求，从而实现高效的并行计算。
技术框架：ConSmax的整体架构包括可微分归一化参数的计算模块、并行处理单元和硬件设计模块。该框架支持混合精度计算，适用于不同的硬件平台。
关键创新：ConSmax的主要创新在于其硬件友好的设计，能够实现无损的非线性操作，并通过位宽分割的查找表（LUT）支持高效的计算，与传统Softmax方法相比，显著提升了并行处理能力。
关键设计：在设计中，ConSmax采用了可调节的归一化参数，优化了损失函数以适应不同的任务需求，并在网络结构上进行了改进，以支持高效的硬件实现。具体参数设置和网络结构细节在实验部分进行了详细说明。
🖼️ 关键图片
📊 实验亮点
实验结果显示，ConSmax在16nm技术下以1250MHz工作频率实现了0.2mW的功耗和0.0008mm²的面积，分别实现了3.35倍和2.75倍的节省。同时，在使用开源EDA工具链的130nm CMOS技术下，ConSmax的功耗为2.69mW，面积为0.007mm²，显示出其在不同技术平台上的优越性。
🎯 应用场景
ConSmax的研究成果在大型语言模型推理、自然语言处理和机器学习等领域具有广泛的应用潜力。其高效的Softmax替代方案能够显著提升模型推理速度和降低功耗，适用于边缘计算和移动设备等资源受限的环境，未来可能推动更高效的AI应用开发。
📄 摘要（原文）
The self-attention mechanism distinguishes transformer-based large language models (LLMs) apart from convolutional and recurrent neural networks. Despite the performance improvement, achieving real-time LLM inference on silicon remains challenging due to the extensive use of Softmax in self-attention. In addition to the non-linearity, the low arithmetic intensity significantly limits processing parallelism, especially when working with longer contexts. To address this challenge, we propose Constant Softmax (ConSmax), a software-hardware co-design that serves as an efficient alternative to Softmax. ConSmax utilizes differentiable normalization parameters to eliminate the need for maximum searching and denominator summation in Softmax. This approach enables extensive parallelization while still executing the essential functions of Softmax. Moreover, a scalable ConSmax hardware design with a bitwidth-split look-up table (LUT) can achieve lossless non-linear operations and support mixed-precision computing. Experimental results show that ConSmax achieves a minuscule power consumption of 0.2mW and an area of 0.0008mm^2 at 1250MHz working frequency in 16nm FinFET technology. For open-source contribution, we further implement our design with the OpenROAD toolchain under SkyWater's 130nm CMOS technology. The corresponding power is 2.69mW and the area is 0.007mm^2. ConSmax achieves 3.35x power savings and 2.75x area savings in 16nm technology, and 3.15x power savings and 4.14x area savings with the open-source EDA toolchain. In the meantime, it also maintains comparable accuracy on the GPT-2 model and the WikiText103 dataset. The project is available at https://github.com/ReaLLMASIC/ConSmax