ICML2025|華為諾亞提出端側大模型新架構MoLE，記憶體搬運代價降低1000倍

2025-05-08 09:32 CVer

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達
點選進入—>【頂會/頂刊】投稿交流群

新增微訊號：CVer2233，小助手會拉你進群！
掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

轉載自：機器之心
Mixture-of-Experts（MoE）在推理時僅啟用每個 token 所需的一小部分專家，憑藉其稀疏啟用的特點，已成為當前 LLM 中的主流架構。然而，MoE 雖然顯著降低了推理時的計算量，但整體引數規模依然大於同等效能的 Dense 模型，因此在視訊記憶體資源極為受限的端側部署場景中，仍然面臨較大挑戰。
現有的主流解決方案是專家解除安裝（Expert Offloading），即將專家模組儲存在下層儲存裝置（如 CPU 記憶體甚至磁碟）中，在推理時按需載入啟用的專家到視訊記憶體進行計算。但這一方法存在兩大主要缺陷：
  1. 由於不同 token 通常啟用的專家不同，每一步推理都需要頻繁載入不同的專家，導致顯著的推理延遲；
  2. 在批次解碼場景中，各 token 可能需要不同的專家，在最壞情況下，甚至需要將一整層的所有專家載入到視訊記憶體中，進一步加劇視訊記憶體壓力並帶來額外的推理延遲。
為了解決上述問題，來自北大和華為諾亞的研究人員提出了 Mixture-of-Lookup-Experts（MoLE），一種在推理階段可重引數化的新型 MoE 架構。
  • 論文連結：https://arxiv.org/pdf/2503.15798
  • 程式碼連結：https://github.com/JieShibo/MoLE
思考
本文的核心思考是，在專家解除安裝方案中，需要將專家模組載入到視訊記憶體，主要是為了在 GPU 上執行高效的矩陣運算。換句話說，如果專家的計算過程能夠繞過矩陣運算的需求，就可以避免將專家權重載入到視訊記憶體，從而根本上規避頻繁載入帶來的開銷。直觀來看，專家模組本質上是一個神經網路，用於建模輸入到輸出的對映。如果能夠在推理前預先計算出所有可能的輸入 – 輸出對應關係，並將其儲存為查詢表，那麼在推理時即可用簡單的查詢操作代替矩陣運算。
一般而言，神經網路所建模的對映通常涉及無限多的輸入 – 輸出對，因此，要實現查詢表替代矩陣運算，首先需要確保專家模組的輸入來自一個離散且有限的集合，並且這一離散化方法能夠適配大規模預訓練任務。其次，由於查詢操作發生在 GPU 之外，還需要保證檢索過程本身不依賴密集計算，避免引入新的效能瓶頸。
基於這一思考，作者注意到，大規模語言模型（LLM）中的 embedding token（即 embedding 層的輸出）天然具備離散且有限的特性，其數量與詞表大小一致，滿足了離散有限要求。並且 embedding token 可以透過 token ID 唯一確定，因此查詢表的檢索可以採用高效的直接定址。因此，MoLE 設計中將專家的輸入由中間特徵改為 embedding token，從而滿足了查詢表構建的所有要求。
訓練階段
在訓練階段，MoLE 相較於傳統的 MoE 架構存在三個主要區別：
  1. 輸入調整：將所有路由專家（routed experts）的輸入由上一層的輸出，改為淺層的 embedding token，以確保專家模組可以被預計算並存儲為查詢表。
  2. 啟用策略：由於查詢表檢索在推理時無需額外計算，MoLE 無需依賴稀疏啟用來控制推理計算量，因此在訓練中選擇啟用所有路由專家。
  3. 損失設計：鑑於不再需要透過稀疏啟用實現負載均衡，MoLE 訓練時僅使用語言建模損失，不再引入額外的負載均衡損失項。
除此之外，MoLE 的其他設計與標準 MoE 保持一致，包括路由（router）模組和共享專家（shared experts），依然使用上一層的輸出作為輸入。計算流程如下

推理階段
在推理前，MoLE 透過預先構建查詢表來完成專家模組的重引數化。具體而言，embedding 層的權重矩陣本身即包含了所有 embedding token 的向量表示，因此可以直接以該權重矩陣作為專家模組的輸入，並透過各個路由專家分別計算對應的輸出。這樣，便可以高效地獲得完整的 token id 到專家輸出的對映集合，用於後續的查詢操作。具體過程如下所示：

