近期,美国麻省理工学院研究科学家、初创公司 BitEnergy AI 创始人罗鸿胤和合作者提出了一种名为线性复杂度乘法(L-Mul,linear-complexity multiplication)的算法。
L-Mul 的最大特点是用整数加法运算逼近浮点数乘法,通过改变数值计算的方式让大模型的计算精度大幅度降低。
用这种“偷懒式计算”可以大幅减少模型的运算量,最高可降低 95% 能耗。
此外,该算法也适用无需极高精度计算的场景。
其潜在的商业价值在于:
对于数据中心而言,显著的能耗减少可以让数据中心在相同的能耗下支持更大算力。
值得一提的是,L-Mul 还可用于具身智能和边缘计算等需要 AI 芯片的场景,比如应用到机器人、笔记本电脑和手机中。
另一方面,L-Mul 为简化设计芯片带来新的思路。
它通过移除浮点数乘法器,让芯片的设计和制造更加简单,帮助芯片制造厂商提高芯片产品的质量和功耗稳定性。
罗鸿胤举例说道:“基于 L-Mul,同样面积的芯片有望能搭载更多的计算单元,有可能用一个 5 纳米的芯片就能达到更先进制程芯片的计算速度。”
有业内人士对该研究评价道:“L-Mul 基本上可以直接替代模型中的所有乘法运算。尽管其基准测试主要集中在推理阶段,但测试使用了诸如 Llama 3.1 8b 等实际模型,并在常规基准测试中展现了几乎相同的性能。”
日前,相关论文以《加法是构建高能效语言模型的全部所需》(Addition is All You Need for Energy-efficient Language Models)为题发表在预印本网站 arXiv 上 [1]。
图丨相关论文(来源:arXiv)
“变乘法为加法”
当下,AI 技术正在加速发展,随着大型神经网络模型的应用越来越广泛,它们在训练和推理过程中的能源消耗问题日益突出。
与此同时,能耗问题也正在成为新数据中心最大的瓶颈,这一点在数据中心实际选址方面可以得到印证。
大型数据中心往往需要充足、稳定的电力供应和较低的运营成本,因而基本会选址在人口稀少和能够提供电力保障的地方。
比如美国国家安全局在美国犹他州布拉夫代尔镇建造了全美最大的数据中心;在中国,华为、腾讯和苹果等将数据中心选在贵州省。
有人因此提出“AI 的尽头是天然气、煤炭和核电站”的观点。
尽管云端数据中心的优化和硬件加速取得了诸多突破,但在浮点数乘法这一核心计算操作上的改进却相对保守。
在浮点数运算中,每个数通常表示为符号位、指数和尾数。乘法操作通常需要对这些组成部分进行相应的运算,包括指数相加和尾数相乘、可能的规格化和舍入步骤。
而 L-Mul 通过省略尾数乘法,仅使用整数加法和一些简单的位移操作来实现相同的计算效果,从而显著降低了计算复杂度和能源消耗。
图丨在两个 fp32 数之间复制正则浮点乘法和线性复杂度乘法(L-Mul)的过程(来源:arXiv)
罗鸿胤和团队以降低模型推理所需的浮点数计算量为出发点,尝试提高模型的推理能效。
在 Transformer 模型中,使用了大量的张量乘法运算。
他推测,模型中的大量组件只关心输入张量各维度间的大小关系,而对其具体的精确数值不太敏感。
例如,注意力机制和负责预测下一个字(token)的解码层都只关心数值较大的维度。
反之,数值较小的乘法结果精确与否,对模型整体性能的影响则可以忽略。
而注意力机制工作方式是,在输入一个高维向量后,它会在高维空间中寻找与之相近的向量。
简单来讲,就像是在地图上标记一个坐标,并查询该坐标附近的餐馆。而“附近”范围以外的某地具体需要开车 5 个小时或是 2 天,对查询者而言没有太大意义。
基于此原理,此前与模型效率相关的研究工作,例如模型量化(减少模型参数的比特位数)或者模型剪枝(减少非零参数量)都取得了不错的效果。
那能不能突破基本运算的桎梏,进行一些更极端的尝试?
于是,研究人员尝试将所有的乘法都替换为加法,并发现经过一定计算后仍然能保持准确。随后,他们开始做理论推导和数值分析,以揭示其保持准确的原因。
罗鸿胤介绍:“L-Mul 的巧妙之处在于,它的数学原理非常直观,同时对应了最简洁的硬件实现;它的算法复杂度低于 8 比特浮点乘法,却能达到更高的计算精度。”
实验验证:L-Mul 在多种任务上表现优异
为了验证 L-Mul 算法的精确性,并探究基于 L-Mul 的大型模型在实际任务中的效果,研究人员在各种基准测试中对不同的模型进行了实验。
作者评估了 Llama-3.1-8b-Instruct、mistral-7b-v0.3-Instruct、Gemma2-2b-It 和 Llava-v1.5-7b 模型,并发现所提出的方法可以在微调或无需训练的情况下,替换 Transformer 层中的不同模块。
在自然语言推理任务中,基于 L-Mul 的注意力机制的性能损失约 0.07%,涵盖了常识、结构化推理和语言理解。
而在视觉任务中,基于 L-Mul 的注意力在视觉问题回答、对象幻觉和自由形式视觉指令任务上提高了 0.12% 的准确率。
图丨比较线性复杂度乘法(L-Mul)的错误水平与不同格式的 8 位 FP 乘法运算的尾数位数(来源:arXiv)
值得关注的是,这些实验结果通过直接切换预训练大模型的标准注意力到基于 L-Mul 的新机制获得,无需额外训练。
误差估计和消融研究表明,在无需训练的情况下,4 位尾数的 L-Mul 可以实现与 float8 e4m3 乘法相当的精度,而 3 位尾数的 L-Mul 超过了 float8 e5m2 乘法。
相关实验还表明,微调能够对 L-Mul 和标准乘法之间的性能差距进行弥补。
在涉及注意力机制、线性变换和逐元素乘积的操作中,如果将所有的乘法操作用 3 位尾数的 L-Mul 操作来替代,那么在微调阶段,这种模型的性能与使用 float8 e4m3 精度作为标准的模型相当。
模型加速领域的一种“非主流”方法
目前,很多大型企业和初创公司都在积极探索模型加速,以提高运算效率和降低成本。
为减少模型计算量,业界正在研制 4 比特芯片,并通过模型量化技术用更少的比特表示相关数值。
另一方面,一些企业通过加速优化芯片的数据读写实现模型速度和能耗的优化,例如谷歌张量处理单元(TPU,Tensor Processing Unit)。
此外,研究人员也在探索非 Transformer 架构。
L-Mul 与当前模型加速领域的主流方法相比,显得“与众不同”。罗鸿胤表示:“我们的方法能与其他提高模型效率的技术进行垂直补充,互相配合进一步提升整体效率,并非互斥的关系。”
例如,在优化输入/输出芯片后,可以用 L-Mul 算法进一步优化计算复杂度;量化后的模型可用 L-Mul 优化模型的乘法;非 Transformer 架构也可以用 L-Mul 的思路加速乘法运算。
图丨比较用 16 位和 8 位张量乘法运算和 L-Mul 近似实现的注意力机制(来源:arXiv)
在理论和数值模拟方面,L-Mul 算法已显示出优异的性能。
虽然目前由于缺乏对应的硬件指令,现有硬件暂不支持浮点数直接进行 L-Mul 运算。但是,只需在硬件层面增加一个新的简单指令就可实现 L-Mul 算法,从而获得显著的能效提升。
据介绍,目前有相关课题组在软件层面通过用中央处理器(CPU,Central Processing Unit)的 fp32 操作,已实现一些能耗降低。
从对 AI 和机器学习领域的更广泛的影响来看,L-Mul 有望减轻 AI 行业的算力垄断问题,降低用户算力成本,并提供更多的算力资源。
罗鸿胤表示:“我们希望从改进数值计算和数值稳定性的角度,提升模型的训练效果、降低存储空间、优化推理效率。”
因此,L-Mul 有望助力提高模型训练效率,减少数值不稳定带来的训练困难,从而提升模型量化、模型剪枝等领域的算法效率和存储空间利用率。
将继续致力于大模型效率的提高
罗鸿胤在清华大学获得学士学位,师从刘知远教授和孙茂松教授。2022 年在麻省理工学院电子工程与计算机科学系获得博士学位,研究方向为语言模型的自我训练。毕业后,他留校继续担任博士后和研究员。
此前,他构建了一个参数量仅为 3.5 亿的小模型,该模型完全使用合成数据训练,并在文本分类任务性能方面的表现超过了拥有 1750 亿参数量级的 GPT-3[2]。之后,他致力于将这种效率提升方法转移到生成模型。
图丨罗鸿胤(来源:罗鸿胤)
目前,罗鸿胤主要聚焦于提高 AI 的效率和推理能力两个方向。
在效率方面,他特别关注将大模型的建模粒度从向量深入到比特层面,通过模型架构与计算架构的协同设计提升 AI 效率;在推理能力方面,他关注模型的编程能力和容错能力。
罗鸿胤表示:“编程可以让模型通过复用推理过程提升效率,而容错能力则使得模型能够一次性输出,减少重复推理的次数,从而节省计算资源。”
在未来的阶段,他计划在可编程阵列逻辑(FPGA,Field Programmable Gate Array)平台进行仿真研究,对数值计算变化后能耗的具体降低程度进行确认,以提供准确的数据支持。
“我们的长期目标是希望通过数值计算研究,解决大模型的各种特性和目前面临的相关挑战,从而促进大模型效率的指数级提高。”罗鸿胤说道。
参考资料:
1.https://arxiv.org/abs/2410.00907
2.http://arxiv.org/pdf/2305.17197
运营/排版:何晨龙