挑戰Scaling Law,Meta發布移動端350M小模型MobileLLM,性能比肩7B LLaMA-v2
從5月和6月幾家科技巨頭的發布會中,我們已經能隱隱感受到AI的一個重要發展趨勢:從云數據中心走向個人用戶,從大型服務器走向筆記本和移動設備。
遵循Scaling Law已經不再是唯一的路徑,模型「以小搏大」的故事不斷上演。
先有微軟更新Phi系列小模型,一個樹莓派即可運行RAG;后有谷歌用27B參數Gemma 2力壓70B的Llama 3。
硬件方面,我們看到了AI功能逐漸與電子產品進行深度集成。
比如微軟臭名昭著的Recall功能,正是他們AI+PC戰略的重要組成部分;蘋果也在Apple Intelligence的大旗下推出用于3B小模型,力求與iOS無縫銜接。
如今LLM的參數量動輒上百億,蘋果3B的參數量已經顯得十分迷你,但對手機這種移動設備來說依舊有很高門檻。
不僅用2-bit和4-bit混合精度壓縮模型(平均每個權重3.5-bit),而且要有至少8G內存和M1芯片才能運行。
Meta最近發表的一篇論文就表明,參數量可以進一步收縮,最新提出的MobileLLM模型參數量小于1B,但性能依舊可觀。
論文地址:https://arxiv.org/abs/2402.14905
LeCun也親自發推為這項研究背書,稱贊了其中一系列精簡參數量的操作。
這篇論文已被ICML 2024接收,模型的訓練代碼也已經在GitHub上開源。
GitHub地址:https://github.com/facebookresearch/MobileLLM
簡介
我們首先做個假設,如果把GPT-4(大約有1萬億參數)以50tokens/s的推理速度部署在生活中,你需要什么樣的硬件?
答案是1億個H100 GPU。別說是移動設備了,家里都放不下。
那如果降低標準,用LLaMA-v2 7B這樣的模型,再加上8-bit量化呢?
簡單計算一下,光存儲模型參數就需要約7GB,但不是存儲空間,而是珍貴的運存空間(DRAM)。
而且DRAM也不能被AI模型全占了,考慮到操作系統和其他應用的運行,LLM的運存占比不能超過10%。
按照圖2的統計,各個品牌最近發布的移動設備一般會配備6~12GB的DRAM。這就意味著,如果要在手機上順利部署,模型的參數量最好能降低到<1B。
不僅是運存,耗電也是一大問題。7B模型的能耗大概是0.7J/token,一個滿電的iPhone大概有50kJ可供揮霍。計算下來,如果生成速度是10tokens/s,手機充滿一次電只夠你和模型對話2小時。
基于上述考慮,用<1B的模型部署在移動端是更理想的選擇,因此MobileLLM的參數量定位在125M/350M,比蘋果的3B模型還少了一個數量級,可謂「迷你中的迷你」。
但是別被Scaling Law局限,參數小不意味著能力弱,模型架構的重要性應該重新進入我們的視線。
MobileLLM不僅在同等大小的模型中達到了SOTA性能,而且提出,架構的深度比寬度更重要。一個「深而窄」的「瘦長」小模型同樣可以學習到抽象概念。
架構與方法
在只有125M/350M參數的情況下,如何在有限范圍內實現架構設計的最優化就成為了重要的問題。
對于<1B的LLM,作者探索出了4種行之有效的架構設計技巧。
1)使用SwiGLU前饋網絡
2)讓網絡整體形狀變得「狹長」,即深而窄
3)重新使用編碼共享(embedding sharing)方法
4)使用組查詢注意力機制(grouped query attention)
在此基礎上,作者還提出了一種塊間層共享(block-wise layer-sharing)方法,能夠在不引入額外內存開銷的情況下進一步提高模型準確率,但代價是增加解碼過程的推理延遲。
這種添加了層共享機制的模型被標記為MobileLLM-LS。
反駁Scaling Law:小模型的架構設計很重要
2020年提出Scaling Law的論文認為,訓練數據量、參數量以及訓練迭代次數才是決定性能的關鍵因素,而模型架構的影響幾乎可以忽視。
然而這篇論文的作者通過對比實驗提出,這個定律對小模型并不適用。
當模型參數固定在125M或者350M時,30~42層的「狹長」模型明顯比12層左右的「矮胖」模型有更優越的性能(圖4),在常識推理、問答、閱讀理解等8個基準測試上都有類似的趨勢。
這其實是非常有趣的發現,因為以往為125M量級的小模型設計架構時,一般都不會疊加超過12層。
為什么要重拾「編碼共享」
「編碼共享」(embedding sharing)方法最開始由OPT這樣的小模型提出,因為小模型中編碼層的參數占到了相當大的比例。
比如,125M模型中要使用上下文長度32k、維度512的編碼,輸入和輸出編碼層就包含了16M的參數,占比達到20%。
相較之下,大模型的編碼層參數量顯得微不足道。比如LLaMA-7B中,這個比例就下降到了3.7%,LLaMA-70B甚至只有0.7%。因此,共享編碼對于LLM來說可有可無。
編碼共享在大模型時代的過氣,不代表這種技術不再適用于小模型,它可以讓模型架構更緊湊、更有效率。
如表1所示,進行編碼共享后,模型在總參數量降低16M的情況下依舊總體維持了原有性能,甚至在某些基準上有提升。
層共享機制
之前提到,論文的實驗結果發現,讓小模型變得「瘦長」有利于性能提升。于是作者想到:如果引入層共享機制,不就相當于保持參數總量不變的同時,增加了模型深度。
實驗證明,這種方法的確可以提升性能,而且論文還對比了不同的層共享方法(圖6),最終權衡設備內存、性能和推理延遲,選擇了即時塊間層共享(immediate block-wise sharing,圖6b)。
評估實驗
作者構建了125M和350M參數的MobileLLM/MobileLLM-LS模型,并在1T的數據集上進行訓練。
預訓練后的模型在多個數據集上進行零樣本測試,包括ARC-easy、ARCchallenge、HellaSwag、 WinoGrande、TQA、RACE等常用基準。
表3展示的是零樣本常識推理方面的測評結果,MobileLLM系列基本實現了全面SOTA,不僅能超越之前發布的OPT、BLOOM等經典模型,也優于最近發布的GPT-neo、Galactica、RWKV等參數更大的模型。
在問答和閱讀理解方面,MobileLLM依舊表現出色(表4)。相比其他模型,125M和325M的MobileLLM在TQA上分別有>6.4分和約10分的提升。
下游任務
除了在基準測試上跑分,論文還考慮到了應用場景部署時對模型多方面的要求,并進行了相應測評。
AlpacaEval和MT-Bench分別測試模型在單輪和多輪聊天任務中的表現,相比其他3個基線模型,MobileLLM依舊是性能最優,而且甚至能用350M的參數超過其他參數>1B模型的表現。
除了對話,在API調用的場景中,MobileLLM的EM分數可以和7B參數的LLaMA-v2相匹配。
此外,MobileLLM與量化(PTQ)的兼容性也很好。經過W8A8量化后,模型的性能只有不到0.5分的下降,并且依舊與層共享機制兼容,因此可以適應更嚴苛硬件條件下的部署。
作者簡介
本文的通訊作者Zechun Liu是Meta Reality Labs的研究科學家。她本科畢業于復旦大學,博士畢業于香港科技大學,加入Meta前曾有兩年多的時間在CMU擔任訪問學者。
Zechun的研究興趣是深度學習在現實場景中的應用,例如資源不足的限制、計算資源和精度之間的權衡等,其中重點關注網絡二值化和量化、網絡通道剪枝、架構設計、知識蒸餾等方面。