讓語言模型「輕裝上陣」。

一直以來，矩陣乘法（MatMul）穩居神經網絡操作的主導地位，其中很大原因歸結為 GPU 專門針對 MatMul 操作進行了優化。這種優化使得 AlexNet 在 ILSVRC2012 挑戰賽中一舉勝出，成為深度學習崛起的歷史性標志。

在這當中，有個值得注意的點是，AlexNet 利用 GPU 來提高訓練速度，超越了 CPU 的能力，至此，GPU 的加入使得深度學習仿佛贏得了「硬件彩票」。

盡管 MatMul 在深度學習中很流行，但不得不承認的是它占據了計算開銷的主要部分，主要表現為 MatMul 在訓練和推理階段消耗大部分執行時間和內存訪問。

目前為止，研究者已經開始借助其他更簡單的操作替代 MatMul，主要有兩種。

• 第一種策略是使用初等運算代替 MatMul，例如，在卷積神經網絡 (CNN) 中，用有符號加法代替乘法；
• 第二種方法是使用二值或三值化量化，將 MatMul 值在累加之前要么翻轉要么清零。比如脈沖神經網絡 (SNN) 使用二值激活，而二值化網絡 BNN 使用量化權重。

在語言建模方面，BitNet 等技術的出現表明量化的可擴展性， 但這種方式仍然保留了昂貴的矩陣 - 矩陣相乘（MMM）的自注意力機制。研究者嘗試過多種努力，但 MatMul 操作在 GPU 上仍然是資源密集型的。

既然 MatMul 占據了 LLM 整體計算成本，且隨著 LLM 向更大的嵌入維度和上下文長度擴展時，這種成本只會增加。這引發了一個問題：是否有可能完全從 LLM 中消除 MatMul 操作？ 

在這項工作中，來自加州大學圣克魯茲分校等機構的研究者證明了 MatMul 操作可以完全從 LLM 中消除，同時在十億參數尺度下保持強大的性能。

• 論文地址：https://arxiv.org/pdf/2406.02528
• 項目地址：https://github.com/ridgerchu/matmulfreellm
• 論文標題：Scalable MatMul-free Language Modeling

實驗表明，該研究提出的 MatMul-free 模型達到了與最先進的 Transformer 相當的性能，后者在推理期間需要更多的內存，規模至少為 2.7B 參數。

此外，論文還研究了擴展定律，發現隨著模型規模的增加，MatMul-free 模型與全精度 Transformer 之間的性能差距逐漸縮小。

研究者還提供了一種高效的 GPU 模型實現方式，在訓練期間相比未優化的基線模型減少了多達 61% 的內存使用。通過在推理時利用優化的內核，模型內存消耗可以比未優化的模型減少超過 10 倍。 

最后，研究者在 FPGA 上構建了一個自定義硬件解決方案，他們以 13W 的功耗處理了十億參數規模的模型，超出了人類可讀的吞吐量，使 LLM 更接近大腦般的效率。

網友看后不禁感嘆道：看來有大事要發生了。

不過，受到計算資源的限制，研究者還沒有在非常大的模型（如參數大于100B的模型）上測試MatMul-free模型的有效性，因此其實戰效果還有待觀察。

方法介紹

該研究構建了首個可擴展的MatMul-free語言模型 (Matmul-free LM)，通過在密集層中使用加法操作以及在自注意力類函數中使用元素級 Hadamard 乘積完成。 

具體來說，三值權重消除了密集層中的 MatMul，類似于 BNN。為了從自注意力中移除 MatMul，該研究優化了門控循環單元 (GRU) ，使其僅依賴于元素級乘積，并表明該模型可與最先進的 Transformer 相媲美，同時消除了所有 MatMul 運算。

論文詳細介紹 MatMul-free 語言模型（LM）的各個組成部分。

具有三值權重的 MatMul-free 密集層

在標準密集層中，輸入和權重矩陣之間的 MatMul 可以表示為：

