Mixture-of-Experts（MoE）在推理時僅激活每個 token 所需的一小部分專家，憑借其稀疏激活的特點，已成為當前 LLM 中的主流架構。然而，MoE 雖然顯著降低了推理時的計算量，但整體參數(shù)規(guī)模依然大于同等性能的 Dense 模型，因此在顯存資源極為受限的端側部署場景中，仍然面臨較大挑戰(zhàn)。

現(xiàn)有的主流解決方案是專家卸載（Expert Offloading），即將專家模塊存儲在下層存儲設備（如 CPU 內存甚至磁盤）中，在推理時按需加載激活的專家到顯存進行計算。但這一方法存在兩大主要缺陷：

1. 由于不同 token 通常激活的專家不同，每一步推理都需要頻繁加載不同的專家，導致顯著的推理延遲；
2. 在批量解碼場景中，各 token 可能需要不同的專家，在最壞情況下，甚至需要將一整層的所有專家加載到顯存中，進一步加劇顯存壓力并帶來額外的推理延遲。

為了解決上述問題，來自北大和華為諾亞的研究人員提出了 Mixture-of-Lookup-Experts（MoLE），一種在推理階段可重參數(shù)化的新型 MoE 架構。

• 論文鏈接：https://arxiv.org/pdf/2503.15798
• 代碼鏈接：https://github.com/JieShibo/MoLE

思考

本文的核心思考是，在專家卸載方案中，需要將專家模塊加載到顯存，主要是為了在 GPU 上執(zhí)行高效的矩陣運算。換句話說，如果專家的計算過程能夠繞過矩陣運算的需求，就可以避免將專家權重加載到顯存，從而根本上規(guī)避頻繁加載帶來的開銷。直觀來看，專家模塊本質上是一個神經網絡，用于建模輸入到輸出的映射。如果能夠在推理前預先計算出所有可能的輸入 - 輸出對應關系，并將其存儲為查找表，那么在推理時即可用簡單的查找操作代替矩陣運算。

一般而言，神經網絡所建模的映射通常涉及無限多的輸入 - 輸出對，因此，要實現(xiàn)查找表替代矩陣運算，首先需要確保專家模塊的輸入來自一個離散且有限的集合，并且這一離散化方法能夠適配大規(guī)模預訓練任務。其次，由于查找操作發(fā)生在 GPU 之外，還需要保證檢索過程本身不依賴密集計算，避免引入新的性能瓶頸。

基于這一思考，作者注意到，大規(guī)模語言模型（LLM）中的 embedding token（即 embedding 層的輸出）天然具備離散且有限的特性，其數(shù)量與詞表大小一致，滿足了離散有限要求。并且 embedding token 可以通過 token ID 唯一確定，因此查找表的檢索可以采用高效的直接尋址。因此，MoLE 設計中將專家的輸入由中間特征改為 embedding token，從而滿足了查找表構建的所有要求。

訓練階段

在訓練階段，MoLE 相較于傳統(tǒng)的 MoE 架構存在三個主要區(qū)別：

1. 輸入調整：將所有路由專家（routed experts）的輸入由上一層的輸出，改為淺層的 embedding token，以確保專家模塊可以被預計算并存儲為查找表。
2. 激活策略：由于查找表檢索在推理時無需額外計算，MoLE 無需依賴稀疏激活來控制推理計算量，因此在訓練中選擇激活所有路由專家。
3. 損失設計：鑒于不再需要通過稀疏激活實現(xiàn)負載均衡，MoLE 訓練時僅使用語言建模損失，不再引入額外的負載均衡損失項。

除此之外，MoLE 的其他設計與標準 MoE 保持一致，包括路由（router）模塊和共享專家（shared experts），依然使用上一層的輸出作為輸入。計算流程如下

推理階段

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

在查找表構建完成后，所有原始的路由專家模塊將被刪除，查找表則被存儲在下層存儲設備中。在推理階段，對于每個 token，根據(jù)其 token ID 直接在查找表中檢索對應的專家輸出，并將檢索到的輸出加載到顯存中，用于后續(xù)的推理計算。整體計算流程如下所示：

復雜度分析

