[논문 리뷰] MedSelect: Selective Labeling for Medical Image Classification Combining Meta-Learning with Deep Reinforcement Learning

Akshay Smit, Damir Vrabac외 4명|arXiv (Cornell University)|2021. 03. 26.

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

한 줄 요약

MedSelect는 대조적 사전학습 임bedding를 활용해 학습 가능한 선택기로, 의료 영상 분류에서 정보가 많은 흉부 X-ray를 선택적으로 레이블링하기 위한 메타학습 및 딥 강화학습 기반 방법을 제안한다. 이 방법은 무작위, 군집화, 메타데이터 기반 전략보다 우수한 성능을 보이며, 예측되지 않은 상태의 조건으로도 효과적으로 일반화된다. 선택된 영상은 임bedding 공간에서 유사성이 높고, 연령 및 성별 분포가 더 균형 잡혀 있다.

ABSTRACT

We propose a selective learning method using meta-learning and deep reinforcement learning for medical image interpretation in the setting of limited labeling resources. Our method, MedSelect, consists of a trainable deep learning selector that uses image embeddings obtained from contrastive pretraining for determining which images to label, and a non-parametric selector that uses cosine similarity to classify unseen images. We demonstrate that MedSelect learns an effective selection strategy outperforming baseline selection strategies across seen and unseen medical conditions for chest X-ray interpretation. We also perform an analysis of the selections performed by MedSelect comparing the distribution of latent embeddings and clinical features, and find significant differences compared to the strongest performing baseline. We believe that our method may be broadly applicable across medical imaging settings where labels are expensive to acquire.

연구 동기 및 목표

• 제한된 전문가 레이블 데이터로 인해 발생하는 의료 영상 분류 과제의 과제를 해결하기 위해, 레이블링 비용을 줄이는 선택적 레이블링 전략을 개발한다.
• 의료 영상에 대한 최적의 선택 정책을 학습함으로써, 적은 수의 레이블 예시로도 모델 성능을 향상시킨다.
• 메타학습 기간 동안 관찰하지 않은 의료 조건에 대해서도 추론 시 일반화할 수 있도록 한다.
• 임베딩, 연령, 성별, 영상 투영 유형 측면에서 선택된 영상의 분포적 성질을 분석하여 선택 편향과 전략의 효과성을 이해한다.
• 실제 임상 환경에서 레이블링이 비용이 많이 들고 시간이 오래 걸리는 상황에서 본 방법의 실용성을 입증한다.

제안 방법

• MedSelect는 대조적 사전학습에서 유도된 이미지 임베딩을 활용해, 1,000장의 레이블이 없는 흉부 X-ray 중에서 어떤 영상을 레이블링할지 결정하는 학습 가능한 딥 강화학습 선택기를 사용한다.
• 선택기는 레이블이 없는 풀에 대해 확률 분포를 출력하며, 이로부터 K장의 영상을 샘플링해 전문가 레이블링을 수행한다.
• 코사인 유사도 기반 비모수적 분류기가 레이블된 지원 세트에 대해 학습되고, AUC를 보상 신호로 사용해 쿼리 세트에서 평가된다.
• 선택기는 정책 기반 강화학습과 backpropagation를 통해 엔드 투 엔드로 학습되며, 쿼리 세트에서 높은 AUC 성능을 최적화하도록 한다.
• 10,000개의 메타학습 작업을 포함한 메타학습 프레임워크를 사용하며, 각 작업은 특정 의료 조건(예: 부종, 기도폐합)에 대한 이진 분류 문제이다.
• 메타학습 기간 동안 관찰하지 않은 조건(부종 및 기도폐합)에 대해 평가하여 제로샷 일반화 능력을 테스트한다.

실험 결과

연구 질문

• RQ1 메타학습을 통해 학습된 딥 강화학습 기반 선택기가, 제한된 레이블을 가진 의료 영상 분류에서 히우리스틱 선택 전략보다 뛰어난 성능을 보일 수 있는가?
• RQ2 MedSelect는 메타학습 기간 동안 관찰하지 않은 의료 조건으로도 일반화되는가? 즉, 제로샷 일반화 능력을 보이는가?
• RQ3 기본 전략과 비교했을 때, 선택된 영상의 잠재적 임베딩 유사도, 연령, 성별, 영상 투영 유형 측면에서 어떤 차이가 있는가?
• RQ4 선택된 영상의 분포가 연령 측면에서 무작위 선택과 유의미하게 다른가? 임베딩 유사성 측면에서도 그렇다면?
• RQ5 선택 전략이 쿼리 세트에서의 모델 성능에 미치는 영향은 무엇인가? AUC 측정 기준으로 평가된다.

주요 결과

• MedSelect는 본래의 조건과 예측되지 않은 조건 모두에서 무작위 선택, 군집 기반 선택, 메타데이터 기반 히우리스틱 전략보다 AUC 성능에서 유의미하게 뛰어나다.
• 보류된 조건(부종 및 기도폐합)에서 MedSelect는 평균 AUC 0.85를 기록했으며, 다음으로 우수한 베이스라인보다 4.2% 포인트 높은 성능(유의수준 p < 0.001)을 보였다.
• MedSelect는 군집 기반 전략보다 유의미하게 임베딩 공간에서 더 가까운 X-ray를 선택한다. 평균 쌍별 L2 거리는 7.76 vs. 8.22이며(p < 0.001), 유의미한 유사도를 보인다.
• 기본 전략이 보여주는 편향과는 달리, MedSelect는 환자의 성별 분포를 더 균형 잡히게 선택하며, 주로 앞면 영상(전향적 영상)을 선호한다.
• MedSelect 하에서 선택된 X-ray 연령의 경험적 분포는 무작위 선택과 유의미하게 다르며, Wasserstein 거리는 13.58로 측정되었고, 군집 기반 전략의 8.43보다 높다(p < 0.001).
• MedSelect 하에서 선택된 임베딩 간의 쌍별 L2 거리 분포 역시 무작위 선택과 유의미하게 다르며, Wasserstein 거리는 2.01로 군집 기반 전략의 1.22보다 높다(p < 0.001).