在人工智慧領域,縮放法則幾乎被視為金科玉律:模型參數越多,效能越強。然而,這種強效是有代價的。動輒數千億參數的模型,不僅讓推論成本(顯示記憶體 VRAM、延遲)暴增至天文數字,更讓環境足跡與能源消耗成為不可忽視的議題。
長久以來,研究人員始終緊盯著 Transformer 架構中最臃腫的環節:前饋神經網路。在現代大型語言模型中,前饋神經網路佔據了超過三分之二的參數量,並貢獻了百分之八十以上的運算量。
有趣的是,生物大腦其實非常節電,僅有極少量的神經元在特定時刻會處於活躍狀態。大型模型同樣具備這種潛力:在採用 ReLU 激活函數的模型中,對於任何給定的輸入,實際上只有極少數的神經元會被啟動。
這便引出了長期困擾產業界的稀疏悖論:既然理論上絕大多數的運算結果皆為零,為何在 GPU 上執行稀疏運算子,反而比密集運算子還要慢?
近日,Sakana AI 與 NVIDIA 發表的最新研究 Sparser, Faster, Lighter Transformer Language Models 正式正面對決此一難題。研究團隊不僅證實模型可達成百分之九十九以上的稀疏度且維持效能無損,更從底層 CUDA 核心層面,設計出名為 TwELL 的全新資料格式,真正將理論上的稀疏轉化為實質的加速。
一、為何稀疏運算在 GPU 上難以提速
要理解此項創新的價值,必須先釐清傳統做法的盲點。GPU 的核心優勢在於極致的平行運算,其架構本是為矩陣乘法此類規律且密集的任務所設計。傳統的稀疏資料格式在處理稀疏矩陣時,必須額外紀錄非零元素的索引與位置。
正如研究指出:由於硬體與軟體堆疊已針對密集運算模式進行大量最佳化,異質工作負載以及實體化、管理稀疏索引所帶來的開銷,一直是阻礙通用計算節省效能的關鍵挑戰。
在門控前饋網路中,稀疏模式是由門控層的激活值所決定。若想使用傳統稀疏運算子,必須先執行一遍門控運算,統計出哪些是非零元素,重新排列索引後,再進行後續的矩陣乘法。這個重新排列的過程往往比節省下來的計算時間更長。換言之,GPU 在等待指令與搬移零散資料上浪費了大量資源。
二、TwELL 格式:專為 GPU 磁貼架構打造的拼圖設計
為突破此悖論,作者提出了 TwELL(Tile-wise ELLPACK)格式。
這是一項極為精妙的工程設計。既然 GPU 偏好以磁貼為單位處理任務,團隊便將稀疏化限制在單一磁貼內進行。TwELL 不再試圖對整個矩陣進行全域壓縮與索引重排,而是改在每個磁貼內部獨立收集非零元素。
此設計的核心優勢在於運算融合:
• 允許在執行門控矩陣乘法的同一個 CUDA 核心尾段,直接將資料具現化為 TwELL 格式。
• 無需全域同步,也免去中間記憶體的讀寫操作。
• 後續的向上投影與向下投影運算子可直接融合進同一管線,讀取這些局部對齊的稀疏資料。
白話而言,TwELL 宛如在生產線上即完成零件分類裝箱,而非等產品全數產出後,再停線進行二次分揀。這種磁貼級的局部處理機制,完美契合了現代 NVIDIA GPU 的硬體特性。
三、推論與訓練的雙重進化
除了推論端的 TwELL 融合運算子,團隊在訓練端亦祭出關鍵解法。大型模型訓練時,顯示記憶體往往是最大瓶頸。若中間層的激活值未加壓縮,將佔據龐大空間。然而稀疏訓練面臨一大隱患:非均勻性。某些 Token 可能激活五百個神經元,有些卻僅有幾個。若以最大值為預留標準,則無法節省記憶體;若以平均值預留,激活量高的 Token 又會導致溢位。
為此,作者設計了一種混合格式來破解難題:
• 絕大多數符合稀疏規律的行列,儲存在緊湊的稀疏矩陣中。
• 極少數激活異常高的長尾行列,則導入至密集的備援緩衝區。
• 此方案既能運用 Tensor Core 的密集運算能力處理高負載任務,又能透過客製化稀疏核心處理輕量任務,實現記憶體佔用與運算速度的雙贏。
四、實驗成果:百分之九十九稀疏度的突破
團隊藉由 L1 正則化誘導模型產生稀疏特性。實驗數據令人振奮:
從數據中可歸納出幾項核心結論:
一、規模效應顯著:模型規模越大,稀疏化帶來的效益越明確。以二十億參數模型為例,推論速度提升百分之二十點五,訓練速度提升百分之二十一點九。
二、記憶體極致優化:十億模型的訓練峰值記憶體從 44.5 GB 降至 33.1 GB,降幅達百分之二十五點五。
三、效能高度保真:在引入輕微 L1 正則化後,模型的平均任務準確率幾乎未見下滑(例如十億模型反而從百分之四十四點六微幅上升至百分之四十四點七)。
研究強調:我們針對大型語言模型稀疏性進行了量化研究,證實簡易的 L1 正則化能誘導出超過百分之九十九的稀疏度,且對下游任務效能的影響微乎其微。
這意味著,我們不僅能提升模型運行效率,更得以在規格較低或顯示記憶體較小的顯示卡上,訓練過往無法負載的大型模型。
五、稀疏性背後的邏輯:模型學會抓重點
研究進一步揭示,大型語言模型究竟在何種情境下轉為稀疏。團隊發現,稀疏度與輸入的資訊熵高度相關。面對高度可預測的 Token(例如網址中的 doi、gov 或常見縮寫),模型分配的激活神經元極少;反之,對於承載關鍵上下文資訊的詞彙(例如特定地名 Vermont、複雜化學名詞 formaldehyde),模型的激活程度顯著攀升。
此外,序列位置亦具影響力。一段序列開頭的前幾個 Token 通常需要最多神經元來建立上下文脈絡,隨著序列增長,稀疏度呈指數型上升。這證明了稀疏大型語言模型確實學會了動態分配運算資源,將算力精準投放於關鍵節點。
六、技術侷限與未來展望
當然,此技術並非毫無門檻。目前該核心程式高度針對 NVIDIA Hopper 架構(例如 H100)進行最佳化,特別是運用了張量記憶體加速器等嶄新特性。這意味著,在舊世代硬體或非 NVIDIA 平台上,效能收益可能會有所折損。此外,正則化係數的微調需極為謹慎,係數過高恐引發神經元壞死的問題。
儘管如此,瑕不掩瑜。Sakana AI 與 NVIDIA 的這項成果為產業指明了一條明路:大型模型的未來,未必是無止境地堆砌硬體運算力,而是邁向更精細化的動態運算分配。
經由開源這些程式碼與核心,作者期盼稀疏化能成為現代基礎模型設計的全新維度。當我們能以更低的能耗、更少的記憶體與更快的速度,跑出同等級的成果時,縮放法則才算真正邁入下一個演化階段。