微軟讓MoE長出多個頭，大幅提升專家激活率
機器之心報道
編輯：Panda W
混合專家（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 形式。
圖 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 的具有以下優勢：
方法
圖 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 的效果，該團隊也爲每個預訓練任務執行了對應的下游任務評估。
以英語爲中心的語言建模
這裡使用了 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 路由到更多不同專家，由此可讓模型對該區域的語義有更細粒度的理解。
該團隊也執行了複雜性和參數分析，詳見原論文。