把MLA和稀疏激活帶到端側!港科大廣州和倫敦大學學院團隊聯合發佈軟硬協同設計的邊緣語言模型PLM

本文由 PLM 團隊撰寫,PLM 團隊是由香港科技大學(廣州)的校長倪明選教授,倫敦大學學院(UCL)AI 中心汪軍教授,香港科技大學(廣州)信息樞紐院長陳雷教授聯合指導。第一作者丹恩程是香港科技大學(廣州)的研究助理,研究方向為端側大模型和高效模型推理;參與成員包括中科院自動化所的孫羅洋博士,曾勇程博士,薑紀文碩士,UCL 吳昕鍵,港科大廣州的博士生肖慶發和趙文欣,港科大的博士後王嘉川以及香港理工的助理教授(研究)李昊洋。通訊作者為丹恩程博士,陳雷教授和汪軍教授。
在大模型「卷參數」的浪潮中,一個亟待解決的課題始終存在:如何讓百億級能力的 AI 跑進手機、嵌入設備、實現萬物互聯?內存牆、算力牆、I/O 牆這三座邊緣設備的「大山」令許多大模型铩羽而歸。
近日,香港科技大學(廣州)倪明選校長和陳雷教授聯合 UCL 汪軍教授團隊以及中科院自動化所團隊聯合提出 PLM(Peripheral Language Model),通過算法 – 系統協同設計,選擇了適合邊緣設備的模型架構。
PLM 是首個結合 MLA 注意力機制與 ReLU2 激活 FFN 的架構。評估結果表明,PLM 的表現優於現有在公開數據上訓練的小型語言模型,性能接近利用 18T tokens 訓練的 Qwen2.5-1.5B 模型,同時保持最低激活參數數量。
其中,PLM 在通用知識理解(ARC)、數學(GSM8K)任務中表現出色,在代碼能力評測(HumanEval)基準上更是以 64.6 分位居榜首。

-
項目論文:https://arxiv.org/abs/2503.12167
-
項目網站:https://www.project-plm.com
-
項目地址:https://github.com/plm-team/PLM
-
模型地址:Hugging Face: https://huggingface.co/PLM-Team
團隊還將模型適配至多種邊緣設備,在 Snapdragon 等芯片上展現出優於同層數模型的吞吐優勢。PLM 團隊不僅開源模型權重,還提供了從架構設計到部署的完整技術報告,並計劃逐步開源訓練數據集及相關代碼腳本。

剖析 PLM:1+1 能否大於 2?

核心模塊一:Multi-head Latent Attention(MLA)—— 把 KV 緩存壓縮到極致
PLM 採用了 Deepseek 提出的 MLA(Multi-Head Latent Attention)注意力機制,並首次將其應用於 2B 參數以下的模型中。為適應端側系統的需求,PLM 對 Deepseek 的 MLA 進行了適度優化,去除了訓練階段用於降低成本的 Q 矩陣壓縮過程,同時保留了 KV 矩陣的 512 維度。此外,PLM 通過解耦的位置編碼機制,確保了模型對長程依賴信息的有效捕捉。
核心模塊二:平方 ReLU 激活 —— 讓計算「稀疏化」
PLM 通過去除門控機制簡化了前饋神經網絡,從而有效降低了計算複雜度和內存消耗。傳統 SwiGLU 激活函數導致 MLP 層計算密集,而 PLM 採用了 ReLU² 作為替代。ReLU² 是一種在性能和稀疏性之間實現最優平衡的激活函數,特別適合稀疏計算場景。其定義如下:

這個設計使得 MLP 層激活稀疏度達到 90.9%,整體計算量減少 26%。此外,從硬件角度出發,零值激活能夠觸髮指令級優化。這一設計理唸成功地將模型與系統的聯合優化整合到大語言模型架構中。
訓練策略:充分利用開源數據
三階段鍛造 PLM-1.8B
PLM 團隊精心設計了一條訓練流水線,僅使用未精細設計與配比的 2.48B 預訓練數據,PLM 性能就達到企業級水平。所有訓練數據均來自開源社區。

預訓練

第一階段累計約 1.65T 的 Token。這個階段中,學習率是屬於預熱和穩定的階段,模型的 Loss 也在 2.3 附近有收斂趨勢。第二階段累計約 550B 的 Token。在這個階段,模型的 Loss 隨著學習的衰減快速下降。第三階段累計約 280B 的 Token。
這個階段里,保持第二階段的最小學習率訓練,模型進行最後的高質量知識吸收,直到 loss 逐漸收斂。整個預訓練中,始終保持中英數據比例 5:2。

SFT 階段
監督微調數據遵循「由淺入深」的數據準備方法,以漸進的方式進行監督微調過程。PLM 的 SFT 訓練分為基本指令微調和高難度指令微調。下面是各個階段的提升效果。

強化學習階段
PLM 在偏好訓練階段沿用了團隊先前提出的 ARIES 訓練方法,以解決經過一般的對齊學習會經多輪自我改進後,性能會顯著下降的問題。

性能實測:真的可以
PLM 採用獨特的模型架構,對比的基線模型涵蓋了當前最先進的 2B 參數量級模型,具體對比如下。

