[論文レビュー] DisCo: Remedy Self-supervised Learning on Lightweight Models with Distilled Contrastive Learning

Yuting Gao, Jiaxin Zhuang+5 more|arXiv (Cornell University)|Apr 19, 2021

Domain Adaptation and Few-Shot Learning参考文献 46被引用数 8

ひとこと要約

本稿では、大規模な教師モデルの最終的コントラスト型埋め込みを、MSE損失を用いて小規模な学生モデルに転送することで、自己教師あり学習（SSL）における軽量モデルの性能を向上させる知識蒸留フレームワーク、DisCoを提案する。この手法は、推論時にパラメータを追加しないまま、線形評価精度を顕著に向上させ、ResNet-101を教師モデルとして用いた場合、EfficientNet-B0のImageNet Top-1精度を22.1%向上させる。

ABSTRACT

While self-supervised representation learning (SSL) has received widespread attention from the community, recent research argue that its performance will suffer a cliff fall when the model size decreases. The current method mainly relies on contrastive learning to train the network and in this work, we propose a simple yet effective Distilled Contrastive Learning (DisCo) to ease the issue by a large margin. Specifically, we find the final embedding obtained by the mainstream SSL methods contains the most fruitful information, and propose to distill the final embedding to maximally transmit a teacher's knowledge to a lightweight model by constraining the last embedding of the student to be consistent with that of the teacher. In addition, in the experiment, we find that there exists a phenomenon termed Distilling BottleNeck and present to enlarge the embedding dimension to alleviate this problem. Our method does not introduce any extra parameter to lightweight models during deployment. Experimental results demonstrate that our method achieves the state-of-the-art on all lightweight models. Particularly, when ResNet-101/ResNet-50 is used as teacher to teach EfficientNet-B0, the linear result of EfficientNet-B0 on ImageNet is very close to ResNet-101/ResNet-50, but the number of parameters of EfficientNet-B0 is only 9.4\%/16.3\% of ResNet-101/ResNet-50. Code is available at https://github. com/Yuting-Gao/DisCo-pytorch.

研究の動機と目的

• モデルサイズが小さくなるにつれて精度が著しく低下する軽量モデルにおける自己教師あり学習（SSL）の性能低下を是正すること。
• 従来の蒸留手法が最適でない知識（例：出力ログティクスや特徴マップ）を転送するのではなく、最も情報量の多い最終埋め込みを転送することによる制限を克服すること。
• 最終的コントラスト型埋め込みを主な知識源として焦点を当てることで、軽量モデルのSSLにおける一般化能力を向上させること。
• 学生モデルのMLPヘッドの隠れ次元が制限されることで生じる「蒸留のボトルネック」——性能の上限——を同定・解消し、埋め込み次元を拡大することで知識伝達能力を向上させること。
• 実用的応用において効率性を保つために、ダウンストリーム展開時に追加パラメータを一切導入しない手法を設計すること。

提案手法

• 教師モデルの最終的コントラスト型埋め込みを、平均二乗誤差（MSE）損失を用いて学生モデルのそれに一致させるように制約することで、大規模な教師モデルから軽量な学生モデルへの知識転送を実現する。
• コントラスト型学習において最も情報を含むとされる、教師モデルの最終的プロジェクションヘッド出力（MLP後）を主な知識信号として用いる。
• 学生モデルの最終埋め込みが、事前学習中に教師モデルの最終埋め込みと一致するように正則化する、修正されたトレーニングパイプラインを導入する。
• 学生モデルのMLPヘッドの隠れ次元を拡大することで、「蒸留のボトルネック」を解消し、知識伝達能力を向上させる。
• 事前学習中に使用する追加のMLPヘッドは、ファインチューニングおよび推論の段階で削除され、モデルの効率性が保たれる。
• MoCo-V2、SwAV、DINOなどの多様なSSLフレームワークにDisCoを統合し、異なるアーキテクチャーや学習パラダイムにおける互換性と一般化能力を検証する。

実験結果

リサーチクエスチョン

• RQ1 大規模な教師モデルの最終的コントラスト型埋め込みを転送することで、自己教師あり学習における軽量な学生モデルの表現品質を顕著に向上させることができるか？
• RQ2 MSE損失を用いて学生モデルの最終埋め込みを教師モデルのそれに一致させることで、従来の蒸留手法よりも優れたダウンストリーム性能が得られるか？
• RQ3 小規模モデルにおける性能劣化は、知識伝達のボトルネックに起因しており、MLPヘッドの次元を拡大することでこの問題が緩和されるか？
• RQ4 MoCo-V2、SwAV、DINOなどの多様なSSLフレームワークに、アーキテクチャの変更なしにDisCoを効果的に適用できるか？
• RQ5 AT、RKD、KDなどの他の蒸留技術と組み合わせることで、さらに性能向上が得られるか？

主な発見

• ResNet-101を教師モデルとして用いた場合、EfficientNet-B0の線形評価Top-1精度が22.1%向上し、66.5%の精度に到達する。
• ResNet-50を教師モデルとして用いた場合、EfficientNet-B0のTop-1精度が19.7%向上し、66.6%の精度に到達する。
• MobileNet-v3 や ResNet-18 を含む、すべてのテスト対象の軽量モデルにおいて、SEED や MoCo-V2 よりも最先端の性能を達成する。
• 学生モデルのMLPヘッドの隠れ次元を拡大することで、性能が単調に向上し、2048次元で最大の向上が観察された。
• AT、RKD、KD などの他の蒸留手法と組み合わせることで、さらなる性能向上が得られ、DisCoの互換性と有効性が裏付けられた。
• MoCo-V2、SwAV、DINO などの多様なSSLフレームワークにおいても、強力な結果を達成しており、一般化能力と互換性の高さが示された。