收斂速度最高8倍，準確率提升超30%！華科發布MoE Jetpack框架 | NeurIPS 2024

作者：新智元 2024-11-21 10:21:06

華中科技大學的研究人員提出了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路由將輸入劃分為高重要性和低重要性兩類，并引導其進入不同計算路徑：高重要性輸入分配至包含更大參數量的核心專家；低重要性輸入則進入包含較小專家的通用路徑。這樣的雙路徑結構通過劃分細粒度的專家增加了專家的數量，優化了資源利用，提升了模型的性能與計算效率。