實驗表明,PLM 表現頗具競爭力,平均分(57.29)位列第三,僅次於 Qwen2.5-1.5B(59.25)和 Yulan-Mini-2.4B(57.51)。PLM-1.8B 在 HumanEval 中獲得了所有模型中的最高分,在 ARC-C、ARC-E、MBPP 和 BoolQ 中排名第二,略遜於行業領先的 Qwen2.5-1.5B,需指出的是,Qwen 系列使用了 18T 閉源語料庫。
另一方面,與 Yulan-Mini-2.4B 相比,PLM-1.8B 在編碼和邏輯推理任務中旗鼓相當。此外,PLM 僅包含 1.8B 參數和 32 層(Yulan-Mini 為 56 層),推理延遲會較低。綜上所述,PLM-1.8B 在基本知識理解、編碼和簡單推理任務中表現強勁且可靠,是一款值得關注的模型。

場景實測:從服務器到樹莓派,全場景通吃
除了基本的模型能力評估,PLM 團隊還在在 5 類硬件平台完成部署驗證,並給出了實際的吞吐量數據。

文章全面評估了不同硬件平台和量化級別的各種邊緣大小 LLM 的推理延遲,包括高性能 GPU(NVIDIA A10、Orin NX)、Apple 的 M3 芯片、Qualcomm 的 Snapdragon 8 Gen 3 和 BCM2712 等嵌入式系統。
評估數據揭示了幾個顯著特徵:MLA 的確增加了計算量,ReLU2 的確可以提升模型推理速度,模型層數會顯著影響端側設備上推理的速度。
PLM 團隊的實驗表明,這些關鍵點恰好觸及了端側計算中最需要關注的內存、算力和 I/O 三個核心維度。PLM 團隊通過模型與硬件的協同設計,在這一領域展現了其獨特的優勢。
算法層面,PLM 做到了稀疏性與低秩的平衡:MLA 壓縮 KV 緩存,ReLU² 激活削減計算,二者互補突破內存 – 算力瓶頸。系統層面,PLM 深度適配 TVM、llama.cpp 等框架,實現高效的量化與編譯優化。
理解 PLM 的 MLA 和稀疏激活
在實際探索中,MLA 的引入會顯著增加計算量。然而,PLM 通過捨棄 Q 矩陣的低秩壓縮來降低推理計算複雜度,並結合稀疏激活函數,成功避免了 MiniCPM3 在預填充或解碼階段的低效問題,從而在特定場景中展現出明顯優勢。

PLM 團隊在較長文本序列上評估了其模型性能,結果表明,當序列長度達到一定閾值後,PLM 的表現優於同深度的 GQA 模型 Fox。
因此,儘管 MLA 增加了計算負載,其對緩存利用率、推理效率和內存消耗的優化,使得 PLM 在邊緣設備上展現出高效、低延遲的性能,為實際應用提供了顯著優勢。
再來看 PLM 的稀疏化設計,在邊緣模型中展現出了更高的普適性與高效性。從系統角度來看,零計算已被高度優化,使得稀疏化在邊緣設備上的部署帶來顯著性能提升。
此外,由於邊緣設備的計算資源有限,模型通常無法完全加載到 GPU 或 RAM,需要 OffLoad 到緩存甚至存儲中。在此情況下,深度學習模型可採用分層加載,將當前所需參數調入計算單元。
因此,最小化每層計算量至關重要。PLM 通過 KV 緩存存儲與稀疏激活減少計算開銷,有效緩解該問題。
PLM 團隊實驗驗證了推理所需的最小參數量。他們對所有模型進行相同稀疏化(即將激活函數後的最小值設為 0),並測試保持建模性能(困惑度下降 1)所需的最少參數量(如下圖)。

具體而言,PLM 團隊繪製了稀疏率(0~1)與困惑度差異的關係圖,以分析神經激活減少的性能成本。不同模型的曲線揭示了各自對稀疏度的敏感性。理論上的「理想點」—— 完全稀疏且困惑度不增加 —— 是無法實現的,因為 MLP 層完全停用將損害模型質量。

在此背景下,PLM 展現出顯著優勢,僅需激活 74.3% 參數,MLP 稀疏率達 90.9%,遠低於同等規模模型,推理所需參數量最低。
PLM 團隊已全面開放資源,為社區提供了一個小型且易於使用的 MLA 模型,使科研工作者能夠在消費級顯卡上開展對 MLA 的研究。同時,PLM 為端側應用廠商提供了一個高性能的端側模型,拓寬了選擇範圍,並支持基於 PLM 稀疏激活架構的高效模型部署與開發。
結語
學術界在大規模模型結構實驗方面面臨諸多挑戰,而堅持從頭預訓練的團隊更是少之又少。PLM 團隊在計算資源有限和數據質量參差不齊的情況下,始終堅信開源社區提供的數據和技術能夠為學術界的持續探索提供強大支持。
未來,PLM 團隊將繼續致力於探索適用於邊緣設備的大模型,訓練更具創新性的架構,並實現更高效的邊緣設備部署。PLM 團隊認為,未來的語言模型不應僅僅是參數的堆砌,而應是效率與智能的精密平衡。PLM的探索,正是向著這一理想邁出的關鍵一步。