マイクロソフトアジアリサーチは、知識蒸留を使用して小さなViTを改善するTinyMIMを提案

1. 研究の動機

マスクモデリング (MIM、MAE) は、非常に効果的な自己教師ありトレーニング方法であることが示されています。ただし、図 1 に示すように、MIM は大規模なモデルでは比較的うまく機能します。モデルが非常に小さい場合 (ViT-T 5M パラメータなど、このようなモデルは現実世界では非常に重要です)、MIM によってモデルの有効性がある程度低下する可能性もあります。たとえば、ImageNet で MAE を使ってトレーニングした ViT-L の分類効果は、ImageNet で通常の教師を使ってトレーニングしたモデルよりも 3.3% 優れていますが、MAE を使ってトレーニングした ViT-T の分類効果は、ImageNet で通常の教師を使ってトレーニングしたモデルよりも 0.6% 低くなります。

本研究では、ViT 構造を変更せずに、他の誘導バイアスを導入するために構造を変更せずに、蒸留を使用して大規模モデルから小規模モデルに知識を転送する TinyMIM を提案しました。

• 論文アドレス: https://arxiv.org/pdf/2301.01296.pdf
• コードアドレス: https://github.com/OliverRensu/TinyMIM

蒸留の目的、データ拡張、正則化、補助損失関数などが蒸留に与える影響を体系的に研究しました。厳密に言えば、ImageNet-1K のみをトレーニング データとして使用し (ImageNet-1K のみでトレーニングされた教師モデルを含む)、ViT-B をモデルとして使用する場合、当社の方法は現在最高のパフォーマンスを達成しています。図に示すように:

私たちの方法 (TinyMIM) を、マスクベースの再構成方法 MAE および最初からトレーニングされた教師あり学習方法 DeiT と比較します。 MAE は、モデルが大きい場合にはパフォーマンスが大幅に向上しますが、モデルが小さい場合には向上が制限され、モデルの最終的な効果に悪影響を与える可能性もあります。私たちの方法 TinyMIM は、さまざまなモデル サイズで大幅な改善を実現します。

私たちの貢献は次のとおりです。

1. 蒸留ターゲット: 1) クラス トークンまたは特徴マップのみを蒸留するよりも、トークン間の関係を蒸留する方が効果的です。2) 中間層を蒸留ターゲットとして使用する方が効果的です。

2. データ拡張とネットワーク正規化: 1) マスクされた画像では効果が悪くなります。2) 生徒モデルではドロップ パスが少し必要ですが、教師モデルでは必要ありません。

3. 補助損失: MIM は補助損失関数としては意味がありません。

4. マクロ蒸留戦略: 順次蒸留 (ViT-B -> ViT-S -> ViT-T) が最も効果的であることがわかりました。

2. 方法

蒸留対象、入力画像、蒸留対象モジュールを体系的に調査します。

2.1 蒸留効果に影響を与える要因

1) 特徴:

a. 中間ブロック機能と出力機能

i=L の場合、Transformer 出力層の機能を参照します。 i < L の場合、Transformer の中間層の機能を参照します。

b. 注意特徴とフィードフォワード層（FFN）特徴

Transformer の各ブロックには、Attention レイヤーと FFN レイヤーがあります。異なるレイヤーを抽出すると、異なる効果が得られます。

c. QKVの機能

Attention 層には Q、K、V の特徴があります。これらの特徴は、Attention メカニズムを計算するために使用されます。また、これらの特徴を直接抽出することも調査しました。

2) 関係

Q、K、Vは注目度マップを計算するために使用され、これらの特徴間の関係も知識蒸留の対象として使用できます。

3) 入力: マスクを含めるかどうか

伝統的な知識の蒸留により、完全なイメージが直接入力されます。私たちの方法は、蒸留されたマスク モデリング モデルを探索することを目的としているため、マスクされた画像が知識蒸留の入力として適切かどうかも探索します。

2.2 知識蒸留手法の比較

1) クラストークンの蒸留:

最も簡単な方法は、DeiT に似た MAE 事前トレーニング済みモデルのクラス トークンを直接抽出することです。

ここで、 は学生モデルのクラス トークンを参照し、 は教師モデルのクラス トークンを参照します。

2) 特徴蒸留:比較のために特徴蒸留[1]を直接参照する。

3) 関係蒸留: この論文でもデフォルトとなっている蒸留戦略を提案する。

3. 実験

3.1 主な実験結果

私たちの方法は ImageNet-1K で事前トレーニングされており、教師モデルも ImageNet-1K で事前トレーニングされています。次に、下流のタスク（分類、セマンティックセグメンテーション）で事前トレーニング済みのモデルを微調整しました。モデルのパフォーマンスを図に示します。

私たちの方法は、特に小規模なモデルの場合、従来の MAE ベースの方法よりも大幅に優れています。具体的には、超小型モデル ViT-T の場合、当社の方法は 75.8% の分類精度を達成し、これは MAE ベースライン モデルよりも 4.2 の改善となります。小型モデル ViT-S では、83.0% の分類精度を達成し、以前の最良の方法よりも 1.4 向上しました。ベース サイズ モデルの場合、当社の方法は MAE ベースライン モデルと以前の最良モデル CAE 4.1 および 2.0 をそれぞれ上回ります。

同時に、図に示すように、モデルの堅牢性もテストしました。

MAE-B と比較すると、TinyMIM-B は ImageNet-A と ImageNet-R でそれぞれ +6.4 と +4.6 向上します。

3.2 アブレーション実験

1) さまざまな関係を精査する

同時に、QK と VV の関係が抽出され、関係を計算するときに Softmax を使用して最良の効果が得られます。

2) 異なる蒸留戦略

TinyMIM の関係抽出アプローチは、さまざまなサイズのモデルにおいて、MAE ベースライン モデル、クラス トークン抽出、および機能マップ抽出よりも優れた結果を実現します。

3) 蒸留中間層

18 番目の層を蒸留すると最良の結果が得られることがわかりました。

IV. 結論

本稿では、小さなモデルがマスク再構成モデ​​リング (MIM) の事前トレーニングのメリットを享受できるようにする最初のモデルである TinyMIM を提案します。マスク再構築をタスクとして採用する代わりに、知識蒸留方式で大規模モデルの関係をシミュレートするようにトレーニングすることで、小規模モデルを事前トレーニングします。 TinyMIM の成功は、蒸留ターゲット、蒸留入力、中間層など、TinyMIM の事前トレーニングに影響を与える可能性のあるさまざまな要素を包括的に研究した結果です。広範囲にわたる実験を通じて、関係蒸留は特徴蒸留やクラスラベル蒸留などよりも優れているという結論に達しました。私たちの方法はシンプルでパフォーマンスに優れているため、将来の研究に強固な基盤を提供できると期待しています。