[논문 리뷰] Intermediate Layers Matter in Momentum Contrastive Self Supervised Learning

Aakash Kaku, Sahana Upadhya외 1명|arXiv (Cornell University)|2021. 10. 27.

AI in cancer detection참고 문헌 28인용 수 9

한 줄 요약

이 논문은 평균 제곱오차(MSE) 또는 Barlow Twins 스타일의 불변성 손실를 사용하여 중간 네트워크 레이어의 특징 비유사도를 최소화함으로써 모멘텀 대비(모코) 자기지도 학습을 향상시키는 방법을 제안한다. 이 방법은 특징을 네트워크의 더 이른 단계에서 더 재사용 가능하고 정보량이 많은 방식으로 학습시켜, 특히 저표본 레이블 환경에서 더 높은 다운스트림 성능을 달성한다. 세 가지 다양한 의료 영상 데이터셋에서 표준 MoCo보다 평균 5%의 성능 향상을 기록했으며, 레이블 데이터가 1%일 때도 유의미한 향상이 있었다.

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 스타일의 손실를 추가로 도입한다.
• 표준 MoCo에서처럼 최종 레이어뿐 아니라 여러 중간 레이어에 이러한 손실를 적용함으로써 MoCo 프레임워크를 확장한다.
• 표준 MoCo와 마찬가지로 음성 키의 모멘텀 업데이트 큐를 사용하지만, 중간 레이어 수준의 대비 목적을 추가한다.
• 특징 재사용 분석을 통해 피니튜닝 전후의 특징 유사도를 측정함으로써 특징 품질을 평가한다.
• 레이어별 탐색을 수행하여 네트워크의 다양한 깊이에서 특징의 정보량을 평가한다.
• 1% 및 6% 레이블 데이터로 피니튜닝한 모델과 100% 레이블 데이터로 학습한 최고 성능 모델 간의 출력 확률 분포 간 콜모고로프–스미르노프(KS) 거리를 비교하여 일반화 성능을 평가한다.

실험 결과

연구 질문

• RQ1 MoCo에서 중간 레이어 표현의 정렬을 향상시키면 의료 영상 데이터셋에서 다운스트림 분류 성능이 향상되는가?
• RQ2 제안된 방법은 표준 MoCo와 비교해 중간 레이어에서 특징 재사용성과 정보량에 어떤 영향을 미치는가?
• RQ3 표현 품질 향상이 저레이블 데이터 환경에서 성능 향상으로 얼마나 잘 이어지는가?
• RQ4 제안된 방법으로 사전학습한 모델의 출력 확률 분포는 1% 레이블 데이터로 피니튜닝한 최고 성능 모델과 비교해 얼마나 유사한가?
• RQ5 중간 레이어 정렬이 다양한 의료 영상 모odalities에서 더 강건하고 일반화 가능한 특징을 유도할 수 있는가?

주요 결과

• MSE 손실를 사용한 제안된 방법은 NIH 흉부 X-ray, 당뇨병 망막병변, 유방암 조직형광 영상 세 가지 의료 영상 데이터셋에서 모두 표준 MoCo를 초월했다. 특히 레이블 비율이 낮을수록 두드러진 성능 향상을 보였다.
• 레이블 데이터가 1%일 때, 데이터셋 간 평균 성능 향상은 표준 MoCo 대비 5%였으며, 이는 저자료 환경에서의 유의미한 향상을 시사한다.
• 중간 레이어 정렬을 적용한 모델는 더 높은 특징 재사용성을 보였으며, 이는 더 높은 품질의, 더 이식 가능한 특징을 학습했다는 것을 의미한다.
• 레이어별 탐색 결과, 제안된 방법에서는 표준 MoCo보다 네트워크의 더 이른 단계에서 학습된 특징이 더 정보량이 많았다.
• 1% 및 6% 레이블 데이터로 피니튜닝한 모델과 100% 레이블 데이터로 학습한 최고 성능 모델 간의 출력 분포 간 콜모고로프–스미르노프(KS) 거리는 제안된 방법에서 더 작았으며, 이는 최적 성능에 더 가까운 성능을 달성했다는 것을 의미한다.
• 성능 향상은 다양한 의료 영상 모달리티에서 일관되게 관찰되어, 이 방법이 특정 데이터 유형을 초월해 일반화됨을 시사한다.