[논문 리뷰] CASS: Cross Architectural Self-Supervision for Medical Image Analysis

Pranav Singh, Elena Sizikova외 1명|arXiv (Cornell University)|2022. 06. 08.

AI in cancer detection인용 수 6

한 줄 요약

CASS는 반응 기반의 대비 학습을 사용하여 CNN과 비전 트랜스포머를 함께 사전 훈련하는 새로운 교차 아키텍처 자기지도 학습 방법을 제안한다. 이는 저데이터 의료 영상에서 표현 학습을 향상시킨다. 1%의 라벨 데이터만으로도 기존 최고 수준의 성능을 달성하며 평균 F1/재현율 3.8% 향상과 함께 훈련 시간을 기존 방법 대비 69% 감소시킨다.

ABSTRACT

Recent advances in deep learning and computer vision have reduced many barriers to automated medical image analysis, allowing algorithms to process label-free images and improve performance. However, existing techniques have extreme computational requirements and drop a lot of performance with a reduction in batch size or training epochs. This paper presents Cross Architectural - Self Supervision (CASS), a novel self-supervised learning approach that leverages Transformer and CNN simultaneously. Compared to the existing state of the art self-supervised learning approaches, we empirically show that CASS-trained CNNs and Transformers across four diverse datasets gained an average of 3.8% with 1% labeled data, 5.9% with 10% labeled data, and 10.13% with 100% labeled data while taking 69% less time. We also show that CASS is much more robust to changes in batch size and training epochs. Notably, one of the test datasets comprised histopathology slides of an autoimmune disease, a condition with minimal data that has been underrepresented in medical imaging. The code is open source and is available on GitHub.

연구 동기 및 목표

• 자기면역 질환과 같은 저대표성 질환에서 라벨이 제한된 의료 영상 분석 문제를 해결한다.
• 기존 자기지도 학습 방법의 계산 비효율성과 하이퍼파rameter 민감성 문제를 해결한다.
• CNN(국소 불변성)과 트랜스포머(전역 주의)의 장점을 활용하는 하이브리드 아키텍처를 개발하여 표현 학습을 향상시킨다.
• 특히 자가면역 및 희귀 질환 영상에서 흔히 발생하는 저데이터 시나리오에서 소량의 라벨 데이터로도 다운스트림 작업에서 성능을 향상시킨다.
• 다양한 배치 크기와 훈련 에포크 수에서 높은 성능을 유지하면서도 계산 효율성을 확보한 방법을 개발한다.

제안 방법

• CASS는 각 입력 영상을 동시에 처리하는 이중 브랜치 아키텍처를 사용한다. 각 브랜치는 별도로 CNN과 비전 트랜스포머를 처리한다.
• 두 아키텍처의 출력 간 일치를 유도하기 위해 대비 손실을 사용하여 양측의 특징 표현을 비교한다.
• 이 방법은 동일한 이미지가 서로 다른 네트워크 아키텍처에서 생성된 출력을 양성 쌍으로 삼는 반응 기반 대비 학습을 활용한다.
• 자기지도 학습 사전 훈련은 각 브랜치에 독립적으로 표준 데이터 증강(자르기, 색상 왜곡, 뒤집기 등)을 적용하여 수행된다.
• 사전 훈련 후, 클래스 불균형 문제를 해결하기 위해 포칼 손실을 사용하여 소량의 라벨 데이터로 다운스트림 분류 작업에 미세조정한다.
• 프레임워크는 자가면역 조직병리학, 피부과, 뇌 종양 MRI, 피부 병변 분류의 네 가지 데이터셋에서 평가된다.

실험 결과

연구 질문

• RQ1 CNN과 트랜스포머 아키텍처를 교차 아키텍처 자기지도 학습으로 조합함으로써 저데이터 의료 영상에서 표현 학습을 향상시킬 수 있는가?
• RQ2 극도로 라벨이 부족한 상황(1%, 10%, 100% 라벨 데이터)에서 CASS는 기존 최고 수준의 자기지도 학습 방법과 비교해 어떻게 성능을 내는가?
• RQ3 기존 방법과 비교해 배치 크기와 훈련 에포크의 변화에 대해 CASS는 어느 정도 견고한가?
• RQ4 CASS는 자가면역 질환과 같은 저대표성 의료 조건에서 더 뛰어난 성능을 내는가?
• RQ5 기존 자기지도 학습 접근법과 비교해 CASS는 성능을 유지하거나 향상시키면서도 훈련 시간을 줄일 수 있는가?

주요 결과

• 네 가지 데이터셋에서 1% 라벨 데이터만으로도 CASS는 기존 최고 수준의 자기지도 학습 방법 대비 평균 F1 및 재현율 3.8% 향상된다.
• 10% 라벨 데이터로도 CASS는 기존 방법 대비 평균 F1 및 재현율 5.9% 향상된다.
• 완전히 지도 학습(100% 라벨 데이터) 조건에서 CASS는 평균 F1 및 재현율 10.13% 향상되며, 뛰어난 일반화 능력을 보여준다.
• CASS는 현재 최고 수준의 방법 대비 평균 69% 훈련 시간을 단축시켜 훨씬 더 효율적이다.
• CASS는 다양한 배치 크기와 훈련 에포크 수에서도 일관된 성능을 유지하며, 기존 방법보다 더 뛰어난 견고성을 보인다.
• 자가면역 질환 조직병리학 데이터셋에서 CASS는 이전 최고 수준 대비 2.55% 향상된 성능을 기록하여, 저데이터 및 저대표성 질환 시나리오에서의 효과성을 입증한다.