為了避免使用基于 MatMul 的密集層，該研究采用 BitNet 來替換包含 MatMul 的密集層，即使用 BitLinear 模塊，這些模塊使用三值權重將 MatMul 操作轉換為純加法操作。當使用三值權重時，權重矩陣 W 中的元素被限制在集合 {-1, 0, +1} 中。帶有三值權重的 MatMul 可以表達為： 

由于三值化權重只能從 {?1, 0, +1} 中取值，因而 MatMul 中的乘法運算可以用簡單的加法或減法運算代替：

因此，三值化 MatMul 可以寫成如下：

算法 1 為量化融合 RMSNorm 和 BitLinear 算法的流程圖：

MatMul-free 語言模型架構

研究人員采用了 Metaformer 的觀點，該觀點認為 Transformer 由兩部分組成：token mixer（用于混合時間信息，即自注意力機制）和 channel mixer（用于混合嵌入 / 空間信息，即前饋網絡，門控線性單元 GLU ）。該架構的高級概覽如圖 2 所示。

自注意力機制是現代語言模型中最常用的 token mixer，它依賴于三個矩陣 Q、K 和 V 之間的矩陣乘法。為了將這些操作轉換為加法，研究人員至少對兩個矩陣進行二值化或三值化處理。

假設所有密集層的權重都是三值的，他們將 Q 和 K 量化，得到一個三值的注意力圖，從而消除自注意力中的乘法操作。但是，以這種方式訓練的模型無法成功收斂。一個可能的解釋是，激活值包含對性能至關重要但難以有效量化的異常值。

為了解決這一挑戰，研究人員探索了不依賴于矩陣乘法的替代方法來混合 token。

通過采用結合了逐元素操作和累積的三值 RNN，可以構建一個MatMul-free 的 token mixer。在各種 RNN 架構中，GRU 因其簡單高效而著稱，它在比長短期記憶網絡（LSTM）使用更少的門控單元和結構更簡單的情況下，實現了類似的性能。因此，研究人員選擇 GRU 作為構建 MatMul-free token mixer 的基礎。

為了實現 MatMul-free 的通道混合，研究人員采用了門控線性單元（GLU），它在許多現代 LLM 中得到了廣泛應用，包括 Llama 、Mistral 和 RWKV。一個適應了 BitLinear 的 GLU 版本可以表達如下：

這里的通道混合器僅由密集層組成，這些層已被三值化累積操作所替代。通過在 BitLinear 模塊中使用三值權重，研究人員可以消除對昂貴 MatMul 的需求，這樣不僅使通道混合器在計算上更加高效，同時還保持了其在跨通道混合信息方面的有效性。

實驗

該研究的重點是在中等規模的語言建模任務上測試 MatMul-free 的語言模型。研究人員將兩種變體的 MatMul-free 語言模型與復現的高級 Transformer 架構（基于 Llama-2 的 Transformer++）在三個模型大小上進行比較：370M、13 億和 27 億參數。

為了公平比較，所有模型都在 SlimPajama 數據集上進行了預訓練，其中 370M 模型訓練了 150 億個 token，而 13 億和 27 億模型各訓練了 1000 億個 token。

模型訓練使用了 8 個 NVIDIA H100 GPU。370M 模型的訓練時間大約為 5 小時，13 億模型為 84 小時，27 億模型為 173 小時。

MatMul-free 語言模型的擴展規律

研究團隊評估了MatMul-free 語言模型和 Transformer++ 在 370M、13 億和 27 億參數模型上的擴展規律，如圖 3 所示。

為便于比較，MatMul-free LM 和 Transformer++ 中的每個操作都被同等對待。但請注意，Transformer++ 中的所有權重和激活都是 BF16 格式，而MatMul-free 語言模型中的 BitLinear 層使用三元參數，激活為 BF16。因此，MatMul-free 語言模型的平均運算成本要低于 Transformer++。

有意思的是，與 Transformer++ 相比，MatMul-free 語言模型的擴展投影顯示出更陡峭的下降趨勢，這表明MatMul-free語言模型在利用額外計算資源提高性能方面更為高效。

