[論文レビュー] Masked Distillation with Receptive Tokens

Tao Huang, Yuan Zhang+5 more|arXiv (Cornell University)|May 29, 2022

Advanced Neural Network Applications被引用数 7

ひとこと要約

本稿では、特徴マップ内の注目領域（PoI）を動的に特定・優先するための学習可能受容フィールドトークンを用いる、密度予測タスク向けの新規知識蒸留フレームワーク、MasKDを提案する。アテンションベースの適応的蒸留マスクを生成し、Dice損失を用いて複数の多様なマスクを統合することで、タスク固有の事前知識を必要とせず、オブジェクト検出およびセマンティックセグメンテーションのベンチマークで最先端の性能を達成する。

ABSTRACT

Distilling from the feature maps can be fairly effective for dense prediction tasks since both the feature discriminability and localization priors can be well transferred. However, not every pixel contributes equally to the performance, and a good student should learn from what really matters to the teacher. In this paper, we introduce a learnable embedding dubbed receptive token to localize those pixels of interests (PoIs) in the feature map, with a distillation mask generated via pixel-wise attention. Then the distillation will be performed on the mask via pixel-wise reconstruction. In this way, a distillation mask actually indicates a pattern of pixel dependencies within feature maps of teacher. We thus adopt multiple receptive tokens to investigate more sophisticated and informative pixel dependencies to further enhance the distillation. To obtain a group of masks, the receptive tokens are learned via the regular task loss but with teacher fixed, and we also leverage a Dice loss to enrich the diversity of learned masks. Our method dubbed MasKD is simple and practical, and needs no priors of tasks in application. Experiments show that our MasKD can achieve state-of-the-art performance consistently on object detection and semantic segmentation benchmarks. Code is available at: https://github.com/hunto/MasKD .

研究の動機と目的

• すべての画素を同等に扱う標準的な特徴蒸留の非効率性を是正する。これは、モデル性能への寄与度が画素ごとに異なるにもかかわらずである。
• 密度予測タスクにおける蒸留領域の選定に、アノテーション済みボクシングボックスやタスク固有の事前知識に依存しない。
• アテンションベースの受容フィールドトークンを用いて、特徴マップ内の微細な動的画素レベルの重要度パターンを学習する。
• 複数の多様なマスクを用いることで複雑な画素依存関係を捉え、同時に訓練の安定性を維持することで、蒸留を強化する。
• 異なるマスクの重要度を考慮するマスク重み付け機構を導入し、学生モデルの性能を向上させる。

提案手法

• 教師の特徴マップ内の注目領域（PoI）を特定するための画素単位のアテンションを計算する、学習可能な受容フィールドトークンを導入する。
• 受容フィールドトークンからの画素単位のアテンションにより、特徴再構築において優先すべき画素を示す蒸留マスクを生成する。
• 複数の受容フィールドトークンを用いて、特徴マップ内での多様で複雑な画素依存関係を捉える。
• 教師ネットワークを固定した状態で、通常のタスク損失を用いて受容フィールドトークンを訓練し、エンドツーエンド最適化を可能にする。
• 複数生成されたマスクの多様性を促進するため、Dice損失を適用する。
• 異なるマスクの蒸留への寄与度に基づき、重要度を適応的に割り当てるマスク重み付けモジュールを導入する。

実験結果

リサーチクエスチョン

• RQ1 特徴マップ内の画素を動的かつアテンションベースで選択することで、密度予測タスクにおける知識蒸留の性能が向上するか？
• RQ2 アノテーションやタスク固有の事前知識に依存せずに、情報量が多く多様な蒸留マスクを学習できるか？
• RQ3 複数の受容フィールドトークンと適応的マスク重み付けは、均一またはヒューリスティックなサンプリングと比較して、どの程度蒸留性能を向上させるか？
• RQ4 提案手法は、オブジェクト検出やセマンティックセグメンテーションといった異なるアーキテクチャやタスクに一般化できるか？
• RQ5 モデルは、教師と学生の特徴表現と整合する意味のある画素依存関係を学習できるか？

主な発見

• COCOオブジェクト検出において、Faster R-CNN-R50に対して2.4 APの向上を達成し、標準的な蒸留ベースラインを顕著に上回る。
• セマンティックセグメンテーションにおいては、COCOでDeepLabV3-R18を2.79 mIoU向上させ、タスク間での強力な一般化性能を示した。
• タスク固有の事前知識を必要とせず、複数のバックボーンアーキテクチャおよびデータセットにおいて一貫した性能向上を達成した。
• アブレーションスタディの結果、複数のマスク、多様性を促進するDice損失、およびマスク重み付けの組み合わせが、蒸留効率を顕著に向上させることを確認した。
• 可視化結果から、学習されたマスクがボクシングボックス全体ではなく、物体内での判別性の高い局所的領域に集中していることが確認され、微細なアテンションが実現している。
• 未学習のアーキテクチャやデータセットに対しても、良好な一般化性能を示しており、実世界の展開環境における頑健性と実用性を裏付けた。