[論文レビュー] The Role of Masking for Efficient Supervised Knowledge Distillation of Vision Transformers

Seung‐Woo Son, Ryu, Jegwang+2 more|arXiv (Cornell University)|Feb 21, 2023

Advanced Neural Network Applications被引用数 5

ひとこと要約

MaskedKDは、教師の推論時に入力画像パッチの25–50%をマスクすることで、ビジョントランスフォーマーの知識蒸留の計算コストを低減する手法である。このマスクは、学生モデルの注意メカニズムから導出されたサリエンシー スコアに基づく。このアプローチにより、学生の精度を損なわずに知識蒸留に必要なFLOPsを最大50%削減でき、合計の訓練FLOPsを最大28%削減し、2.0×–3.7×のスループット向上を達成する。

ABSTRACT

Knowledge distillation is an effective method for training lightweight vision models. However, acquiring teacher supervision for training samples is often costly, especially from large-scale models like vision transformers (ViTs). In this paper, we develop a simple framework to reduce the supervision cost of ViT distillation: masking out a fraction of input tokens given to the teacher. By masking input tokens, one can skip the computations associated with the masked tokens without requiring any change to teacher parameters or architecture. We find that masking patches with the lowest student attention scores is highly effective, saving up to 50% of teacher FLOPs without any drop in student accuracy, while other masking criterion leads to suboptimal efficiency gains. Through in-depth analyses, we reveal that the student-guided masking provides a good curriculum to the student, making teacher supervision easier to follow during the early stage and challenging in the later stage.

研究の動機と目的

• 知識蒸留中に大規模なビジョントランスフォーマー（ViT）教師の高い計算コストに対処すること。
• 学生モデルの性能を低下させることなく、教師モデルを実行するために必要なFLOPs（蒸留コスト）を低減すること。
• 教師の予測を事前に計算して保存する必要を回避すること。これはメモリを食い、データオーグメンテーションの下で精度を低下させる要因となる。
• 自己教師あり学習や音声ベースのトランスフォーマー設定においても、効率的な蒸留を可能にすること。これは教師ありViT訓練にとどまらない拡張を目的としている。

提案手法

• MaskedKDは、学生モデルの最終マルチヘッドアテンション層から計算されるパッチサリエンシー メトリックに基づいてマスクするパッチを選択する。既存のアテンションスコアを再利用することで、ほとんど追加コストが発生しない。
• サリエンシー メトリックは、クラストークンを除くパッチ間アテンションスコアの平均値を用い、学生の予測に最も関連する画像特徴を特定・保持する。
• サリエンシー スコアが最も低いパッチがマスクされ、コアな特徴は保持されつつ、教師モデルに供給されるトークン数が削減される。
• この手法は蒸留の順伝播中に行われ、教師モデルのアーキテクチャの変更や再訓練を必要としない。
• マスキングは教師の順伝播にのみ適用され、学生は完全な画像を処理するため、特徴学習の整合性が保たれる。
• このアプローチは、ヴァナイルKD、DINO、パッチレベル蒸留を含む、さまざまなViT蒸留手法と互換性がある。

実験結果

リサーチクエスチョン

• RQ1 ViT蒸留の計算コストを、学生のトップ-1精度を損なわずに低減できるか？
• RQ2 学生の注意メカニズムからの情報のみを用いて、教師モデルへの入力パッチをマスクすることで、顕著なFLOP削減が達成できるか？
• RQ3 他のマスキング戦略（ランダムマスキングやDINOに基づくサリエンシー）と比較して、注意に基づくパッチサリエンシーは、蒸留性能の維持においてどのように優れているか？
• RQ4 MaskedKDは自己教師あり学習や視覚外のトランスフォーマー課題（例：音声処理）へも拡張可能か？
• RQ5 性能を維持しつつ最大の効率を得るために、マスクするパッチの最適な割合は何か？

主な発見

• MaskedKDは、多様なViTアーキテクチャと蒸留手法において、学生のトップ-1精度に何ら損なわれることなく、教師モデルへの入力パッチ数を25–50%削減した。
• この手法により、蒸留に必要なFLOPs（すなわち教師の推論コスト）が50%削減され、DeiT-SをDeiT-Tiに蒸留する際、合計訓練FLOPsが28%削減された。
• 教師の計算が削減されたことにより、スループットが2.0×から3.7×向上した。
• パッチ間アテンションスコアを用いたマスキングは性能を低下させ、提案されたサリエンシー メトリックがより効果的であることを確認した。
• DINOのアテンションスコアに基づくマスキングでさえ、より高い計算コストにもかかわらずMaskedKDを上回らなかった。これは、学生の注意に基づくマスキングが優れていることを示している。
• 上位k個のサリエンシー パッチ（最も重要なパッチ）を保持する戦略が、ランダムマスキングや下位k個のマスキングよりも優れていた。これは、サリエンシーと性能の間の相関関係を裏付けた。