「在移動設(shè)備上運(yùn)行 LLM？可能需要 Meta 的一些技巧?！箘倓?，圖靈獎得主 Yann LeCun 在個人社交平臺表示。

他所宣傳的這項研究來自 Meta 最新論文《 MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases 》，在眾多作者中也有我們熟悉的來自 Meta FAIR 田淵棟。

田淵棟表示：「我們的 MobileLLM 預(yù)訓(xùn)練模型（125M/350M），性能達(dá)到 SoTA，特別是在聊天 / API 調(diào)用方面表現(xiàn)出色。此外，本工作中的一個有趣研究是跨 Transformer 層的權(quán)重共享，這樣不僅節(jié)省了參數(shù)，還減少了推理過程中的延遲?！?/span>

論文地址：https://arxiv.org/pdf/2402.14905.pdf

現(xiàn)階段大語言模型（LLM）已經(jīng)滲透到人類生活的各個方面，尤其是以 ChatGPT 等為代表的模型，這類研究主要在云環(huán)境中運(yùn)行。

然而領(lǐng)先的模型如 ChatGPT4 的參數(shù)量已經(jīng)超過了 1 萬億。我們設(shè)想這樣一個場景，這個場景廣泛依賴 LLM，不僅用于前端的會話界面，也用于后端操作，如推薦系統(tǒng)，覆蓋人類約 5% 的時間。在這一假設(shè)場景中，假如以 GPT-4 每秒處理 50 個 token 的速率來計算，則需要部署大約一億個 H100 GPU，每個 GPU 的計算能力為 60 TFLOPs/s。這種計算規(guī)模，還不包括通信和數(shù)據(jù)傳輸?shù)拈_銷，就已經(jīng)與 160 個 Meta 規(guī)模的公司相當(dāng)。隨之而來的能源消耗和二氧化碳排放將帶來巨大的環(huán)境挑戰(zhàn)。

因此，最好的解決方案是縮小 LLM 的規(guī)模。

此外，在當(dāng)前的移動技術(shù)領(lǐng)域，由于主內(nèi)存（DRAM）容量的限制，將像 LLaMAv2 7B 這樣的 LLM 與 8 位權(quán)重整合起來代價過高。移動設(shè)備中普遍的內(nèi)存層結(jié)構(gòu)如圖 2 所示。隨著 DRAM 容量從 iPhone 15 的 6GB 到 Google Pixel 8 Pro 的 12GB 不等，一個移動應(yīng)用不應(yīng)超過 DRAM 的 10%，因?yàn)?DRAM 需要與操作系統(tǒng)和其他應(yīng)用程序共享。這一要求促進(jìn)了部署小于十億參數(shù) LLM 更進(jìn)一步的研究。

基于上述考量，來自 Meta 的研究者專注于設(shè)計參數(shù)少于十億的高質(zhì)量 LLM，這是在移動端部署 LLM 比較好的解決方案。

與強(qiáng)調(diào)數(shù)據(jù)和參數(shù)數(shù)量在決定模型質(zhì)量方面的關(guān)鍵作用的普遍觀點(diǎn)相反，Meta 強(qiáng)調(diào)了模型架構(gòu)對少于十億（sub-billion）規(guī)模 LLM 的重要性。

基于深而窄的架構(gòu)，加上嵌入共享和分組查詢注意力機(jī)制，Meta 建立了一個強(qiáng)大的基線網(wǎng)絡(luò)，稱為 MobileLLM，與之前的 125M/350M 最先進(jìn)模型相比，其準(zhǔn)確率顯著提高了 2.7%/4.3% 。這也說明了與縮放定律（scaling law）相反，該研究證明對于小型 LLM 來說深度比寬度更重要，一個深而窄的模型結(jié)構(gòu)在捕獲抽象概念方面更為出色。

此外，Meta 還提出了一種及時逐塊權(quán)重共享（ immediate block-wise weight sharing）方法，該方法不會增加模型大小，所得模型表示為 MobileLLM-LS，其準(zhǔn)確率比 MobileLLM 125M/350M 進(jìn)一步提高了 0.7%/0.8%。此外，在下游任務(wù)中，例如 Chat 和 API 調(diào)用，MobileLLM 模型家族顯著優(yōu)于同等規(guī)模的模型。在 API 調(diào)用任務(wù)中，與規(guī)模較大的 LLaMA-v2 7B 相比，MobileLLM 甚至實(shí)現(xiàn)了相媲美的分?jǐn)?shù)。

看到這項研究后，網(wǎng)友紛紛表示「我們應(yīng)該向 Meta 致敬，很高興看到這個領(lǐng)域的活躍玩家。該機(jī)構(gòu)通過使用低于 10 億參數(shù)的模型，并且 350M 8 位模型的能源消耗僅為 0.035 J/token  ，要是部署在 iPhone 上的話，可以支持用戶一整天的會話使用?！?/span>

改進(jìn)十億以下參數(shù)規(guī)模的 LLM 設(shè)計

