MH-MoE 能最佳化幾乎所有專家，實現起來非常簡單。
混合專家（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 的具有以下優勢：
專家啟用率更高且擴充套件性更好。MH-MoE 能最佳化幾乎所有專家，從而可以緩解專家啟用率低的問題並大幅提升更大專家的使用率，如圖 1a 所示實現了 90.71% 的啟用率，這能讓模型能力獲得更高效的擴充套件。
具有更細粒度的理解能力。MH-MoE 採用的多頭機制會將子 token 分配給不同的專家，從而可以聯合關注來自不同專家的不同表徵空間的資訊，最終獲得更好更細粒度的理解能力。舉個例子，如圖 1b 的明亮區域所示，子 token 會被分配給更多樣化的一組專家，這有助於捕獲語義豐富的資訊。
可實現無縫整合。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 是最終的困惑值。
據此可以看出：
相比於基準，MH-MoE 的困惑度總是更低，這說明其能更有效地學習；
在三個不同的設定中，MH-MoE 的困惑度是最低的；
當專家數量增多時，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 路由到更多不同專家，由此可讓模型對該區域的語義有更細粒度的理解。
該團隊也執行了複雜性和引數分析，詳見原論文。