第673章 矩陣乘法無需相乘,速度提升100倍
在一篇被ICML2021接收的論文中,MIT的一位計算機科學博士生及其業界大佬導師為矩陣乘法引入了一種基於學習的算法,該算法具有一個有趣的特性——需要的乘加運算為零。在來自不同領域的數百個矩陣的實驗中,這種學習算法的運行速度是精確矩陣乘積的100倍,是當前近似方法的10倍。
矩陣乘法是機器學習中最基礎和計算密集型的操作之一。因此,研究社區在高效逼近矩陣乘法方面已經做了大量工作,比如實現高速矩陣乘法庫、設計自定義硬件加速特定矩陣的乘法運算、計算分佈式矩陣乘法以及在各種假設下設計高效逼近矩陣乘法(AMM)等。
在MIT計算機科學博士生DavisBlalock及其導師JohnGuttag教授發表的論文《MultiplyingMatricesWithoutMultiplying》中,他們為逼近矩陣乘法任務引入了一種基於學習的算法,結果顯示該算法顯著優於現有方法。在來自不同領域的數百個矩陣的實驗中,這種學習算法的運行速度是精確矩陣乘積的100倍,是當前近似方法的10倍。這篇論文入選了機器學習頂會ICML2021。
此外,在一個矩陣提前已知的常見情況下,研究者提出的方法還具有一個有趣的特性——需要的乘加運算(multiply-adds)為零。
這些結果表明,相較於最近重點進行了大量研究與硬件投入的稀疏化、因式分解和/或標量量化矩陣乘積而言,研究者所提方法中的核心操作——哈希、求平均值和byteshuffling結合可能是更有前途的機器學習構建塊。
也有網友表示:「這篇論文為實現更高效的AI打開了一扇門。」
對於有網友提到的「該研究在硬件實現方面似乎很有發展前景」,一作本人現身reddit並給出了回復:「我們的編碼表示是密集矩陣,所以佈局和訪問模式看上去基本與GEMM內核相同,也就意味着可以很容易地使用脈動陣列或修正張量核心來實現。在x86上,一般只需要一個vpshufb-add指令和一個4-bit解包指令就可以了。」
具體來說,該研究專註於AMM任務,並假設矩陣是高的(tall),並且相對密集,存在於單個機器內存中。在這種設置下,研究者遇到的主要挑戰是在給定保真度水平下最小化近似線性運算所需的計算時間。這種設置會很自然地出現在機器學習和數據挖掘中,當一個數據矩陣A的行是樣本,而一個線性算子B希望應用這些樣本,B可以是一個線性分類器、線性回歸器,或嵌入矩陣,以及其他可能性。
舉例來說,考慮一個近似softmax分類器的任務,以預測來自神經網絡嵌入的圖像標籤。在這裏,A的行是每個圖像的嵌入,B的列是每個類的權值向量。分類是通過計算乘積AB並在結果的每一行中取argmax來執行的。圖1結果表明,在CIFAR-10和CIFAR-100數據集上,使用該研究的方法及其最佳性能競爭對手的方法近似AB的結果。
該研究所用方法與傳統方法背離,傳統的AMM方法構造矩陣V_A,V_B∈R^(D×d),d<
通常,V_A、V_B是稀疏的,包含某種採樣方案,或者具有其他結構,使得這些投影操作比密集矩陣乘法更快。簡而言之,這些方法使用線性函數對A和B進行預處理,並將問題簡化為低維空間中的精確矩陣乘法。
該研究提出MADDNESS方法,該方法採用非線性預處理函數,將問題簡化為查表。此外,在B提前已知的情況下,即將訓練好的線性模型應用於新數據等情況時,MADDNESS不需要任何乘-加運算。該方法與用於相似性搜索的矢量量化方法密切相關。然而,該研究沒有使用太多的乘-加量化函數,而是引入了一系列不需要乘-加的量化函數。
一個高效的學習矢量量化函數族,可以在單個CPU線程中每秒編碼超過100GB的數據。
一種用於低位寬整數(low-bitwidthintegers)的高速求和算法,可避免upcasting、飽和和溢出。
基於這些函數的近似矩陣乘法算法。數百個不同矩陣的實驗表明,該算法明顯優於現有替代方案。並且還具有理論質量保證。
為了評估MADDNESS的有效性,研究者用c++和Python實現了該算法和其他幾個現有算法。評估結果如下。
MADDNESS到底有多快?
研究者首先分析了MADDNESS的原始速度。在圖3中,他們為各種矢量量化方法計算g(A)函數的時間,結果表明,MADDNESS比現有方法快兩個數量級,其吞吐量隨行的長度而增加。
從上圖可以看出,Bolt(藍色虛線)是與MADDNESS最接近的競爭對手。研究者還使用與Bolt相同的基線分析了聚合函數f(·,·)的速度。如圖4所示,他們基於average的、matrix-aware的聚合方法明顯快於Bolt基於upcasting的方法。
前面說過,研究者在廣泛使用的CIFAR-10和CIFAR-100數據集上近似線性分類器。如圖5所示,MADDNESS顯著優於所有現有方法,幾乎達到了與精確乘法相同的準確率,但比精確乘法快了一個數量級。而且,MADDNESS是在硬件支持較差的情況下實現了這種性能。
為了評估該方法在更大、多樣性更強的數據集上的表現,研究者在來自UCRTimeSeriesArchive的數據集上訓練了kernel分類器。結果如下圖6所示,MADDNESS在給定準確率的情況下明顯快於替代方案。
為了測試MADDNESS的極限,研究者對將小濾波器應用於圖像的各種技術能力進行了基準測試。結果如下圖7所示,只有MADDNESS比精確矩陣乘積更有優勢。
DavisBlalock是麻省理工學院的博士生,由JohnGuttag教授指導。他的主要研究方向是設計高性能的機器學習算法,以減少人們在機器學習速度、準確性、私隱和安全性之間的妥協。他於2014年獲得弗吉尼亞大學學士學位,2016年獲得麻省理工學院碩士學位。他是QualcommInnovationFellow、NSFGraduateResearchFellow和BarryM.GoldwaterScholar。
JohnGuttag是麻省理工學院計算機科學與電氣工程教授、ACMFellow。他領導着麻省理工學院計算機科學和人工智能實驗室(CSAIL)的臨床與應用機器組。該小組負責開發和應用先進的機器學習和計算機視覺技術,以解決各種臨床相關問題,目前的研究項目包括預測和減少不良醫療事件,幫助患者匹配治療方法和治療者,以及醫學成像。
。