如表所示，在推理階段，MoLE 的計算過程中僅保留了共享專家模塊，因此只有共享專家被激活并參與計算，其整體計算量與具有相同激活參數(shù)量的 Dense 模型和傳統(tǒng) MoE 模型相當。相比之下，MoLE 在推理時僅需傳輸專家輸出的結果向量，而傳統(tǒng) MoE 需要傳輸中間維度 D_r 的專家權重矩陣，因此 MoLE 的傳輸開銷相比 MoE 減少了數(shù)個量級。在存儲開銷方面，對于端側部署的模型，詞表大小 | V | 通常在數(shù)萬左右，與 D_r 為相同數(shù)量級，因此 MoLE 查找表的存儲需求與單個專家模塊的大小處于同一數(shù)量級，不會帶來顯著額外的存儲負擔。

實驗結果

本文在 Pile 數(shù)據(jù)集的 100B-token 子集上訓練了 160M、410M、1B 激活參數(shù)量的 Dense、MoE 和 MoLE 模型。對于 MoE 和 MoLE 模型，控制兩者的訓練階段參數(shù)量相等。由于實驗中以及近期 OLMoE 的結果都發(fā)現(xiàn)共享專家會降低 MoE 的性能，我們對 MoE 只采用了路由專家。MoLE 的專家大小與 Dense 的 FFN 保持一致，而 MoE 由于需要激活兩個專家，其專家大小為 dense FFN 的一半，但專家數(shù)量是 MoLE 的兩倍。

實驗結果表明 MoLE 在相同訓練參數(shù)量和推理激活參數(shù)量（即顯存使用量）下，具有與 MoE 相當?shù)男阅埽啾?Dense 有顯著提升。與專家卸載的 MoE 相比，MoLE 減少了千倍以上的傳輸開銷。

在 V100 上進行的評測結果表明，在顯存用量一定的前提下，MoLE 的推理延遲與 Dense 基本一致，顯著優(yōu)于專家卸載的 MoE。在批量解碼場景下，隨著 batch size 的增加，MoE 的推理延遲迅速上升，而 MoLE 與 Dense 模型的延遲則基本保持穩(wěn)定，進一步展現(xiàn)了 MoLE 在高吞吐量推理任務中的優(yōu)勢。

此外，消融實驗表明，MoLE 的訓練確實不需要輔助損失。

在專家數(shù)量提升時，模型性能也會提升

然而，如果僅增大專家的隱層維度，由于查找表的大小保持不變，當專家規(guī)模增大到一定程度時，推理性能將受限于查找表的固定大小，最終達到飽和。

作者通過將一個 MoE 模型逐步修改為 MoLE 模型，系統(tǒng)性地探索了 MoLE 各組成部分對性能的影響。實驗結果表明，使用淺層的 embedding token 作為專家輸入確實會削弱模型的表達能力，這是由于輸入中缺乏豐富的上下文信息所致。然而，激活所有專家有效彌補了這一損失，使得 MoLE 最終能夠達到與 MoE 相當?shù)男阅芩健?/p>

需要注意的是，路由專家的輸入不包含上下文信息，并不意味著專家無法影響模型對上下文的處理。實際上，專家可以通過改變其輸出，從而間接影響后續(xù)注意力層的輸入，實現(xiàn)對上下文的建模。此外，共享專家和路由仍然接收包含上下文信息的輸入，進一步保障了模型對上下文理解能力的保留。

最后，作者發(fā)現(xiàn)查找表中仍然存在較大程度的冗余。即使將查找表壓縮至 3-bit 精度（例如使用 NF3 格式），模型性能依然能夠基本保持不變。這表明查找表在存儲開銷上仍具有進一步壓縮和優(yōu)化的潛力。

總結

綜上，本文提出了 MoLE，一種面向端側部署優(yōu)化的新型 MoE 架構。通過將專家輸入改為淺層的 embedding token，并采用查找表替代傳統(tǒng)的矩陣運算，MoLE 有效解決了傳統(tǒng) MoE 在推理階段面臨的顯存開銷大、傳輸延遲高的問題。實驗結果表明，MoLE 在保持與 MoE 相當性能的同時，大幅降低了推理延遲，尤其在批量解碼場景中展現(xiàn)出顯著優(yōu)勢。