研究者介紹了從十億以下參數(shù)規(guī)模的基線模型到新的 SOTA 模型的演進(jìn)之路（如下圖 3 所示）。他們分別研究了 125M 和 350M 參數(shù)規(guī)模的模型，并在這兩個規(guī)模下展示了一致的改進(jìn)。對于模型尺寸成為主要制約因素的設(shè)備用例而言，如何有效地分配有限的權(quán)重參數(shù)變得比以往更加重要。

研究者首先通過測試四種有益于十億以下規(guī)模 LLM 的模型設(shè)計方法，提出了一個名為MobileLLM 的強(qiáng)大基線模型。這四種模型設(shè)計方法包括 1）采用 SwiGLU FFN，2）強(qiáng)制使用深和薄的架構(gòu)，3）重新審視嵌入共享方法，4）利用分組查詢注意力。

接下來，研究者開發(fā)了一種直接的逐塊層共享方法，基于該方法可以進(jìn)一步提高準(zhǔn)確度，而不產(chǎn)生任何額外的內(nèi)存開銷，并在內(nèi)存有限的 LM 解碼過程中產(chǎn)生很小的延遲開銷。他們將具有層共享的模型表示為 MobileLLM-LS。

訓(xùn)練設(shè)置

研究者在 32 個 A100 GPU 上進(jìn)行實(shí)驗(yàn)，其中每個 GPU 的批大小為 32。他們在 0.25T 的 tokens 上執(zhí)行了 120k 次迭代的探索性實(shí)驗(yàn)。下文中表 3 和表 4 報告了在 1T 的 tokens 上執(zhí)行 480k 次迭代訓(xùn)練的 top 模型。

層共享

關(guān)于層深度與寬度影響的研究結(jié)果表明，更深的層有利于小型 transformer 模型。這促使本文研究層共享作為增加隱藏層數(shù)量而不增加存儲成本的策略。這種方法在模型大小成為主要制約因素的場景中尤其有用。

令人驚訝的是，實(shí)驗(yàn)結(jié)果表明，通過簡單地復(fù)制 transformer 塊就可以提高準(zhǔn)確度而無需任何架構(gòu)修改或擴(kuò)大模型尺寸。研究者進(jìn)一步探究三種不同的權(quán)重共享策略，具體如下圖 6 所示。

下表 2 結(jié)果表明，重復(fù)層共享策略在立即塊重復(fù)、全面重復(fù)（repeat all-over）和反向共享策略中產(chǎn)生了最佳性能。

不過，考慮到硬件內(nèi)存的層級結(jié)構(gòu)（如圖 2），用于計算的 SRAM 通常限制在了 20M 左右。該容量通常僅夠容納單個 transformer 塊。因此，將共享權(quán)重放入緩存中并立即計算兩次則無需在 SRAM 和 DRAM 之間傳輸權(quán)重，提高了自回歸推理的整體執(zhí)行速度。

研究者在模型設(shè)計中選擇了直接的分塊共享策略，并將提出的帶有層共享的模型表示為 MobileLLM-LS。

實(shí)驗(yàn)結(jié)果

該研究進(jìn)行實(shí)驗(yàn)比較了模型在零樣本（zero-shot）常識推理任務(wù)、問答和閱讀理解任務(wù)上的性能。

零樣本常識推理任務(wù)的實(shí)驗(yàn)結(jié)果如下表 3 所示：

在問答和閱讀理解任務(wù)上，該研究采用 TQA 問答基準(zhǔn)和 RACE 閱讀理解基準(zhǔn)來評估預(yù)訓(xùn)練模型，實(shí)驗(yàn)結(jié)果如下表 4 所示：

為了驗(yàn)證將模型用于設(shè)備上應(yīng)用程序的有效性，該研究評估了模型在兩個關(guān)鍵任務(wù)上的性能：聊天和 API 調(diào)用。

針對聊天任務(wù)，該研究在兩個基準(zhǔn)上進(jìn)行了評估實(shí)驗(yàn)：AlpacaEval（單輪聊天基準(zhǔn)）和 MT-Bench（多輪聊天基準(zhǔn)），實(shí)驗(yàn)結(jié)果如下表 5 所示：

在 API 調(diào)用方面，如下表 6 所示，MobileLLM-350M 表現(xiàn)出與 LLaMA-v2 7B 相當(dāng)?shù)?EM_intent 和 EM_structure，其中 EM_intent 越高，表明模型對用戶計劃調(diào)用 API 的預(yù)測就越準(zhǔn)確，而 EM_structure 反映了預(yù)測 API 函數(shù)內(nèi)內(nèi)容的熟練程度。

該研究進(jìn)一步在 MobileLLM 和 MobileLLM-LS 模型上針對每個 token 進(jìn)行最小 / 最大訓(xùn)練后量化 (PTQ) 實(shí)驗(yàn)，模型大小分別為 125M 和 350M，在 0.25T token 上進(jìn)行訓(xùn)練，實(shí)驗(yàn)結(jié)果如下圖 7 所示：

模型加載、初始化和執(zhí)行時間如下表 7 所示：

更多技術(shù)細(xì)節(jié)請參閱原論文。