矩陣乘法新突破!XX^T原來(lái)可以更快!RL助力搜索,世界紀(jì)錄又被提升了5%
深圳市大數(shù)據(jù)研究院、香港中文大學(xué)(深圳)研究團(tuán)隊(duì)最新研究發(fā)現(xiàn), 這類特殊的矩陣乘法可以進(jìn)一步加速,并在強(qiáng)化學(xué)習(xí)與組合優(yōu)化技術(shù)的結(jié)合下發(fā)掘出了一種新的算法,節(jié)省 5% 的乘法數(shù)量。
- 論文標(biāo)題:XXt Can Be Faster
- 論文鏈接:https://arxiv.org/abs/2505.09814
該成果在國(guó)際社交媒體平臺(tái) X 引發(fā)熱烈討論,并引起 MIT、斯坦福、哈佛及 Google DeepMind 科學(xué)家的廣泛關(guān)注。
背景
矩陣乘法優(yōu)化堪稱計(jì)算機(jī)科學(xué)領(lǐng)域的「珠穆朗瑪峰」。自 1969 年 Strassen 算法橫空出世以來(lái),這個(gè)充滿組合爆炸可能性的數(shù)學(xué)迷宮就持續(xù)考驗(yàn)著人類智慧的邊界。
Google DeepMind 為此專門投入四年心血,先后推出 AlphaTensor、AlphaEvolve 等機(jī)器學(xué)習(xí)系統(tǒng)來(lái)攻克這一難題。這就像短跑運(yùn)動(dòng)員將百米紀(jì)錄從 9.58 秒推進(jìn)到 9.57 秒——每個(gè) 0.01 秒的突破背后,都是對(duì)計(jì)算理論極限的重新定義。
(矩陣乘以自身的轉(zhuǎn)置)這類特殊的矩陣乘法廣泛存在于各類數(shù)據(jù)科學(xué)的實(shí)際應(yīng)用中,實(shí)際應(yīng)用包括:
- 5G 與自動(dòng)駕駛定制芯片設(shè)計(jì)
- 線性回歸與數(shù)據(jù)分析
- 大語(yǔ)言模型訓(xùn)練算法(Muon、SOAP)
這類操作每分鐘在全球執(zhí)行數(shù)萬(wàn)億次,假如能減少該操作的計(jì)算量,對(duì)能耗開銷可以帶來(lái)相當(dāng)可觀的節(jié)省。令人驚訝的是,相比于普適的矩陣乘法 AB,研究者對(duì)于
這類的特殊矩陣乘法的關(guān)注少之又少。Google DeepMind 的 AlphaTensor、AlphaEvolve 探索了帶有特殊結(jié)構(gòu)的 AB 矩陣乘法,但他們尚未匯報(bào)任何關(guān)于
的結(jié)果。
通過(guò)觀察 運(yùn)算的特殊結(jié)構(gòu),該團(tuán)隊(duì)發(fā)現(xiàn)
的計(jì)算確實(shí)存在加速空間!
主要貢獻(xiàn)
在 AI 技術(shù)的輔助下,研究團(tuán)隊(duì)發(fā)掘了新算法(RXTX),以讓 這一常見的底層操作減少 5% 的運(yùn)算量,這可以進(jìn)一步轉(zhuǎn)換成節(jié)省 5% 的能耗以及時(shí)間(特別的,能耗開銷主要由乘法運(yùn)算數(shù)量決定)。值得一提的是,RXTX 的 5% 加速不僅對(duì)超大規(guī)模矩陣成立,對(duì)小規(guī)模矩陣也成立,比如:RXTX 對(duì) 4x4 矩陣 X 僅需 34 次乘法運(yùn)算。此前最先進(jìn)的 Strassen 算法需要 38 次乘法(減少 10% 運(yùn)算量)。
乘法運(yùn)算量復(fù)雜度分析
研究團(tuán)隊(duì)對(duì)乘法運(yùn)算量的復(fù)雜度進(jìn)行了分析。分析結(jié)果表明,RXTX 的漸進(jìn)常數(shù) 26/41≈0.63,較先前最優(yōu)值 2/3≈0.66 降低 5%。
總運(yùn)算量(乘法+加法)復(fù)雜度分析
研究團(tuán)隊(duì)進(jìn)一步提供了總運(yùn)算量(乘法+加法)的復(fù)雜度分析。分析結(jié)果表明,當(dāng) n≥256 時(shí),RXTX 的總加法與乘法次數(shù)也少于現(xiàn)有最優(yōu)方案,且漸進(jìn)意義下約有 5% 的穩(wěn)定提升。
核心技術(shù)
該方法屬于基于神經(jīng)網(wǎng)絡(luò)的大鄰域搜索方法框架:
- 利用強(qiáng)化學(xué)習(xí)策略生成候選雙線性乘積
- 構(gòu)建組合問(wèn)題一(MILP-A):將目標(biāo)表達(dá)式構(gòu)建為候選乘積的線性組合
- 構(gòu)建組合問(wèn)題二(MILP-B):篩選能完整表達(dá)
結(jié)果的最小乘積集
這是 DeepMind 的 AlphaTensor 方法的一種變體——通過(guò)使用組合求解器,行動(dòng)空間被縮小了一百萬(wàn)倍。以下為研究團(tuán)隊(duì)提供的 2*2 矩陣的簡(jiǎn)單例子:
總結(jié)
本文針對(duì) 這類特殊矩陣乘法提出了創(chuàng)新性加速方法,通過(guò)引入 AI 方法設(shè)計(jì)出新型算法「RXTX」,成功實(shí)現(xiàn)了總運(yùn)算量 5% 的優(yōu)化。這一突破不僅從理論上拓展了人類對(duì)計(jì)算復(fù)雜度邊界的認(rèn)識(shí),也為相關(guān)領(lǐng)域的算法優(yōu)化提供了新的研究范式。
鑒于 矩陣在多個(gè)學(xué)科領(lǐng)域的基礎(chǔ)性作用,本研究成果有望為實(shí)際應(yīng)用場(chǎng)景帶來(lái)顯著的能耗優(yōu)化。然而,新算法的工程化應(yīng)用仍面臨硬件適配和內(nèi)存管理等關(guān)鍵挑戰(zhàn),其產(chǎn)業(yè)化落地尚需學(xué)術(shù)界與工業(yè)界的持續(xù)協(xié)同攻關(guān)。要實(shí)現(xiàn)新算法的全方面落地,仍然面臨諸多挑戰(zhàn),可謂任重道遠(yuǎn)。