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

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

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

具體來說，就是在優(yōu)化時只有一小部分專家會被激活，如圖 1a 所示（8.33% 的激活率），這會導(dǎo)致在學(xué)習(xí)應(yīng)對復(fù)雜任務(wù)的大量專家時，會出現(xiàn)性能次優(yōu)和效果不佳的問題。

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

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

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

圖 2 展示了 MH-MoE 的工作流程?？梢钥吹剑?dāng)輸入單個 token 時，MH-MoE 會將其分成 4 個子 token，進(jìn)而激活 4 個專家，而 SMoE 僅激活 1 個專家。

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

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

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

MH-MoE 的具有以下優(yōu)勢：

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

方法

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

多頭混合專家

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

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

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

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

對于路由方法，這篇論文關(guān)注的重點方法是 top-k 路由，也就是激活路由分?jǐn)?shù)最大的 k 個專家。然后讓這些激活的專家處理子 token。

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

然后，通過一個 token 合并操作將所得整合結(jié)果轉(zhuǎn)換回原始 token 形式。

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

訓(xùn)練目標(biāo)

MH-MoE 的訓(xùn)練目標(biāo)是最小化兩個損失：針對具體任務(wù)的損失和輔助性的負(fù)載平衡損失。

實驗

實驗設(shè)置

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

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

實驗任務(wù)有三個：以英語為中心的語言建模、多語言語言建模、掩碼式多模態(tài)建模。

更多有關(guān)數(shù)據(jù)集和模型架構(gòu)的設(shè)置請參閱原論文。

困惑度評估

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

據(jù)此可以看出：

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

這些結(jié)果表明 MH-MoE 在多種預(yù)訓(xùn)練范式下都有更優(yōu)的學(xué)習(xí)效率和語言表征能力。

下游任務(wù)評估

為了驗證 MH-MoE 的效果，該團隊也為每個預(yù)訓(xùn)練任務(wù)執(zhí)行了對應(yīng)的下游任務(wù)評估。

以英語為中心的語言建模

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

可以看到，相比于 Dense 模型，X-MoE 有明顯優(yōu)勢，這說明較大的模型能讓 SMoE 模型（如 X-MoE）受益?？傮w而言，MH-MoE 在所有基準(zhǔn)上都表現(xiàn)最佳。

多語言語言建模

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

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

掩碼式多模態(tài)建模

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

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

消融研究

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

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

分析

專家激活分析

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

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

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

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

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

分析細(xì)粒度理解能力

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

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

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

對于圖像數(shù)據(jù)，他們分析的是不同圖塊的散度層級在訓(xùn)練過程中的變化情況，結(jié)果見圖 8。

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

該團隊也執(zhí)行了復(fù)雜性和參數(shù)分析，詳見原論文。