编辑 | 萝卜皮
通常，矩阵乘法 (MatMul) 在大型语言模型（LLM）总体计算成本中占据主导地位。随着 LLM 扩展到更大的嵌入维度和上下文长度，这方面的成本只会增加。
加州大学、LuxiTech 和苏州大学的研究人员声称开发出一种新方法，通过消除过程中的矩阵乘法来更有效地运行人工智能语言模型。这从根本上重新设计了目前由 GPU 芯片加速的神经网络操作方式。
研究人员描述了如何在不使用 MatMul 的情况下创建一个自定义的 27 亿参数模型，性能与当前最先进的 Transformer 模型相当。
该研究以「Scalable MatMul-free Language Modeling」为题，于 2024 年 6 月 4 日发布在 arXiv 预印平台。
矩阵乘法是当今大多数神经网络计算任务的核心，而 GPU 特别擅长快速执行数学运算，因为它们可以并行执行大量乘法运算。
这种能力甚至让 Nvidia 在两周前短暂地成为了全球最有价值的公司；该公司目前占据数据中心 GPU 市场约 98% 的份额，这些 GPU 通常用于为 ChatGPT 和 Google Gemini 等 AI 系统提供支持。
图示：370M 中无 MatMul 的 Transformer++ 和新方法的训练步骤损失。（来源：论文）
在最新的研究中，加州大学、LuxiTech 和苏州大学的研究人员展示了 LLM 中可以完全消除 MatMul 操作，同时在十亿参数规模下保持强劲性能。
他们通过在密集层中使用加性运算和逐元素 Hadamard 积来实现类似自注意的功能，开发了第一个可扩展的无 MatMul 语言模型 (Matmul-free LM)。
具体而言，研究人员利用三元权重消除了密集层中的 MatMul，类似于 BNN。为了从自注意力中移除 MatMul，研究人员优化了门控循环单元 (GRU)，使其仅依赖于元素级乘积。
图示：Matmul-free LM 概述。（来源：论文）
为了评估他们的方法，研究人员将他们的 MatMul-free LM 与复制的 Llama-2 样式模型（他们称之为「Transformer++」）进行了比较，涉及三种模型大小：3.7 亿、13 亿和 27 亿参数。所有模型均在 SlimPajama 数据集上进行了预训练，其中较大的模型分别在 1000 亿个标记上进行了训练。
不含 MatMul 的 LM 在多个基准任务上与 Llama 2 基线相比取得了具有竞争力的性能，包括回答问题、常识推理和物理理解。
实验表明，该团队提出的无 MatMul 模型的性能与最先进的 Transformer 模型相当，后者在推理过程中需要更多内存。
为了量化轻量级模型的硬件优势，除了定制的 FPGA 加速器外，研究人员还提供了优化的 GPU 实现。通过在三元密集层的 GPU 实现中使用融合内核，与 GPU 上未优化的基线相比，训练速度加快了 25.6%，内存消耗减少了高达 61.0%。
此外，通过采用低位优化的 CUDA 内核，当模型扩展到 13B 参数时，推理速度提高了 4.57 倍，内存使用量减少了 10 倍。
为了正确量化该架构的效率，研究人员在 FPGA 上构建了一个自定义硬件解决方案，该解决方案利用了 GPU 无法处理的轻量级操作。
研究人员演示了如何在 GPU 上以每秒 23.8 个 token 的速度运行 13 亿个参数的模型；该方法以 13 瓦的功耗（不计算 GPU 的功耗）处理了十亿参数规模的模型，超出了人类可读的吞吐量，使 LLM 更接近类似大脑的效率。
这项工作不仅展示了 LLM 在保持有效运行的情况下可以被剥离到何种程度，而且还指出了未来加速器在处理下一代轻量级 LLM 时应该优化的操作类型。
不过需要明确的是，拥有 27 亿个参数的 Llama-2 模型与目前市场上最好的 LLM（例如 GPT-4）相差甚远，据估计 GPT-4 总共拥有超过 1 万亿个参数。因此，这里还没有在这里讨论 ChatGPT 级别的处理能力。
参数数量通常意味着模型的复杂性（以及大致上的能力）更高，研究人员一直在寻找用更少的参数实现更高级别 LLM 性能的方法。
研究人员表示，他们在实验中观察到的缩放规律表明，无 MatMul 的 LM 在非常大规模下的表现也可能优于传统 LLM。
研究人员预测，他们的方法在理论上可以与标准 LLM 相媲美，并且超越其在 10²³ FLOPS 左右的规模上的性能，这大致相当于 Meta 的 Llama-3 8B 或 Llama-2 70B 等模型所需的训练计算量。
然而，该团队也指出他们的工作有局限性。由于计算限制，无 MatMul 的 LM 尚未在超大规模模型（例如 1000 亿多个参数）上进行测试。他们呼吁拥有更多资源的机构投资扩大规模并进一步开发这种轻量级的语言建模方法。
论文链接：https://arxiv.org/abs/2406.02528
相关报道：https://arstechnica.com/information-technology/2024/06/researchers-upend-ai-status-quo-by-eliminating-matrix-multiplication-in-llms/
