新智元報道
編輯：LRST
【新智元導讀】華中科技大學的研究人員提出了MoE Jetpack框架，通Checkpoint Recycling方法和SpheroMoE結構，將密集啟用模型的預訓練權重微調為混合專家（MoE）模型，從而免去了MoE模型的預訓練過程，大幅提升了MoE在下游任務中的精度和收斂速度。
混合專家模型（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的廣泛研究與應用提供了有力支援。
參考資料：
https://arxiv.org/abs/2406.04801
