[논문 리뷰] Highly Efficient Representation and Active Learning Framework for Imbalanced Data and its Application to COVID-19 X-Ray Classification

Heng Hao, Sima Didari외 3명|arXiv (Cornell University)|2021. 02. 24.

COVID-19 diagnosis using AI참고 문헌 41인용 수 4

한 줄 요약

이 논문은 자기지도 학습 기반 표현 학습과 가우시안 프로세스 기반 활성 학습을 융합하여, 불균형한 흉부 X선 영상 분류를 위한 데이터 및 레이블 효율적인 프레임워크를 제안한다. 비지도 특징 학습과 불확실성 인식 예측을 활용함으로써, 레이블이 약 10%인 경우에도 최신 기술 수준의 성능을 달성하며, 코로나19 X선 영상 분류에서 이전 모델보다 4% 높은 성능을 보였다.

ABSTRACT

We propose a highly data-efficient classification and active learning framework for classifying chest X-rays. It is based on (1) unsupervised representation learning of a Convolutional Neural Network and (2) the Gaussian Process method. The unsupervised representation learning employs self-supervision that does not require class labels, and the learned features are proven to achieve label-efficient classification. GP is a kernel-based Bayesian approach that also leads to data-efficient predictions with the added benefit of estimating each decision's uncertainty. Our novel framework combines these two elements in sequence to achieve highly data and label efficient classifications. Moreover, both elements are less sensitive to the prevalent and challenging class imbalance issue, thanks to the (1) feature learned without labels and (2) the Bayesian nature of GP. The GP-provided uncertainty estimates enable active learning by ranking samples based on the uncertainty and selectively labeling samples showing higher uncertainty. We apply this novel combination to the data-deficient and severely imbalanced case of COVID-19 chest X-ray classification. We demonstrate that only $\sim 10\%$ of the labeled data is needed to reach the accuracy from training all available labels. Its application to the COVID-19 data in a fully supervised classification scenario shows that our model, with a generic ResNet backbone, outperforms (COVID-19 case by 4\%) the state-of-the-art model with a highly tuned architecture. Our model architecture and proposed framework are general and straightforward to apply to a broader class of datasets, with expected success.

연구 동기 및 목표

• 코로나19 X선 영상 분류에서 데이터 부족과 심각한 클래스 불균형 문제를 해결하기 위해.
• 자기지도 학습을 활용하여 최소한의 레이블 데이터를 요구하는 표현 학습 방법을 개발하기 위해.
• 불확실성 추정을 가우시안 프로세스를 통해 통합하여 활성 학습을 수행함으로써 정보가 풍부한 샘플을 우선순위로 선정하기 위해.
• 큰 레이블 데이터셋에 의존하는 것의 의존도를 줄이면서도 높은 분류 정확도를 유지하기 위해.
• 기타 불균형된 의료 영상 분류 과제에 일반화 가능한 프레임워크를 구축하기 위해.

제안 방법

• 클래스 레이블이 필요 없이 흉부 X선 영상에서 대비 기반 사전 훈련을 통한 자기지도 표현 학습을 활용한다.
• 사전 훈련된 ResNet 기반 모델을 사용하여 비지도 방식으로 X선 영상에서 깊은 특징을 추출한다.
• 예측과 함께 불확실성 추정을 제공하는 가우시안 프로세스 분류기를 적용하여 불확실성 기반 활성 학습을 가능하게 한다.
• 비레이블된 샘플을 불확실성 기준으로 순위를 매기고, 반복적인 레이블링 과정에서 가장 불확실한 샘플을 인간의 레이블링 대상으로 선정한다.
• 비지도 특징 추출기와 GP 분류기를 순차적 파이프라인으로 조합하여 데이터 효율성을 향상시킨다.
• GP의 베이지안 성격과 레이블에 의존하지 않는 특징 학습을 활용하여 클래스 불균형에 대한 민감도를 감소시킨다.

실험 결과

연구 질문

• RQ1 자기지도 표현 학습은 흉부 X선 영상 분류에서 레이블 데이터의 필요성을 크게 줄일 수 있는가?
• RQ2 가우시안 프로세스를 통한 불확실성 추정은 불균형 데이터셋에서 활성 학습 성능을 어떻게 향상시키는가?
• RQ3 자기지도 특징를 가진 사전 훈련된 CNN 모델이 제한된 레이블 데이터로 인해 완전히 지도 학습된 모델보다 얼마나 뛰어나게 성능을 내는가?
• RQ4 자기지도 학습과 GP 기반 활성 학습의 조합이 의료 영상에서의 클래스 불균형 영향을 완화하는가?
• RQ5 이 프레임워크는 다른 데이터가 부족하고 불균형한 의료 영상 분류 과제로 일반화 가능한가?

주요 결과

• 제안된 프레임워크는 레이블 데이터의 약 10%만으로도 완전히 지도 학습한 결과와 유사한 성능을 달성한다.
• 일반적인 ResNet 기반 모델을 사용하여 코로나19 X선 영상 분류 과제에서 최신 기술 수준의 모델보다 4% 높은 성능을 기록한다.
• 레이블 없는 특징 학습과 베이지안 불확실성 추정 덕분에 클래스 불균형에 대해 강건함을 입증한다.
• GP 불확실성 기반 활성 학습은 정보가 풍부한 샘플을 효과적으로 우선순위로 선정하여 레이블링 노력의 감소를 이룬다.
• 제한된 레이블링 조건에서도 높은 정확도를 유지함으로써, 이 프레임워크의 데이터 효율성이 뚜렷하게 드러난다.
• 이 방법은 잘 일반화되며, 유사한 과제를 가진 다른 의료 영상 데이터셋에 쉽게 적용할 수 있다.