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

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

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

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

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

近日，微軟研究院和清華大學(xué)提出了多頭混合專家（MH-MoE）。顧名思義，MH-MoE 采用了多頭機(jī)制，可將每個(gè)輸入 token 分成多個(gè)子 token。然后將這些子 token 分配給一組多樣化的專家并行處理，之后再無(wú)縫地將它們整合進(jìn)原來(lái)的 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)輸入單個(gè) token 時(shí)，MH-MoE 會(huì)將其分成 4 個(gè)子 token，進(jìn)而激活 4 個(gè)專家，而 SMoE 僅激活 1 個(gè)專家。

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

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

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

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

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

方法

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

多頭混合專家

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

首先，通過(guò)一個(gè)多頭層將輸入 token 序列投射成一個(gè)新序列。

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

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

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

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

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

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

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

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

實(shí)驗(yàn)

實(shí)驗(yàn)設(shè)置

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

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

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

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

困惑度評(píng)估

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

據(jù)此可以看出：

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

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

下游任務(wù)評(píng)估

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

以英語(yǔ)為中心的語(yǔ)言建模

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

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

多語(yǔ)言語(yǔ)言建模

他們?cè)诳缯Z(yǔ)言自然語(yǔ)言推理（XNLI）語(yǔ)料庫(kù)（14 種語(yǔ)言）上評(píng)估了新的多語(yǔ)言語(yǔ)言模型。評(píng)估框架依然是 LLM Evaluation Harness，同樣使用了零樣本設(shè)置。結(jié)果見(jiàn)表 3。

MH-MoE 依然表現(xiàn)最佳，這體現(xiàn)了多頭機(jī)制在建?？缯Z(yǔ)言自然語(yǔ)言方面的有效性。

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

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

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

消融研究

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

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

分析

專家激活分析

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

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

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

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

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

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

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

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

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

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

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

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