[IEEE] 논문읽기: 전이학습 for covid19 / An Uncertainty-Aware Transfer Learning-Based Framework for COVID-19 Diagnosis

오늘 읽어본 논문은

An Uncertainty-Aware Transfer Learning-Based Framework for COVID-19 Diagnosis 입니다.

논문은 아래의 링크에서 읽을 수 있습니다.

https://ieeexplore.ieee.org/document/9353390 

An Uncertainty-Aware Transfer Learning-Based Framework for COVID-19 Diagnosis

 

이 논문에서는 Covid-19 감염 조기 진단을 위한 의료 이미지의 Transfer learning 프레임워크를 제시합니다. 

 

Transfer learning?

Transfer learning은 자신이 풀고자 하는 task의 데이터가 충분히 많지 않을때 사용할 수 있는 방법으로, 현재의 task와 비슷하면서 다른, 방대한 데이터로 미리 학습되어있는 모델을 이용하는 방법입니다. 간단히 말해 사전 학습 모델의 일부를 이용하면서, 끝부분만 우리가 해결하고자 하는 task의 문제로 바꾸어 원하는 결과를 얻는 방법이라 할 수 있습니다.

사전 학습 모델을 이용함으로써 새로운 모델을 만들때 학습 속도도 더 빨라지게 되고, 가지고있는 데이터 양에 비해 예측의 정확도도 향상될것을 기대할 수 있습니다. 특히 의료 데이터는 다른 도메인의 데이터보다 얻기 힘들기 때문에(비용 등의 문제) 전이학습은 유용하고 중요한 학습 방법이라 할 수 있습니다.

 

논문 요약

X-ray 및 CT 이미지 dataset에 대하여, 사전학습 모델로 널리 사용되는 VGG16, ResNet50, DenseNet121, InceptionResNetV2 및 CNN 을 적용해 이미지에서 feature를 추출하고, 이를 기계 학습 및 통계 모델링 기술을 통해 Classification하여 Covid-19의 음성/양성을 판별하는 연구입니다. 또한 모델을 통한 분류 결과의 불확실성을 함께 계산하여 보고합니다. 결론적으로는 두 데이터셋 모두 linearSVM 모델과 neural network model에서 accuracy, sensitivity, specificity, ROC curve, AUC 값이 높은 성능을 보였다고 하며, X-ray에 비해 CT 이미지의 예측 불확실성 추정치가 높았다고 합니다.

논문에서 사용된 사전학습모델 (VGG16, ResNet50, DenseNet121, InceptionResNetV2) 의 데이터는 1.4 백만개의 이미지, 1000개의 label 로 구성된 방대한 데이터셋인 ImageNet 입니다.

연구진은 ImageNet data를 학습한 모델의 Image detection task와 Covid-19 classification 사이에 Fundamental similarities 가 있다고 생각하여 연구를 진행했습니다.

 

Dataset

연구에서 이용한 데이터셋은 흉부 X-ray 이미지(2D)와, CT 이미지(2D)입니다.

사전 학습 모델이 2D data를 input으로 받는 모델이다보니, 3D 이미지인 여러장의 CT data에서, 한 장의 CTdata만 이용해 2D로 학습시킨것으로 보입니다.

 

Chest X-Ray Data set
✓ 코로나 양성 이미지 25장
✓ 코로나 음성 이미지 75장
✓ non-Covid : consist of other unhealthy conditions
→ non-Covid does not necessarily infer a healthy lower respiratory system

CT Data set
✓ 코로나 양성 이미지 349장
✓ 코로나 음성 이미지 397장
✓ carry more information compared to chest X-rays

 

 

위에도 적어놓았다시피, 흉부 X-ray dataset의 '음성' label 이미지는 단순히 건강한 사람(무질병자)의 이미지가 아니라, Covid가 아닌 다른 질병에 걸린 환자의 dataset도 포함합니다. 즉, 질병자와 무질병자를 나누는 모델이 아니라 코로나 환자와 비코로나 환자를 나누는 모델을 만들고자 하였다는 것입니다. 

 

Framework

연구에서는 CNN을 이용한 사전학습모델로 상기 4가지 모델을 사용했고, Classifier로는 다음의 8가지 모델을 사용했습니다. 최종 모델의 앞부분은 사전학습모델의 Convolution layer를 그대로 가져와서 가중치를 사용하였고, 뒷부분은 Covid classification을 위한 classifier를 사용하였습니다. 

 

Visualization

Performance Metrics (Percentage)

