混合專家模型（MoE, Mixture of Experts）是一種通過動態(tài)激活網(wǎng)絡(luò)的部分結(jié)構(gòu)來提升計算效率的架構(gòu)，可以在保持相對穩(wěn)定的計算成本的前提下大幅增加參數(shù)量，從而有效提升模型性能。

這一特性使得MoE能夠兼顧模型的規(guī)模與效率，已廣泛應(yīng)用于各種大規(guī)模任務(wù)。

然而，MoE模型通常需要在大型數(shù)據(jù)集上預(yù)訓(xùn)練以獲得理想性能，導(dǎo)致其對時間和計算資源的需求極高，這也限制了其在深度學(xué)習(xí)社區(qū)中的普及性。

為解決這一問題，華中科技大學(xué)的研究人員提出了MoE Jetpack框架，利用密集模型的預(yù)訓(xùn)練權(quán)重（Dense checkpoints）來微調(diào)出視覺混合專家模型（MoE，Mixture of Experts）。

目前，這項工作已被NeurIPS 2024接收。

論文標題：MoE Jetpack: From Dense Checkpoints to Adaptive Mixture of Experts for Vision Tasks

論文地址： https://arxiv.org/abs/2406.04801

代碼地址： https://github.com/Adlith/MoE-Jetpack

MoE Jetpack框架的核心創(chuàng)新包括：

1. Checkpoint recycling：通過采樣密集模型權(quán)重產(chǎn)生差異化的專家，組成MoE模型的初始化權(quán)重，從而加速模型收斂、提升性能，并避免大規(guī)模的MoE模型預(yù)訓(xùn)練。

2. SpheroMoE Layer：通過調(diào)整MoE結(jié)構(gòu)，利用交叉注意力機制進行專家分配，將query和key投影到超球空間以提升微調(diào)過程的穩(wěn)定性，并通過一系列專家正則化方法有效緩解MoE模型微調(diào)過程中的過擬合現(xiàn)象。

實驗結(jié)果表明，MoE Jetpack在多個數(shù)據(jù)集和網(wǎng)絡(luò)結(jié)構(gòu)上實現(xiàn)了顯著的性能提升。在ImageNet-1K上，模型收斂速度提升2倍，準確率提高了2.8%；在小規(guī)模數(shù)據(jù)集上，收斂速度可達8倍提升，準確率提升超過30%。

圖1 （a）MoE Jetpack將密集預(yù)訓(xùn)練權(quán)重轉(zhuǎn)化為MoE模型的初始化權(quán)重，在性能提升的同時保持等效的FLOPs。(b) 未預(yù)訓(xùn)練的ViT、微調(diào)的ViT、未預(yù)訓(xùn)練的Soft MoE 與MoE Jetpack在多個視覺數(shù)據(jù)集上的性能比較。

在多個下游數(shù)據(jù)集上的實驗表明，該框架能夠高效利用預(yù)訓(xùn)練權(quán)重，實現(xiàn)更快的收斂速度和更優(yōu)的性能表現(xiàn)。

研究方法

MoE Jetpack的核心由兩個階段組成：Checkpoint Recycling（用于MoE模型的初始化）和SpheroMoE層（用于微調(diào)MoE模型），如下圖所示。

圖2 Checkpoint Recycling和SpheroMoE結(jié)構(gòu)

Checkpoint Recycling：作為MoE Jetpack的基礎(chǔ)階段，Checkpoint Recycling通過將預(yù)訓(xùn)練的密集模型權(quán)重轉(zhuǎn)換為高質(zhì)量的MoE初始化權(quán)重，使新模型在性能和收斂速度上都得以提升。

具體地，Checkpoint Recycling從密集權(quán)重的多層感知器（MLP）中采樣出部分權(quán)重構(gòu)建專家層，以確保專家的多樣性和靈活性。

本文比較了四種主要的權(quán)重回收策略：

1. 重要性采樣（Importance-Based Weight Sampling）：重要性采樣是 MoE Jetpack 默認的權(quán)重采樣方法，通過計算輸出特征的均值，選擇top-d' 個最重要的特征維度，同時根據(jù)隱藏單元的激活值大小進行獨立采樣。這種方法確保每個專家包含關(guān)鍵的特征和隱藏單元，有助于提升模型的初始化質(zhì)量和訓(xùn)練效率。

2. 圖劃分法（Co-Activation Graph Partitioning）：圖劃分法通過構(gòu)建共激活圖，將常一起激活的隱藏單元分組，使用圖劃分算法（如 Metis），將密集權(quán)重的隱藏單元劃分成多個子圖，組合不同的子圖形成不同專家層，確保專家專注于不同的功能區(qū)域，提升模型的特征表達能力。