在查詢表構建完成後，所有原始的路由專家模組將被刪除，查詢表則被儲存在下層儲存裝置中。在推理階段，對於每個 token，根據其 token ID 直接在查詢表中檢索對應的專家輸出，並將檢索到的輸出載入到視訊記憶體中，用於後續的推理計算。整體計算流程如下所示：
複雜度分析
如表所示，在推理階段，MoLE 的計算過程中僅保留了共享專家模組，因此只有共享專家被啟用並參與計算，其整體計算量與具有相同啟用引數量的 Dense 模型和傳統 MoE 模型相當。相比之下，MoLE 在推理時僅需傳輸專家輸出的結果向量，而傳統 MoE 需要傳輸中間維度 D_r 的專家權重矩陣，因此 MoLE 的傳輸開銷相比 MoE 減少了數個量級。在儲存開銷方面，對於端側部署的模型，詞表大小 | V | 通常在數萬左右，與 D_r 為相同數量級，因此 MoLE 查詢表的儲存需求與單個專家模組的大小處於同一數量級，不會帶來顯著額外的儲存負擔。
實驗結果
本文在 Pile 資料集的 100B-token 子集上訓練了 160M、410M、1B 啟用引數量的 Dense、MoE 和 MoLE 模型。對於 MoE 和 MoLE 模型，控制兩者的訓練階段引數量相等。由於實驗中以及近期 OLMoE 的結果都發現共享專家會降低 MoE 的效能，我們對 MoE 只採用了路由專家。MoLE 的專家大小與 Dense 的 FFN 保持一致，而 MoE 由於需要啟用兩個專家，其專家大小為 dense FFN 的一半，但專家數量是 MoLE 的兩倍。
實驗結果表明 MoLE 在相同訓練引數量和推理啟用引數量（即視訊記憶體使用量）下，具有與 MoE 相當的效能，相比 Dense 有顯著提升。與專家解除安裝的 MoE 相比，MoLE 減少了千倍以上的傳輸開銷。
在 V100 上進行的評測結果表明，在視訊記憶體用量一定的前提下，MoLE 的推理延遲與 Dense 基本一致，顯著優於專家解除安裝的 MoE。在批次解碼場景下，隨著 batch size 的增加，MoE 的推理延遲迅速上升，而 MoLE 與 Dense 模型的延遲則基本保持穩定，進一步展現了 MoLE 在高吞吐量推理任務中的優勢。
此外，消融實驗表明，MoLE 的訓練確實不需要輔助損失。
在專家數量提升時，模型效能也會提升
然而，如果僅增大專家的隱層維度，由於查詢表的大小保持不變，當專家規模增大到一定程度時，推理效能將受限於查詢表的固定大小，最終達到飽和。
作者透過將一個 MoE 模型逐步修改為 MoLE 模型，系統性地探索了 MoLE 各組成部分對效能的影響。實驗結果表明，使用淺層的 embedding token 作為專家輸入確實會削弱模型的表達能力，這是由於輸入中缺乏豐富的上下文資訊所致。然而，啟用所有專家有效彌補了這一損失，使得 MoLE 最終能夠達到與 MoE 相當的效能水平。
需要注意的是，路由專家的輸入不包含上下文資訊，並不意味著專家無法影響模型對上下文的處理。實際上，專家可以透過改變其輸出，從而間接影響後續注意力層的輸入，實現對上下文的建模。此外，共享專家和路由仍然接收包含上下文資訊的輸入，進一步保障了模型對上下文理解能力的保留。
最後，作者發現查詢表中仍然存在較大程度的冗餘。即使將查詢表壓縮至 3-bit 精度（例如使用 NF3 格式），模型效能依然能夠基本保持不變。這表明查詢表在儲存開銷上仍具有進一步壓縮和最佳化的潛力。
總結
綜上，本文提出了 MoLE，一種面向端側部署最佳化的新型 MoE 架構。透過將專家輸入改為淺層的 embedding token，並採用查詢表替代傳統的矩陣運算，MoLE 有效解決了傳統 MoE 在推理階段面臨的視訊記憶體開銷大、傳輸延遲高的問題。實驗結果表明，MoLE 在保持與 MoE 相當效能的同時，大幅降低了推理延遲，尤其在批次解碼場景中展現出顯著優勢。

何愷明在MIT授課的課件PPT下載
在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

CVPR 2025 論文和程式碼下載

在CVer公眾號後臺回覆：CVPR2025，即可下載CVPR 2025論文和程式碼開源的論文合集

ECCV 2024 論文和程式碼下載

在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集
CV垂直方向和論文投稿交流群成立
掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-垂直方向和論文投稿微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者論文投稿+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群

▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）人數破萬！如果你想要了解最新最快最好的CV/DL/AI論文、實戰專案、行業前沿、從入門到精通學習教程等資料，一定要掃描下方二維碼，加入CVer知識星球！最強助力你的科研和工作！

▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請點贊和在看


相關文章