因此，MatMul-free 語言模型的擴展曲線預計將在大約 10^23 次浮點運算（FLOPs）處與 Transformer++ 的擴展曲線相交。這個計算規模大致相當于訓練 Llama-3 8B（使用 1.5 萬億個 token 訓練）和 Llama-2 70B（使用 2 萬億個 token 訓練）所需的 FLOPs，這表明MatMul-free 語言模型不僅在效率上勝出，而且擴展時在損失方面也可能表現更好。

下游任務

學習率是語言模型訓練中一個關鍵的超參數，當模型處于三元 / 二元權重狀態時，對學習率變得更加敏感。為了確定最優學習率，研究人員使用 370M 模型，批量大小為 50k 個 token，在 1.5e?3 到 3e?2 的范圍內進行了搜索。這次搜索的結果如圖 4 (c) 所示。

結果顯示，當學習率從 1.5e?3 增加到 1e?2 時，最終訓練損失單調遞減。只有當學習率超過 2e?2 時，模型才表現出不穩定。這一發現表明，以前使用三元權重的作品，如使用 1.5e?3 學習率的 BitNet，可能不是最優的，更高的學習率有可能帶來更好的性能。

這些發現與 Deepseek LLM 的觀察結果一致，后者發現傳統大型語言模型（LLMs）的最佳學習率實際上比大多數 LLM 訓練設置中通常報告的值要大。

有趣的是，研究人員還觀察到，與使用較小學習率訓練的模型相比，訓練初期使用較大學習率訓練的模型，在訓練后期的階段訓練損失下降得更快。

研究人員根據訓練時間和內存使用情況評估了他們提出的融合型 BitLinear 和傳統型 BitLinear 實現，如圖 4 (a-b) 所示。

實驗表明，他們的融合操作器在更大的批量大小下，能夠帶來更快的訓練速度，并減少內存消耗。當批量大小為 2 的 8 次方時，1.3B 參數模型的訓練速度從每次迭代 1.52 秒提高到 1.21 秒，比 Vanilla 實現快了 25.6%。

此外，內存消耗從 82GB 減少到 32GB，內存使用減少了 61.0%。隨著批量大小的增加，融合實現的性能顯著提高，允許同時處理更多的樣本，并減少總迭代次數。

圖 4 (d) 展示了不同模型大小下，所提出的MatMul-free語言模型與 Transformer++ 在 GPU 推理內存消耗和延遲方面的比較。

在MatMul-free語言模型中，研究人員采用 BitBLAS 進行加速，以進一步提高效率。評估是在批量大小為 1，序列長度為 2048 的情況下進行的。

MatMul-free語言模型在所有模型大小上都顯示出比 Transformer++ 更低的內存使用和延遲。

對于單層，MatMul-free 語言模型只需要 0.12 GB 的 GPU 內存，并且實現了 3.79 毫秒的延遲，而 Transformer++ 消耗了 0.21 GB 的內存，并且有 13.87 毫秒的延遲。隨著模型大小的增加，MatMul-free 語言模型的內存和延遲優勢變得更加明顯。

值得注意的是，對于大于 2.7B 的模型大小，結果是使用隨機初始化的權重進行模擬的。對于最大的 13B 參數模型，MatMul-free 語言模型僅使用 4.19 GB 的 GPU 內存，并且有 695.48 毫秒的延遲，而 Transformer++ 需要 48.50 GB 的內存，并表現出 3183.10 毫秒的延遲。

這些結果突出了MatMul-free語言模型所實現的效率增益，使其成為大規模語言建模任務的有希望的方法，特別是在推理期間。

為了測試MatMul-free語言模型功耗和有效性，研究人員使用 SystemVerilog 創建了一個 FPGA 加速器。概覽如圖 5 所示。

表 2 顯示了單塊（N = 1）單核實現的資源利用、功耗和性能。

通過使用完整的 512 位 DDR4 接口，并對占據核心處理時間 99% 的 TMATMUL 功能單元進行并行化，預計在保持相同時鐘速率的情況下，無需額外優化或流水線化，可以實現大約 64 倍的進一步加速，如表 3 所示。