3. 均勻采樣（Uniform Weight Selection）：均勻采樣在特征維度和隱藏單元上均勻選擇權(quán)重，保證每個 MoE 專家層均衡分布初始化權(quán)重。此方法實現(xiàn)簡單，但不考慮特征重要性，因此性能提升效果較為一般。

4. 隨機采樣（Random Weight Sampling）：隨機采樣在特征維度和隱藏單元中隨機抽取，生成專家層的初始化權(quán)重。該方法實現(xiàn)簡便，但由于沒有關(guān)注特征的重要性，性能通常較低。

Checkpoint Recycling引入的計算開銷幾乎可以忽略，同時顯著提升模型性能并與收斂速度。

SpheroMoE層

在MoE模型權(quán)重初始化后，SpheroMoE層進一步優(yōu)化微調(diào)過程。SpheroMoE 層的引入解決了 MoE 模型在優(yōu)化過程中面臨的數(shù)值不穩(wěn)定、專家過度專一等問題。它通過三種機制來提升模型在下游任務(wù)的性能和穩(wěn)定性：

1. 超球路由機制：利用cross attention結(jié)構(gòu)，將輸入動態(tài)分配給 MoE 模型的不同專家。這種機制首先對隨機初始化的專家查詢（Q, query）和輸入鍵（K, key）進行歸一化投影（L2 Norm）至超球空間，以確保數(shù)值穩(wěn)定性，并通過余弦相似度來選擇輸入對應(yīng)的專家。最終輸出由各個專家的結(jié)果組合而成，保證 MoE 模型的輸出特征和密集模型之間的分布一致性。

2. 自適應(yīng)雙路徑MoE：為提升計算效率，SpheroMoE路由將輸入劃分為高重要性和低重要性兩類，并引導(dǎo)其進入不同計算路徑：高重要性輸入分配至包含更大參數(shù)量的核心專家；低重要性輸入則進入包含較小專家的通用路徑。這樣的雙路徑結(jié)構(gòu)通過劃分細粒度的專家增加了專家的數(shù)量，優(yōu)化了資源利用，提升了模型的性能與計算效率。

圖3 自適應(yīng)雙路徑MoE

3. 專家正則化：為避免專家層過度專注于特定輸入或出現(xiàn)過度特化，本文引入可學(xué)習(xí)的軟溫度參數(shù)，用以調(diào)整softmax的平滑程度以精確控制輸入的分配和輸出的組合。

此外，使用專家隨機失活機制能有效防止模型對特定專家的依賴。

這些設(shè)計使MoE Jetpack在下游任務(wù)微調(diào)中不僅具備了更快的收斂速度，還實現(xiàn)了顯著的性能提升。

實驗結(jié)果

本文在 ViT 和 ConvNeXt 兩種典型網(wǎng)絡(luò)結(jié)構(gòu)以及八個圖像分類任務(wù)上進行了廣泛實驗。實驗結(jié)果表明，MoE Jetpack 在性能上顯著優(yōu)于隨機初始化的 Soft MoE 模型，并且利用密集權(quán)重的微調(diào)效果明顯超過直接微調(diào)經(jīng)過預(yù)訓(xùn)練的密集模型。

表1 MoE Jetpack基于ViT和ConvNeXt在8個下游數(shù)據(jù)集上的性能表現(xiàn)

此外，論文還對MoE Jetpack的多種配置進行了深入研究，系統(tǒng)分析了不同專家數(shù)量、不同原始網(wǎng)絡(luò)尺寸大小等因素對模型性能的影響。

表2 各種參數(shù)量的MoE Jetpack都展示出了顯著的性能提升

下圖展示了微調(diào)的全過程中MoE Jetpack對加速模型收斂速度和提升模型性能方面的效果，突顯了其作為 MoE 模型預(yù)訓(xùn)練替代方法的潛力。

圖4 MoE Jetpack帶來了收斂速度提升

專家注意力圖展示了不同的專家關(guān)注圖像的不同區(qū)域，各司其職。專家貢獻分布圖表明，核心專家和普通專家在不同層次的貢獻差異顯著，展示了模型的自適應(yīng)路由調(diào)度機制。

圖5 專家注意力圖和專家貢獻分布圖

總結(jié)

MoE Jetpack是一個創(chuàng)新框架，旨在將預(yù)訓(xùn)練的密集模型權(quán)重高效轉(zhuǎn)換為MoE模型。通過提出Checkpoint Recycling技術(shù)，MoE Jetpack能夠有效繼承密集模型的知識；引入的SpheroMoE 層，顯著提升微調(diào)過程的穩(wěn)定性和性能。

該框架不僅降低了MoE模型的訓(xùn)練成本和硬件要求，還減少了對環(huán)境的影響，使得研究者在普通計算資源下也能輕松使用混合專家模型，為MoE的廣泛研究與應(yīng)用提供了有力支持。