VGGnet 을 이용해 각 data에서 25088개의 feature를 추출하고, 이를 PCA를 통해 2개의 feature로 줄인 그림입니다. (설명력은 논문에 기재되지 않은것으로 보임.) 각 포인트는 환자 한 명의 데이터입니다. 즉, 환자 한 명의 데이터를 2차원 평면에 시각화한 그림입니다. 흉부 X-ray data는 음성과 양성 데이터의 구분이 꽤 뚜렷한 반면에 CT data는 그러한 경계가 없다는 것을 확인할 수 있습니다.

 

Visualize the effectiveness of extracted features

gradient-weighted class activation mapping 그림입니다. Backpropagating 한 결과를 HeatMap 으로 그려 X-ray 사진에 HeatMap을 덧씌운 그림입니다. 모델이 데이터를 분류할때 어느 부분이 중요하게 영향을 끼쳤는지 확인이 가능합니다.

 

 

Evaluation

연구진은 100 회에 걸쳐 Train 및 Evaluation 을 진행하였고, Accuracy, Sensitivity, Specificity, AUC value 값을 평가 지표로서 활용하였습니다. 아래의 표는 모든 예측(100회) 의 평가지표의 평균을 낸 값입니다. 

pre-trained 모델로 ResNet50 + 분류기 모델로 Linear SVM을 사용한 조합이 각 데이터셋에서 높은 성과를 보인것을 확인할 수 있습니다.

ROC-AUC score

각 모델별 ROC curve 입니다. linear SVM 과 neural Net 이 가장 높은 AUC 값을 가지는것을 확인할 수 있습니다. 

또한 classifier의 성능은 pre-trained CNN 모델에 따라 다양하게 나타나고 있습니다.

연구진은 추가적으로, 각 모델별로 추출되는 feature 및 파라미터의 갯수 또한 비교하였는데요. 그래프의 원은 pre-trained CNN 모델이며, 원의 크기는 Accuracy를 의미합니다. VGG16 모델은 적은 파라미터 및 feature 갯수에 비해 높은 정확도, AUC 값을 가지고 있고(Most informative& Discriminative), InceptionResNetV2는 높은 파라미터 및 feature 갯수에 비해 낮은 정확도, AUC 값을 가지고 있습니다. (Least informative & Discriminative features)

이러한 자료를 통해, 큰 네트워크를 가진 모델이더라도 반드시 more informative and discriminative features 를 추출하는것은 아님을 알 수 있습니다.

 

predictive Uncertainty Estimation

Uncertainty 는 2가지가 있습니다.

Aleatoric Uncertainty는, 어쩔 수 없이 발생하는 오류로 학습을 통해 줄일 수 없고,

Epistemic Uncertainty는 모델이 모르기때문에 발생하는 오류로 학습을 통해 줄일 수 있습니다.

연구진은 추가적으로, 모델의 Epistemic Uncertainty를 줄이기 위해 이를 추정해 나타내고자 하였습니다.

그리고 이를 위해 Ensemble model을 만들었습니다.

Ensemble 모델은 1개의 Hidden layer 를 가진 20개의 독립적인 neural network로 구성되어있고, hidden layer의 neuron 는 랜덤하게 50~400 개 중 선택되었습니다.

불확실성은 다음과 같이 계산되었습니다. (p 는 softmax 함수와 비슷해 보이네요..)

 

이를 다시 PCA를 통해 dimension을 줄여 2차원으로 만들었고, 시각화하여 나타내었습니다.

밝은색(녹색)일수록 높은 확실성, 어두운색(남색)일수록 높은 불확실성을 나타내고,

불확실성이 높다는 것은 앙상블 모델 내부의 classifier들의 분류가 서로 완전히 일치하지 않다는 것을 의미합니다.

X-ray의 Epistemic Uncertainty 는 상대적으로 낮고, CT 이미지는 꽤 높은것을 확인할 수 있습니다.

 

앞서 말씀드렸듯 CT 이미지는 원래 흉부를 관통하여 찍는 3D 이미지인데요, 이를 모두 사용하지않고 임의로 한 장을 선택해 분류를 진행하였으니 x-ray에 비해 불확실성이 높게 나타난 것이 아닌가 싶습니다. (가령 Covid를 분류하기 위한 설명력 높은 feature를 추출할 수 있는 위치에서 찍은 이미지가 아니었다던가..) 

 

Further Exploration

마지막으로, 연구진은 이 모델을 개선하기 위해 fine-tuning, hybrid models(feature 간의 통합 등), ensembles, more comprehensive Uncertainty measure 도출 등 다양한 방법을 제시하였습니다.

 

 

끝!