混合專家（MoE）是個好方法，支持著現在一些非常優秀的大模型，比如谷歌家的 Gemini 1.5 以及備受關注的 Mixtral 8x7B。

稀疏混合專家（SMoE）可在不顯著增加訓練和推理成本的前提下提升模型的能力。比如 Mixtral 8×7B 就是一個 SMoE 模型，其包含 8 個專家（共 7B 參數），而其表現卻可以超過或比肩 LLaMA-2 70B 和 GPT-3.5。

但是，它也有兩個問題。一是專家激活率低 —— 也就是搞不好會出現下圖這種情況：

具體來說，就是在優化時只有一小部分專家會被激活，如圖 1a 所示（8.33% 的激活率），這會導致在學習應對復雜任務的大量專家時，會出現性能次優和效果不佳的問題。

二是無法細粒度地分析單個 token 的多重語義概念，比如多義詞和具有多重細節的圖塊。

近日，微軟研究院和清華大學提出了多頭混合專家（MH-MoE）。顧名思義，MH-MoE 采用了多頭機制，可將每個輸入 token 分成多個子 token。然后將這些子 token 分配給一組多樣化的專家并行處理，之后再無縫地將它們整合進原來的 token 形式。

• 論文標題：Multi-Head Mixture-of-Experts
• 論文地址：https://arxiv.org/pdf/2404.15045
• 代碼地址：https://github.com/yushuiwx/MH-MoE

圖 2 展示了 MH-MoE 的工作流程。可以看到，當輸入單個 token 時，MH-MoE 會將其分成 4 個子 token，進而激活 4 個專家，而 SMoE 僅激活 1 個專家。

如圖 2 所示，分配給專家 3 和 2 的子 token 包含對圖塊內每個角色動作的詳細理解，而分配給專家 1 和 4 的子 token 則顯式地建模了錯誤的同源詞「camera」的語義。

專家處理完成后，再將子 token 無縫地重新整合進原來的 token 形式，由此可以避免后續非并行層（例如注意力層）的任何額外計算負擔，同時還集成從多個專家捕獲的語義信息。

這樣的操作可讓 MH-MoE 從整體上關注來自不同專家內不同表征空間的信息，從而可以加深上下文理解能力，同時提升專家激活率。該項目的代碼也將發布。

MH-MoE 的具有以下優勢：

1. 專家激活率更高且擴展性更好。MH-MoE 能優化幾乎所有專家，從而可以緩解專家激活率低的問題并大幅提升更大專家的使用率，如圖 1a 所示實現了 90.71% 的激活率，這能讓模型能力獲得更高效的擴展。
2. 具有更細粒度的理解能力。MH-MoE 采用的多頭機制會將子 token 分配給不同的專家，從而可以聯合關注來自不同專家的不同表征空間的信息，最終獲得更好更細粒度的理解能力。舉個例子，如圖 1b 的明亮區域所示，子 token 會被分配給更多樣化的一組專家，這有助于捕獲語義豐富的信息。
3. 可實現無縫整合。MH-MoE 實現起來非常簡單，而且與其它 SMoE 優化方法（如 GShard）無關，反而可以將它們整合起來一起使用以獲得更好的性能。

方法

圖 3 給出了 MH-MoE 的整體架構，其使用了多頭機制將每個 token 分拆為子 token，然后將這些子 token 路由給不同的專家。

多頭混合專家

為了能清楚說明，這里僅描述單層 MH-MoE。

首先，通過一個多頭層將輸入 token 序列投射成一個新序列。

之后，沿 token 維度將新序列中的每個 token 分拆為多個子 token，并根據原始 token 序列并行排布這些子 token，進而構成一個新的特征空間。

然后將所有這些子 token 輸送給一個門控函數。將特定子 token 路由到第 p 個專家的門控值的計算方式為：

對于路由方法，這篇論文關注的重點方法是 top-k 路由，也就是激活路由分數最大的 k 個專家。然后讓這些激活的專家處理子 token。

之后，按子 token 原來的順序重新排布并整合所得結果。

然后，通過一個 token 合并操作將所得整合結果轉換回原始 token 形式。

最后，使用一個融合層將轉換后的結果投射成多個特征的有效整合形式，此時這些特征已捕獲了不同專家表征空間的詳細信息。這樣便可得到單層 MH-MoE 的最終輸出。

訓練目標

MH-MoE 的訓練目標是最小化兩個損失：針對具體任務的損失和輔助性的負載平衡損失。

實驗

實驗設置

為了進行比較，該研究的實驗采用了兩種基準模型：(1) Dense，這是沒有整合稀疏激活的并行模塊（SMoE 層）的 Transformer 解碼器。(2) X-MoE，基于 Chi et al. (2022) 的論文《On the representation collapse of sparse mixture of experts》提出的方法的實現。

