「在移動設備上運行 LLM？可能需要 Meta 的一些技巧。」剛剛，圖靈獎得主 Yann LeCun 在個人社交平臺表示。

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

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

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

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

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

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

此外，在當前的移動技術領域，由于主內存（DRAM）容量的限制，將像 LLaMAv2 7B 這樣的 LLM 與 8 位權重整合起來代價過高。移動設備中普遍的內存層結構如圖 2 所示。隨著 DRAM 容量從 iPhone 15 的 6GB 到 Google Pixel 8 Pro 的 12GB 不等，一個移動應用不應超過 DRAM 的 10%，因為 DRAM 需要與操作系統和其他應用程序共享。這一要求促進了部署小于十億參數 LLM 更進一步的研究。

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

與強調數據和參數數量在決定模型質量方面的關鍵作用的普遍觀點相反，Meta 強調了模型架構對少于十億（sub-billion）規模 LLM 的重要性。

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

此外，Meta 還提出了一種及時逐塊權重共享（ immediate block-wise weight sharing）方法，該方法不會增加模型大小，所得模型表示為 MobileLLM-LS，其準確率比 MobileLLM 125M/350M 進一步提高了 0.7%/0.8%。此外，在下游任務中，例如 Chat 和 API 調用，MobileLLM 模型家族顯著優于同等規模的模型。在 API 調用任務中，與規模較大的 LLaMA-v2 7B 相比，MobileLLM 甚至實現了相媲美的分數。

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

改進十億以下參數規模的 LLM 設計

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

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

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

訓練設置

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

層共享

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

令人驚訝的是，實驗結果表明，通過簡單地復制 transformer 塊就可以提高準確度而無需任何架構修改或擴大模型尺寸。研究者進一步探究三種不同的權重共享策略，具體如下圖 6 所示。

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

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

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

實驗結果

該研究進行實驗比較了模型在零樣本（zero-shot）常識推理任務、問答和閱讀理解任務上的性能。

零樣本常識推理任務的實驗結果如下表 3 所示：

在問答和閱讀理解任務上，該研究采用 TQA 問答基準和 RACE 閱讀理解基準來評估預訓練模型，實驗結果如下表 4 所示：

為了驗證將模型用于設備上應用程序的有效性，該研究評估了模型在兩個關鍵任務上的性能：聊天和 API 調用。

針對聊天任務，該研究在兩個基準上進行了評估實驗：AlpacaEval（單輪聊天基準）和 MT-Bench（多輪聊天基準），實驗結果如下表 5 所示：

在 API 調用方面，如下表 6 所示，MobileLLM-350M 表現出與 LLaMA-v2 7B 相當的 EM_intent 和 EM_structure，其中 EM_intent 越高，表明模型對用戶計劃調用 API 的預測就越準確，而 EM_structure 反映了預測 API 函數內內容的熟練程度。

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

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

更多技術細節請參閱原論文。