[논문 리뷰] A Multi-Stage Attentive Transfer Learning Framework for Improving COVID-19 Diagnosis

Yi Liu, Shuiwang Ji|arXiv (Cornell University)|2021. 01. 14.

Domain Adaptation and Few-Shot Learning참고 문헌 38인용 수 6

한 줄 요약

이 논문은 자연 이미지와 의료 영상에서의 지도 학습 및 자기지도 학습을 활용하여 흉부 CT 영상에서의 코로나19 진단을 향상시키는 다단계 주의 기반 전이 학습 프레임워크를 제안한다. 또한 폐 부위 구조를 강조하는 새로운 자기지도 학습 방법을 도입한다. 결과적으로 주의 메커니즘이 컨볼루션 레이어보다 전이 가능성이 뛰어나, ResNet-50 및 ResNet-101 모델에서 다양한 지표에서 성능 향상을 이룬다.

ABSTRACT

Computed tomography (CT) imaging is a promising approach to diagnosing the COVID-19. Machine learning methods can be employed to train models from labeled CT images and predict whether a case is positive or negative. However, there exists no publicly-available and large-scale CT data to train accurate models. In this work, we propose a multi-stage attentive transfer learning framework for improving COVID-19 diagnosis. Our proposed framework consists of three stages to train accurate diagnosis models through learning knowledge from multiple source tasks and data of different domains. Importantly, we propose a novel self-supervised learning method to learn multi-scale representations for lung CT images. Our method captures semantic information from the whole lung and highlights the functionality of each lung region for better representation learning. The method is then integrated to the last stage of the proposed transfer learning framework to reuse the complex patterns learned from the same CT images. We use a base model integrating self-attention (ATTNs) and convolutional operations. Experimental results show that networks with ATTNs induce greater performance improvement through transfer learning than networks without ATTNs. This indicates attention exhibits higher transferability than convolution. Our results also show that the proposed self-supervised learning method outperforms several baseline methods.

연구 동기 및 목표

• 코로나19 진단을 위한 정확한 딥 러닝 모델을 훈련시키기 위해 레이블이 부족한 CT 데이터의 한계를 해결하기 위해.
• 의료 영상 전이 학습에서 주의 메커니즘이 컨볼루션 레이어보다 지식 전이 효율성이 높은지 조사하기 위해.
• 다양한 크기의 폐 영역 표현을 포착하는 데 목적이 있는 새로운 자기지도 학습 방법을 개발하여 CT 영상 표현 학습을 향상시키기 위해.
• 자연 이미지, 의료 영상, 동일한 CT 데이터 등 다양한 데이터 도메인에서 지도 학습 및 자기지도 학습을 통합한 통합적이고 다단계의 전이 학습 프레임워크를 설계하기 위해.
• 주의 강화 네트워크가 원천 작업에서 목표 작업인 코로나19 진단으로 지식을 전이하는 데 효과적인지 입증하기 위해.

제안 방법

• 프레임워크는 세 단계로 구성된다: 자연 이미지에서의 지도 전이 학습(STL-N), 의료 영상에서의 지도 전이 학습(STL-M), 동일한 CT 영상에서의 자기지도 전이 학습(SSTL-M).
• 폐 부위 구조의 생물학적 지식을 활용하여 다중 척도 표현을 학습하는 새로운 자기지도 학습 방법을 제안하여 전반적인 폐 기능과 지역적 기능 정보를 모두 포착한다.
• ResNet 아키텍처에 주의 모듈(ATTNs)을 통합하여 전이 학습에서 주의 레이어와 컨볼루션 레이어의 전이 가능성 비교를 수행한다.
• 컨볼루션 연산과 주의 메커니즘을 조합한 기본 모델을 사용하여 계층적 특징을 추출하며, 설명 가능성을 위해 Grad-CAM을 활용해 주의 맵을 시각화한다.
• 단계별 전이 학습을 적용한다: 먼저 STL-N 및 STL-M에서 사전 훈련을 수행하고, 이후 목표 코로나19 진단 작업에서 미세조정을 수행한다.
• 전이 가능성을 정량적으로 평가하기 위해 LEEP 점수를 계산하며, 다양한 네트워크 변종에서 각 사전 훈련 단계의 성능 향상 정도를 비교한다.

실험 결과

연구 질문

• RQ1 의료 영상 진단을 위한 전이 학습에서 주의 기반 특징 학습이 컨볼루션 특징 학습보다 더 효과적으로 전이되는가?
• RQ2 폐 부위 구조의 기능을 강조하는 자기지도 학습 방법이 자원이 제한된 CT 데이터셋에서 표현 학습을 향상시킬 수 있는가?
• RQ3 자연 이미지, 의료 영상, 자기지도 학습 기반 CT 데이터를 활용한 다단계 전이 학습이 최종 진단 성능에 어떤 영향을 미치는가?
• RQ4 표준 컨볼루션 레이어에 비해 주의 레이어가 CT 스캔에서 감염 부위 탐지에 얼마나 향상되는가?
• RQ5 주의 레이어가 다양한 네트워크 깊이와 사전 훈련 전략에서 일관되게 성능 향상을 이끌어내는가?

주요 결과

• 주의 레이어(ATTNs)를 포함한 네트워크는 표준 ResNet보다 전이 학습을 통해 유의미하게 높은 성능 향상을 달성하며, STL-M 단계에서 1개의 ATTN을 추가했을 때 ResNet-50에서 F1 점수 4.7% 향상되었다.
• ResNet-50에서 5개의 ATTN을 추가하면 STL-N 성능이 정확도 2.8%, F1 점수 4.7%, AUC 1.5% 향상되어 모든 지표에서 일관된 향상이 확인되었다.
• 제안된 자기지도 학습 방법은 폐 CT 영상 표현 학습에서 몇몇 최첨단 기준 모델보다 뛰어난 성능을 보이며, 더 나은 후속 작업 진단을 가능하게 했다.
• Grad-CAM을 활용한 정성적 결과 분석에서 주의 레이어가 CT 스캔에서 심하게 감염된 부위를 성공적으로 강조하는 반면, 컨볼루션 레이어는 때로는 감염되지 않은 영역에 중요도를 잘못 할당하는 경향을 보였다.
• LEEP 점수 분석 결과, 주의 메커니즘이 원천 데이터와 목표 데이터의 분포 차이에 관계없이 전이 학습의 각 단계에서 더 큰 성능 향상을 이끌어내는 것으로 확인되었다.
• 본 연구는 주의 메커니즘이 의료 영상 분석을 위한 전이 학습에서 컨볼루션 레이어보다 더 높은 전이 가능성을 보이는 최초의 실증적 증거를 제공한다.