實驗中的 MH-MoE 基于 X-MoE 并使用了與其一樣的設置。

實驗任務有三個：以英語為中心的語言建模、多語言語言建模、掩碼式多模態建模。

更多有關數據集和模型架構的設置請參閱原論文。

困惑度評估

他們在兩種專家設置（8 個專家和 32 個專家）下研究了所有預訓練模型和預訓練任務的驗證困惑度曲線。圖 4 給出了困惑度趨勢，表 1 是最終的困惑值。

據此可以看出：

1. 相比于基準，MH-MoE 的困惑度總是更低，這說明其能更有效地學習；
2. 在三個不同的設置中，MH-MoE 的困惑度是最低的；
3. 當專家數量增多時，MH-MoE 的困惑度會下降，這說明隨著專家數量增多，其表征學習能力會提升，模型也能從中受益。

這些結果表明 MH-MoE 在多種預訓練范式下都有更優的學習效率和語言表征能力。

下游任務評估

為了驗證 MH-MoE 的效果，該團隊也為每個預訓練任務執行了對應的下游任務評估。

以英語為中心的語言建模

這里使用了 9 個不同的零樣本評估基準，可以評估模型解決多種不同自然語言任務的能力，比如常識推理、一般語言理解和知識理解。評估框架為 LLM Evaluation Harness。結果見表 2。

可以看到，相比于 Dense 模型，X-MoE 有明顯優勢，這說明較大的模型能讓 SMoE 模型（如 X-MoE）受益。總體而言，MH-MoE 在所有基準上都表現最佳。

多語言語言建模

他們在跨語言自然語言推理（XNLI）語料庫（14 種語言）上評估了新的多語言語言模型。評估框架依然是 LLM Evaluation Harness，同樣使用了零樣本設置。結果見表 3。

MH-MoE 依然表現最佳，這體現了多頭機制在建模跨語言自然語言方面的有效性。

掩碼式多模態建模

他們也在社區廣泛使用的視覺 - 語言理解和生成基準上執行了評估，包括視覺問答、視覺推理和圖像描述。評估結果見表 4。

可以看到，MH-MoE 在這三個任務上有著全面的優勢。這些結果表明 MH-MoE 具有更強的視覺信息理解能力，這也驗證了新提出的多頭機制在捕獲視覺數據中的不同語義和詳細信息方面的有效性。

消融研究

為了驗證 MH-MoE 各組件和參數的效果，該團隊也進行了消融研究。他們研究的內容包括頭的數量、多層感知器層（包括多頭層和融合層）、token 拆分與融合操作、MLP 層的數量。

表 5、6、7 給出了研究結果。整體而言，MH-MoE 各組件的效果得到了驗證，并且他們也得到了一些有趣的結果，比如從表 7 可以看出單層 MLP 足以實現 token 分割和融合。

分析

專家激活分析

最后該團隊還通過可視化分析等方法對 MH-MoE 進行了分析。

圖 5 給出了 X-MoE 和 MH-MoE 中專家激活的分布情況。

可以看到，MH-MoE 的專家激活率明顯更高，并且隨著頭的數量 h 增大，專家激活的頻率也會上升。

圖 6 則對比了 X-MoE 和 MH-MoE 的可擴展性（專家數量從 8 擴展到 256）。

可以看到 MH-MoE 的優勢非常明顯，并且 X-MoE 的下游性能會在專家數為 64 時達到飽和，而 MH-MoE 卻還能繼續提升。

分析細粒度理解能力

為了進一步分析多頭機制對 MH-MoE 的幫助，該團隊更深入地分析了其理解多樣且復雜的語義信息的能力，比如理解語言中的多義詞和錯誤同源詞（記為 PF token）以及圖像中的信息豐富的區域。

對于語言數據，他們計算和比較了從 PF token 和非 PF token 拆分出的子 token 的散度層級（即這些子 token 路由到的不同專家的數量）。結果見圖 7。

可以看到相比于非 PF token，PF token 的散度分布明顯靠右。這說明，在 MH-MoE 的推理過程中，PF token 會將其子 token 路由到更多不同專家，從而會捕獲到與非 PF token 不同的語義信息，實現更好的多義詞和錯誤同源詞建模。

對于圖像數據，他們分析的是不同圖塊的散度層級在訓練過程中的變化情況，結果見圖 8。

有趣的是，可以看到隨著訓練步驟增多，高頻紋理區域（即有豐富語義信息的區域）的散度層級會逐漸增大，而低頻紋理區域的散度層級則會逐漸降低。這表明在訓練過程中，MH-MoE 傾向于將具有復雜紋理的區域的 token 路由到更多不同專家，由此可讓模型對該區域的語義有更細粒度的理解。

該團隊也執行了復雜性和參數分析，詳見原論文。