ICML 2025 Spotlight｜華為諾亞提出端側(cè)大模型新架構(gòu)MoLE，內(nèi)存搬運(yùn)代價(jià)降低1000倍
Mixture-of-Experts（MoE）在推理時(shí)僅激活每個(gè) token 所需的一小部分專(zhuān)家，憑借其稀疏激活的特點(diǎn)，已成為當(dāng)前 LLM 中的主流架構(gòu)。然而，MoE 雖然顯著降低了推理時(shí)的計(jì)算量，但整體參數(shù)規(guī)模依然大于同等性能的 Dense 模型，因此在顯存資源極為受限的端側(cè)部署場(chǎng)景中，仍然面臨較大挑戰(zhàn)。
現(xiàn)有的主流解決方案是專(zhuān)家卸載（Expert Offloading），即將專(zhuān)家模塊存儲(chǔ)在下層存儲(chǔ)設(shè)備（如 CPU 內(nèi)存甚至磁盤(pán)）中，在推理時(shí)按需加載激活的專(zhuān)家到顯存進(jìn)行計(jì)算。但這一方法存在兩大主要缺陷：
- 由于不同 token 通常激活的專(zhuān)家不同，每一步推理都需要頻繁加載不同的專(zhuān)家，導(dǎo)致顯著的推理延遲；
- 在批量解碼場(chǎng)景中，各 token 可能需要不同的專(zhuān)家，在最壞情況下，甚至需要將一整層的所有專(zhuān)家加載到顯存中，進(jìn)一步加劇顯存壓力并帶來(lái)額外的推理延遲。
為了解決上述問(wèn)題，來(lái)自北大和華為諾亞的研究人員提出了 Mixture-of-Lookup-Experts（MoLE），一種在推理階段可重參數(shù)化的新型 MoE 架構(gòu)。
- 論文鏈接：https://arxiv.org/pdf/2503.15798
- 代碼鏈接：https://github.com/JieShibo/MoLE
思考
本文的核心思考是，在專(zhuān)家卸載方案中，需要將專(zhuān)家模塊加載到顯存，主要是為了在 GPU 上執(zhí)行高效的矩陣運(yùn)算。換句話說(shuō)，如果專(zhuān)家的計(jì)算過(guò)程能夠繞過(guò)矩陣運(yùn)算的需求，就可以避免將專(zhuān)家權(quán)重加載到顯存，從而根本上規(guī)避頻繁加載帶來(lái)的開(kāi)銷(xiāo)。直觀來(lái)看，專(zhuān)家模塊本質(zhì)上是一個(gè)神經(jīng)網(wǎng)絡(luò)，用于建模輸入到輸出的映射。如果能夠在推理前預(yù)先計(jì)算出所有可能的輸入 - 輸出對(duì)應(yīng)關(guān)系，并將其存儲(chǔ)為查找表，那么在推理時(shí)即可用簡(jiǎn)單的查找操作代替矩陣運(yùn)算。
一般而言，神經(jīng)網(wǎng)絡(luò)所建模的映射通常涉及無(wú)限多的輸入 - 輸出對(duì)，因此，要實(shí)現(xiàn)查找表替代矩陣運(yùn)算，首先需要確保專(zhuān)家模塊的輸入來(lái)自一個(gè)離散且有限的集合，并且這一離散化方法能夠適配大規(guī)模預(yù)訓(xùn)練任務(wù)。其次，由于查找操作發(fā)生在 GPU 之外，還需要保證檢索過(guò)程本身不依賴(lài)密集計(jì)算，避免引入新的性能瓶頸。
基于這一思考，作者注意到，大規(guī)模語(yǔ)言模型（LLM）中的 embedding token（即 embedding 層的輸出）天然具備離散且有限的特性，其數(shù)量與詞表大小一致，滿足了離散有限要求。并且 embedding token 可以通過(guò) token ID 唯一確定，因此查找表的檢索可以采用高效的直接尋址。因此，MoLE 設(shè)計(jì)中將專(zhuān)家的輸入由中間特征改為 embedding token，從而滿足了查找表構(gòu)建的所有要求。
訓(xùn)練階段
在訓(xùn)練階段，MoLE 相較于傳統(tǒng)的 MoE 架構(gòu)存在三個(gè)主要區(qū)別：
- 輸入調(diào)整：將所有路由專(zhuān)家（routed experts）的輸入由上一層的輸出，改為淺層的 embedding token，以確保專(zhuān)家模塊可以被預(yù)計(jì)算并存儲(chǔ)為查找表。
- 激活策略：由于查找表檢索在推理時(shí)無(wú)需額外計(jì)算，MoLE 無(wú)需依賴(lài)稀疏激活來(lái)控制推理計(jì)算量，因此在訓(xùn)練中選擇激活所有路由專(zhuān)家。
- 損失設(shè)計(jì)：鑒于不再需要通過(guò)稀疏激活實(shí)現(xiàn)負(fù)載均衡，MoLE 訓(xùn)練時(shí)僅使用語(yǔ)言建模損失，不再引入額外的負(fù)載均衡損失項(xiàng)。
除此之外，MoLE 的其他設(shè)計(jì)與標(biāo)準(zhǔn) MoE 保持一致，包括路由（router）模塊和共享專(zhuān)家（shared experts），依然使用上一層的輸出作為輸入。計(jì)算流程如下
推理階段
在推理前，MoLE 通過(guò)預(yù)先構(gòu)建查找表來(lái)完成專(zhuān)家模塊的重參數(shù)化。具體而言，embedding 層的權(quán)重矩陣本身即包含了所有 embedding token 的向量表示，因此可以直接以該權(quán)重矩陣作為專(zhuān)家模塊的輸入，并通過(guò)各個(gè)路由專(zhuān)家分別計(jì)算對(duì)應(yīng)的輸出。這樣，便可以高效地獲得完整的 token id 到專(zhuān)家輸出的映射集合，用于后續(xù)的查找操作。具體過(guò)程如下所示：
在查找表構(gòu)建完成后，所有原始的路由專(zhuān)家模塊將被刪除，查找表則被存儲(chǔ)在下層存儲(chǔ)設(shè)備中。在推理階段，對(duì)于每個(gè) token，根據(jù)其 token ID 直接在查找表中檢索對(duì)應(yīng)的專(zhuān)家輸出，并將檢索到的輸出加載到顯存中，用于后續(xù)的推理計(jì)算。整體計(jì)算流程如下所示：
復(fù)雜度分析
如表所示，在推理階段，MoLE 的計(jì)算過(guò)程中僅保留了共享專(zhuān)家模塊，因此只有共享專(zhuān)家被激活并參與計(jì)算，其整體計(jì)算量與具有相同激活參數(shù)量的 Dense 模型和傳統(tǒng) MoE 模型相當(dāng)。相比之下，MoLE 在推理時(shí)僅需傳輸專(zhuān)家輸出的結(jié)果向量，而傳統(tǒng) MoE 需要傳輸中間維度 D_r 的專(zhuān)家權(quán)重矩陣，因此 MoLE 的傳輸開(kāi)銷(xiāo)相比 MoE 減少了數(shù)個(gè)量級(jí)。在存儲(chǔ)開(kāi)銷(xiāo)方面，對(duì)于端側(cè)部署的模型，詞表大小 | V | 通常在數(shù)萬(wàn)左右，與 D_r 為相同數(shù)量級(jí)，因此 MoLE 查找表的存儲(chǔ)需求與單個(gè)專(zhuān)家模塊的大小處于同一數(shù)量級(jí)，不會(huì)帶來(lái)顯著額外的存儲(chǔ)負(fù)擔(dān)。
實(shí)驗(yàn)結(jié)果
本文在 Pile 數(shù)據(jù)集的 100B-token 子集上訓(xùn)練了 160M、410M、1B 激活參數(shù)量的 Dense、MoE 和 MoLE 模型。對(duì)于 MoE 和 MoLE 模型，控制兩者的訓(xùn)練階段參數(shù)量相等。由于實(shí)驗(yàn)中以及近期 OLMoE 的結(jié)果都發(fā)現(xiàn)共享專(zhuān)家會(huì)降低 MoE 的性能，我們對(duì) MoE 只采用了路由專(zhuān)家。MoLE 的專(zhuān)家大小與 Dense 的 FFN 保持一致，而 MoE 由于需要激活兩個(gè)專(zhuān)家，其專(zhuān)家大小為 dense FFN 的一半，但專(zhuān)家數(shù)量是 MoLE 的兩倍。
實(shí)驗(yàn)結(jié)果表明 MoLE 在相同訓(xùn)練參數(shù)量和推理激活參數(shù)量（即顯存使用量）下，具有與 MoE 相當(dāng)?shù)男阅?，相?Dense 有顯著提升。與專(zhuān)家卸載的 MoE 相比，MoLE 減少了千倍以上的傳輸開(kāi)銷(xiāo)。
在 V100 上進(jìn)行的評(píng)測(cè)結(jié)果表明，在顯存用量一定的前提下，MoLE 的推理延遲與 Dense 基本一致，顯著優(yōu)于專(zhuān)家卸載的 MoE。在批量解碼場(chǎng)景下，隨著 batch size 的增加，MoE 的推理延遲迅速上升，而 MoLE 與 Dense 模型的延遲則基本保持穩(wěn)定，進(jìn)一步展現(xiàn)了 MoLE 在高吞吐量推理任務(wù)中的優(yōu)勢(shì)。
此外，消融實(shí)驗(yàn)表明，MoLE 的訓(xùn)練確實(shí)不需要輔助損失。
在專(zhuān)家數(shù)量提升時(shí)，模型性能也會(huì)提升
然而，如果僅增大專(zhuān)家的隱層維度，由于查找表的大小保持不變，當(dāng)專(zhuān)家規(guī)模增大到一定程度時(shí)，推理性能將受限于查找表的固定大小，最終達(dá)到飽和。
作者通過(guò)將一個(gè) MoE 模型逐步修改為 MoLE 模型，系統(tǒng)性地探索了 MoLE 各組成部分對(duì)性能的影響。實(shí)驗(yàn)結(jié)果表明，使用淺層的 embedding token 作為專(zhuān)家輸入確實(shí)會(huì)削弱模型的表達(dá)能力，這是由于輸入中缺乏豐富的上下文信息所致。然而，激活所有專(zhuān)家有效彌補(bǔ)了這一損失，使得 MoLE 最終能夠達(dá)到與 MoE 相當(dāng)?shù)男阅芩健?/span>
需要注意的是，路由專(zhuān)家的輸入不包含上下文信息，并不意味著專(zhuān)家無(wú)法影響模型對(duì)上下文的處理。實(shí)際上，專(zhuān)家可以通過(guò)改變其輸出，從而間接影響后續(xù)注意力層的輸入，實(shí)現(xiàn)對(duì)上下文的建模。此外，共享專(zhuān)家和路由仍然接收包含上下文信息的輸入，進(jìn)一步保障了模型對(duì)上下文理解能力的保留。
最后，作者發(fā)現(xiàn)查找表中仍然存在較大程度的冗余。即使將查找表壓縮至 3-bit 精度（例如使用 NF3 格式），模型性能依然能夠基本保持不變。這表明查找表在存儲(chǔ)開(kāi)銷(xiāo)上仍具有進(jìn)一步壓縮和優(yōu)化的潛力。
總結(jié)
綜上，本文提出了 MoLE，一種面向端側(cè)部署優(yōu)化的新型 MoE 架構(gòu)。通過(guò)將專(zhuān)家輸入改為淺層的 embedding token，并采用查找表替代傳統(tǒng)的矩陣運(yùn)算，MoLE 有效解決了傳統(tǒng) MoE 在推理階段面臨的顯存開(kāi)銷(xiāo)大、傳輸延遲高的問(wèn)題。實(shí)驗(yàn)結(jié)果表明，MoLE 在保持與 MoE 相當(dāng)性能的同時(shí)，大幅降低了推理延遲，尤其在批量解碼場(chǎng)景中展現(xiàn)出顯著優(yōu)勢(shì)。
