混合專家模型（MoE, Mixture of Experts）是一種通過動態激活網絡的部分結構來提升計算效率的架構，可以在保持相對穩定的計算成本的前提下大幅增加參數量，從而有效提升模型性能。

這一特性使得MoE能夠兼顧模型的規模與效率，已廣泛應用于各種大規模任務。

然而，MoE模型通常需要在大型數據集上預訓練以獲得理想性能，導致其對時間和計算資源的需求極高，這也限制了其在深度學習社區中的普及性。

為解決這一問題，華中科技大學的研究人員提出了MoE Jetpack框架，利用密集模型的預訓練權重（Dense checkpoints）來微調出視覺混合專家模型（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框架的核心創新包括：

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

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

實驗結果表明，MoE Jetpack在多個數據集和網絡結構上實現了顯著的性能提升。在ImageNet-1K上，模型收斂速度提升2倍，準確率提高了2.8%；在小規模數據集上，收斂速度可達8倍提升，準確率提升超過30%。

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

在多個下游數據集上的實驗表明，該框架能夠高效利用預訓練權重，實現更快的收斂速度和更優的性能表現。

研究方法

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

圖2 Checkpoint Recycling和SpheroMoE結構

Checkpoint Recycling：作為MoE Jetpack的基礎階段，Checkpoint Recycling通過將預訓練的密集模型權重轉換為高質量的MoE初始化權重，使新模型在性能和收斂速度上都得以提升。

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

本文比較了四種主要的權重回收策略：

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

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

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

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

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

SpheroMoE層

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

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

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

圖3 自適應雙路徑MoE

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

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

這些設計使MoE Jetpack在下游任務微調中不僅具備了更快的收斂速度，還實現了顯著的性能提升。

實驗結果

本文在 ViT 和 ConvNeXt 兩種典型網絡結構以及八個圖像分類任務上進行了廣泛實驗。實驗結果表明，MoE Jetpack 在性能上顯著優于隨機初始化的 Soft MoE 模型，并且利用密集權重的微調效果明顯超過直接微調經過預訓練的密集模型。

表1 MoE Jetpack基于ViT和ConvNeXt在8個下游數據集上的性能表現

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

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

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

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

專家注意力圖展示了不同的專家關注圖像的不同區域，各司其職。專家貢獻分布圖表明，核心專家和普通專家在不同層次的貢獻差異顯著，展示了模型的自適應路由調度機制。

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

總結

MoE Jetpack是一個創新框架，旨在將預訓練的密集模型權重高效轉換為MoE模型。通過提出Checkpoint Recycling技術，MoE Jetpack能夠有效繼承密集模型的知識；引入的SpheroMoE 層，顯著提升微調過程的穩定性和性能。

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