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

一直以來，矩陣乘法（MatMul）穩(wěn)居神經(jīng)網(wǎng)絡(luò)操作的主導(dǎo)地位，其中很大原因歸結(jié)為 GPU 專門針對 MatMul 操作進(jìn)行了優(yōu)化。這種優(yōu)化使得 AlexNet 在 ILSVRC2012 挑戰(zhàn)賽中一舉勝出，成為深度學(xué)習(xí)崛起的歷史性標(biāo)志。

在這當(dāng)中，有個值得注意的點(diǎn)是，AlexNet 利用 GPU 來提高訓(xùn)練速度，超越了 CPU 的能力，至此，GPU 的加入使得深度學(xué)習(xí)仿佛贏得了「硬件彩票」。

盡管 MatMul 在深度學(xué)習(xí)中很流行，但不得不承認(rèn)的是它占據(jù)了計(jì)算開銷的主要部分，主要表現(xiàn)為 MatMul 在訓(xùn)練和推理階段消耗大部分執(zhí)行時間和內(nèi)存訪問。

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

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

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

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

在這項(xiàng)工作中，來自加州大學(xué)圣克魯茲分校等機(jī)構(gòu)的研究者證明了 MatMul 操作可以完全從 LLM 中消除，同時在十億參數(shù)尺度下保持強(qiáng)大的性能。

• 論文地址：https://arxiv.org/pdf/2406.02528
• 項(xiàng)目地址：https://github.com/ridgerchu/matmulfreellm
• 論文標(biāo)題：Scalable MatMul-free Language Modeling

實(shí)驗(yàn)表明，該研究提出的 MatMul-free 模型達(dá)到了與最先進(jìn)的 Transformer 相當(dāng)?shù)男阅?，后者在推理期間需要更多的內(nèi)存，規(guī)模至少為 2.7B 參數(shù)。

此外，論文還研究了擴(kuò)展定律，發(fā)現(xiàn)隨著模型規(guī)模的增加，MatMul-free 模型與全精度 Transformer 之間的性能差距逐漸縮小。

研究者還提供了一種高效的 GPU 模型實(shí)現(xiàn)方式，在訓(xùn)練期間相比未優(yōu)化的基線模型減少了多達(dá) 61% 的內(nèi)存使用。通過在推理時利用優(yōu)化的內(nèi)核，模型內(nèi)存消耗可以比未優(yōu)化的模型減少超過 10 倍。 

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

網(wǎng)友看后不禁感嘆道：看來有大事要發(fā)生了。

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

方法介紹

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

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

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

具有三值權(quán)重的 MatMul-free 密集層

在標(biāo)準(zhǔn)密集層中，輸入和權(quán)重矩陣之間的 MatMul 可以表示為：

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

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

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

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

MatMul-free 語言模型架構(gòu)

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

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

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

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

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

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

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

實(shí)驗(yàn)

該研究的重點(diǎn)是在中等規(guī)模的語言建模任務(wù)上測試 MatMul-free 的語言模型。研究人員將兩種變體的 MatMul-free 語言模型與復(fù)現(xiàn)的高級 Transformer 架構(gòu)（基于 Llama-2 的 Transformer++）在三個模型大小上進(jìn)行比較：370M、13 億和 27 億參數(shù)。

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

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

MatMul-free 語言模型的擴(kuò)展規(guī)律

研究團(tuán)隊(duì)評估了MatMul-free 語言模型和 Transformer++ 在 370M、13 億和 27 億參數(shù)模型上的擴(kuò)展規(guī)律，如圖 3 所示。

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

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

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

下游任務(wù)

學(xué)習(xí)率是語言模型訓(xùn)練中一個關(guān)鍵的超參數(shù)，當(dāng)模型處于三元 / 二元權(quán)重狀態(tài)時，對學(xué)習(xí)率變得更加敏感。為了確定最優(yōu)學(xué)習(xí)率，研究人員使用 370M 模型，批量大小為 50k 個 token，在 1.5e?3 到 3e?2 的范圍內(nèi)進(jìn)行了搜索。這次搜索的結(jié)果如圖 4 (c) 所示。

結(jié)果顯示，當(dāng)學(xué)習(xí)率從 1.5e?3 增加到 1e?2 時，最終訓(xùn)練損失單調(diào)遞減。只有當(dāng)學(xué)習(xí)率超過 2e?2 時，模型才表現(xiàn)出不穩(wěn)定。這一發(fā)現(xiàn)表明，以前使用三元權(quán)重的作品，如使用 1.5e?3 學(xué)習(xí)率的 BitNet，可能不是最優(yōu)的，更高的學(xué)習(xí)率有可能帶來更好的性能。

這些發(fā)現(xiàn)與 Deepseek LLM 的觀察結(jié)果一致，后者發(fā)現(xiàn)傳統(tǒng)大型語言模型（LLMs）的最佳學(xué)習(xí)率實(shí)際上比大多數(shù) LLM 訓(xùn)練設(shè)置中通常報(bào)告的值要大。

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

研究人員根據(jù)訓(xùn)練時間和內(nèi)存使用情況評估了他們提出的融合型 BitLinear 和傳統(tǒng)型 BitLinear 實(shí)現(xiàn)，如圖 4 (a-b) 所示。

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

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

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

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

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

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

值得注意的是，對于大于 2.7B 的模型大小，結(jié)果是使用隨機(jī)初始化的權(quán)重進(jìn)行模擬的。對于最大的 13B 參數(shù)模型，MatMul-free 語言模型僅使用 4.19 GB 的 GPU 內(nèi)存，并且有 695.48 毫秒的延遲，而 Transformer++ 需要 48.50 GB 的內(nèi)存，并表現(xiàn)出 3183.10 毫秒的延遲。

這些結(jié)果突出了MatMul-free語言模型所實(shí)現(xiàn)的效率增益，使其成為大規(guī)模語言建模任務(wù)的有希望的方法，特別是在推理期間。

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

表 2 顯示了單塊（N = 1）單核實(shí)現(xiàn)的資源利用、功耗和性能。

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