[論文レビュー] Intermediate Layers Matter in Momentum Contrastive Self Supervised Learning

Aakash Kaku, Sahana Upadhya+1 more|arXiv (Cornell University)|Oct 27, 2021

AI in cancer detection参考文献 28被引用数 9

ひとこと要約

本論文では、Momentum Contrastive (MoCo)自己教師付き学習を、中間層の特徴量の類似度を最小化することで向上させる手法を提案する。平均二乗誤差（MSE）またはBarlow Twins風の不変性損失を用いる。この手法により、ネットワークの初期段階でより再利用可能で情報量の多い特徴量を学習でき、特にラベルが少ない状況下でも性能が向上し、3つの多様な医用画像データセットで標準的なMoCoを上回る。1%のラベルデータでの平均的な性能向上は5%に達する。

ABSTRACT

We show that bringing intermediate layers' representations of two augmented versions of an image closer together in self-supervised learning helps to improve the momentum contrastive (MoCo) method. To this end, in addition to the contrastive loss, we minimize the mean squared error between the intermediate layer representations or make their cross-correlation matrix closer to an identity matrix. Both loss objectives either outperform standard MoCo, or achieve similar performances on three diverse medical imaging datasets: NIH-Chest Xrays, Breast Cancer Histopathology, and Diabetic Retinopathy. The gains of the improved MoCo are especially large in a low-labeled data regime (e.g. 1% labeled data) with an average gain of 5% across three datasets. We analyze the models trained using our novel approach via feature similarity analysis and layer-wise probing. Our analysis reveals that models trained via our approach have higher feature reuse compared to a standard MoCo and learn informative features earlier in the network. Finally, by comparing the output probability distribution of models fine-tuned on small versus large labeled data, we conclude that our proposed method of pre-training leads to lower Kolmogorov-Smirnov distance, as compared to a standard MoCo. This provides additional evidence that our proposed method learns more informative features in the pre-training phase which could be leveraged in a low-labeled data regime.

研究の動機と目的

• 中間層表現の整合性を向上させることで、医用画像分野における自己教師付き表現学習が改善するかどうかを調査すること。
• 医用AIにおけるラベルデータの限界を克服するために、微調整用の事前学習品質を向上させること。
• より良い特徴量品質が一般化性能にどのように寄与するか、特にラベルデータが少ない状況下での有効性を評価すること。
• 線形プローブを超える特徴量品質の評価のため、特徴量の再利用性と層ごとの情報量を分析すること。
• 微調整後の出力確率分布を比較し、モデルが最適性能にどれほど近いかを評価すること。

提案手法

• 中間層特徴量の間の平均二乗誤差（MSE）と、特徴量間の相互相関を最小化するBarlow Twins風の損失を追加する2つの損失項を導入する。
• 標準的なMoCoとは異なり、最終層だけでなく複数の中間層に対してもこれらの損失を適用する。
• 標準的なMoCoと同様に、負例キューをモーメンタム更新するが、中間層レベルの対照的な目的関数を追加する。
• 特徴量再利用解析を用いて、微調整前後における特徴量の類似度を測定し、特徴量品質を評価する。
• 層ごとのプロービングを実施し、ネットワーク内の異なる深さにおける特徴量の情報量を評価する。
• 1%と6%のラベルデータで微調整したモデルと100%ラベルデータで微調整した最良のモデルとの出力確率分布の間のコルモゴロフ–スミルノフ（KS）距離を比較し、一般化性能を評価する。

実験結果

リサーチクエスチョン

• RQ1 MoCoにおける中間層表現の整合化は、医用画像データセットにおける下流分類性能の向上に寄与するか？
• RQ2 提案手法は、標準的なMoCoと比較して、ネットワークの各層における特徴量の再利用性と情報量にどのような影響を与えるか？
• RQ3 特徴量品質の向上が、特にラベルデータが少ない状況下で、どの程度性能向上に結びつくか？
• RQ4 提案手法で事前学習されたモデルの出力確率分布は、1%ラベルデータで微調整した最良のモデルと比較して、最適性能にどれほど近いか？
• RQ5 中間層の整合化は、多様な医用画像モodalitiesにおいて、より頑健で一般化可能な特徴量を生成できるか？

主な発見

• MSE損失を用いた提案手法は、NIH Chest X-ray、糖尿病網膜症、乳がんヒストパスロジーの3つの医用画像データセットすべてで、標準的なMoCoを上回った。特にラベル割合が低い状況で顕著な向上が見られた。
• 1%のラベルデータでの平均的な性能向上は、標準的なMoCoと比較して5%に達し、ラベルが少ない状況下での顕著な改善を示した。
• 中間層の整合化を施したモデルは、より高い特徴量再利用性を示し、品質が高く、より再利用可能な特徴量を学習していた。
• 層ごとのプロービングにより、提案手法では標準的なMoCoと比較して、ネットワークの初期段階で学習された特徴量がより情報量が多くなっていることが明らかになった。
• 1%と6%のラベルデータで微調整したモデルと100%ラベルデータで微調整した最良のモデルとの出力分布の間のコルモゴロフ–スミルノフ（KS）距離が、提案手法では小さく抑えられ、最適性能に近い出力を得ていることが示された。
• 得られた向上効果は、多様な医用画像モダリティにわたり一貫しており、本手法が特定のデータタイプに限定されず、一般化可能であることが示された。