[論文レビュー] Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy

Pingzhi Li, Zhenyu Zhang+5 more|arXiv (Cornell University)|Oct 2, 2023

Domain Adaptation and Few-Shot LearningComputer Science被引用数 3

ひとこと要約

本稿では、ルーティング方針に従ったエキスパート統合によるスパースMixture-of-Experts（SMoE）の最初の統合、次に低ランクおよび構造的スパarsityを用いた統合エキスパートの圧縮を実行する、MC-SMoEという新規フレームワークを提案する。活性化周波数とルーティング類似度を活用することで、M-SMoEは冗長性とメモリ使用量を低減し、8つのベンチマークで性能の著しい低下なしに最大80%のメモリ削減と20%のFLOPs削減を達成する。

ABSTRACT

Sparsely activated Mixture-of-Experts (SMoE) has shown promise to scale up the learning capacity of neural networks, however, they have issues like (a) High Memory Usage, due to duplication of the network layers into multiple copies as experts; and (b) Redundancy in Experts, as common learning-based routing policies suffer from representational collapse. Therefore, vanilla SMoE models are memory inefficient and non-scalable, especially for resource-constrained downstream scenarios. In this paper, we ask: Can we craft a compact SMoE model by consolidating expert information? What is the best recipe to merge multiple experts into fewer but more knowledgeable experts? Our pilot investigation reveals that conventional model merging methods fail to be effective in such expert merging for SMoE. The potential reasons are: (1) redundant information overshadows critical experts; (2) appropriate neuron permutation for each expert is missing to bring all of them in alignment. To address this, we propose M-SMoE, which leverages routing statistics to guide expert merging. Specifically, it starts with neuron permutation alignment for experts; then, dominant experts and their "group members" are formed; lastly, every expert group is merged into a single expert by utilizing each expert's activation frequency as their weight for merging, thus diminishing the impact of insignificant experts. Moreover, we observed that our proposed merging promotes a low dimensionality in the merged expert's weight space, naturally paving the way for additional compression. Hence, our final method, MC-SMoE (i.e., Merge, then Compress SMoE), further decomposes the merged experts into low-rank and structural sparse alternatives. Extensive experiments across 8 benchmarks validate the effectiveness of MC-SMoE. For instance, our MC-SMoE achieves up to 80% memory and a 20% FLOPs reduction, with virtually no loss in performance.

研究の動機と目的

• リソース制約のあるデプロイメント環境において、特に高いメモリ使用量とエキスパートの冗長性を有するスパースMixture-of-Experts（SMoE）モデルに対処すること。
• 性能の低下なしに、エキスパート情報がより少ないがより知識豊富なエキスパートに効果的に統合可能かどうかを調査すること。
• 従来のモデル統合技術よりもSMoE設定において優れた性能を示す統合戦略を開発すること。
• 統合されたエキスパートが本質的により圧縮可能であるかどうかを調査し、さらなる効率向上を可能にするかを検討すること。

提案手法

• M-SMoEはルーティング統計を用いてエキスパート統合をガイドする：まず、特徴の対応を保証するため、置換を用いてエキスパート間のニューロンを整列させる。
• ルーティング方針の類似度に基づいて、支配的エキスパートとその「グループメンバー」を特定することで、エキスパートグループを形成する。
• 各グループ内のエキスパートは、活性化周波数を重みとして平均化することで統合され、重要度の低いエキスパートの影響を低減する。
• 統合されたエキスパートは、より低い次元の重み空間に位置することが判明し、さらなる圧縮が可能である。
• MC-SMoEは統合エキスパートに低ランクおよび構造的スパarsityを適用し、追加のパラメータおよびFLOPs削減を達成する。
• 本手法は、スイッチベース-32 SMoEモデルを用いて8つのNLPベンチマークでエンドツーエンドで評価される。

実験結果

リサーチクエスチョン

• RQ1 SMoEにおけるエキスパート情報は、性能の損失なしに、より少ないがより知識豊富なエキスパートに効果的に統合可能か？
• RQ2 従来のモデル統合手法がSMoE設定で失敗する理由は何か？主な失敗モードは何か？
• RQ3 ルーティング方針をどのように活用してエキスパート統合をガイドし、冗長エキスパートと支配的エキスパートを特定できるか？
• RQ4 エキスパートの統合は、自然に低次元の重み空間を形成するのか？その結果、さらなる圧縮が可能になるか？
• RQ5 統合と圧縮の両者が、メモリおよびFLOPs効率性を達成するために果たす相対的な貢献度は何か？

主な発見

• MC-SMoEは、スイッチベース-32 SMoEモデルにおいて、8つのベンチマークで最小限の性能低下で最大80%のメモリ削減と20%のFLOPs削減を達成する。
• M-SMoEは、すべてのタスクでC-SMoE（圧縮のみ）を上回り、性能維持の観点から、統合が圧縮単体よりも優れていることを示している。
• 周波数重み付き統合は、一様重みおよびフィッシャー重み付き統合よりも一貫して優れており、タスク平均で1.5–2%の精度向上を達成する。
• M-SMoEによる統合エキスパートは、より低い内因的次元性を示し、低ランクおよび構造的スパarsityによる効果的な後続圧縮が可能である。
• 後続のSMoEレイヤーでは、ヒートマップ可視化における残存パラメータの分布から、より高い圧縮可能性が示されている。
• アブレーションスタディにより、周波数に配慮した統合が、冗長エキスパートを抑制し、重要な知識が覆い隠されるのを防いでいることが確認された。