[論文レビュー] One Student Knows All Experts Know: From Sparse to Dense

Fuzhao Xue, Xiaoxin He+3 more|arXiv (Cornell University)|Jan 26, 2022

COVID-19 diagnosis using AI被引用数 5

ひとこと要約

本論文では、スパースなMixture-of-Experts（MoE）モデルから知識を蒸留する新しいフレームワーク「知識統合」を提案する。2段階の訓練プロセスを用いる：まずSVD-KGまたはTop-KGによる知識収集を行い、次に知識蒸留を実行する。OneSは1500万パラメータでImageNetのトップ1精度78.4％を達成し、MoEの61.7％の利点を維持し、MoEと比較して3.7倍高速な推論を実現。NLPタスクではベースラインを最大51.7％上回る性能を発揮する。

ABSTRACT

Human education system trains one student by multiple experts. Mixture-of-experts (MoE) is a powerful sparse architecture including multiple experts. However, sparse MoE model is easy to overfit, hard to deploy, and not hardware-friendly for practitioners. In this work, inspired by the human education model, we propose a novel task, knowledge integration, to obtain a dense student model (OneS) as knowledgeable as one sparse MoE. We investigate this task by proposing a general training framework including knowledge gathering and knowledge distillation. Specifically, to gather key knowledge from different pre-trained experts, we first investigate four different possible knowledge gathering methods, \ie summation, averaging, Top-K Knowledge Gathering (Top-KG), and Singular Value Decomposition Knowledge Gathering (SVD-KG) proposed in this paper. We then refine the dense student model by knowledge distillation to offset the noise from gathering. On ImageNet, our OneS preserves $61.7\%$ benefits from MoE and achieves $78.4\%$ top-1 accuracy ImageNet with only $15$M parameters. On four natural language processing datasets, OneS obtains $88.2\%$ MoE benefits and outperforms the best baseline by $51.7\%$ using the same architecture and training data. In addition, compared with the MoE counterpart, OneS can achieve $3.7 imes$ inference speedup due to less computation and hardware-friendly architecture.

研究の動機と目的

• スパースなMoEモデルの実用的限界（過学習、導入困難、ハードウェア効率の悪さ）を解消するため、その知識を密度型で展開可能な学生モデルに移転すること。
• 1つの密度型モデルが複数の事前学習済みエキスパートから学ぶという、人間の教育モデル（複数の教師から学ぶ）を模倣する、新たなタスク「知識統合」を提案すること。
• 知識収集と知識蒸留を組み合わせた一般化可能な訓練フレームワークを開発し、事前学習済みMoEから密度型学生モデルを有効に初期化・最適化すること。
• 得られた密度型モデル（OneS）がMoEの性能を同等または上回りつつ、はるかに効率的で導入が容易であることを実証すること。

提案手法

• 知識収集は4つの方法で実施：和、平均、Top-K知識収集（Top-KG）、および新規の特異値分解知識収集（SVD-KG）。これらによりMoEエキスパートから重要な知識を抽出・統合する。
• SVD-KGはエキスパートの重みをSVDで分解し、最も重要な特異成分を保持。これらの成分を密度型学生モデルの前向き伝搬ネットワーク（FFN）層の初期化に使用する。
• Top-KGはエキスパート全体で最も活性化度の高いk個のニューロンまたは重み成分を選択し、学生モデルの初期化に集中して高影響度の知識を反映させる。
• 知識蒸留は、完全なMoE教師モデルのソフトラベルを用いて実施。知識収集段階で生じるノイズを低減し、一般化性能を向上させる。
• フレームワークは2段階の訓練プロセスを採用：まず知識収集により学生モデルを初期化し、次にソフトラベルと交差エントロピー損失を用いた知識蒸留により微調整する。
• 学生モデルのアーキテクチャは標準的な密度型モデルと同一であるため、ハードウェア互換性が保たれ、効率的な推論が可能となる。

実験結果

リサーチクエスチョン

• RQ1 事前学習済みMoEからの知識を用いて、性能を維持しつつ導入性を向上させた密度型学生モデルを効果的に初期化できるか？
• RQ2 和、平均、Top-KG、SVD-KGのうち、どの知識収集手法がMoEエキスパートから密度型学生モデルを初期化する際に最も優れた性能を発揮するか？
• RQ3 知識収集後、知識蒸留が学生モデルの性能をどの程度向上できるか、特に初期化段階のノイズ低減に寄与するか？
• RQ4 最終的な密度型学生モデルは、精度、パラメータ効率、推論速度の観点でMoEと比較してどの程度の差異を示すか？
• RQ5 知識統合型密度モデルは、視覚およびNLPベンチマークにおいて、既存の蒸留法およびスパースモデルベースラインを上回る性能を発揮するか？

主な発見

• ImageNetでは、OneSは1500万パラメータでトップ1精度78.4％を達成し、MoEの61.7％の性能利点を維持した。
• 4つのNLPベンチマークでは、OneSはパラメータ数がMoEの46％にとどまりながら、MoEの88.2％の性能利点を達成し、最良のベースラインを51.7％上回った。
• 計算量の削減とハードウェアに優しい密度型アーキテクチャのおかげで、OneSはMoEの3.7倍高速な推論を実現した。
• アブレーションスタディにより、知識収集と知識蒸留の両方が不可欠であることが確認された。両方を省くと顕著な性能低下が生じ、特に知識収集が初期段階での性能向上に最も寄与した。
• 同等の訓練エポック数で、OneSはベースラインを上回り、さらに訓練を延長しても性能が向上を続ける一方、MoEベースのモデル（例：WideNet）は性能が頭打ちになる傾向を示した。
• SQuAD1.1 や SST-2 のような小規模データセットでも、OneSはMoE教師モデルを上回った。これは、密度型で汎化性の高い学生構造が過学習を軽減している可能性を示唆している。