近期,美國麻省理工學院研究科學家、初創公司 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
營運/排版:何晨龍