[논문 리뷰] Label-Assemble: Leveraging Multiple Datasets with Partial Labels

Mintong Kang, Li, Bowen외 5명|arXiv (Cornell University)|2021. 09. 25.

COVID-19 diagnosis using AI인용 수 4

한 줄 요약

Label-Assemble는 부분적으로 레이블이 부여된 의료 데이터셋을 활용하기 위해 부정적 및 관련 질환 레이블을 보조 신호로 간주하는 혁신적인 프레임워크를 제안한다. 클래스 쿼리, 의사 레이블링, 일致성 정규화를 통해, 전체 레이블이 필요로 하지 않는 이질적인 공공 데이터셋에서 병리학적 또는 공간적으로 관련된 레이블을 조립함으로써, 최신 기술 수준의 성능을 달성한다. 이는 코로나19 진단 정확도를 99.3%로 향상시키고, PDAC 검출 민감도를 84.0%로 높인다.

ABSTRACT

The success of deep learning relies heavily on large labeled datasets, but we often only have access to several small datasets associated with partial labels. To address this problem, we propose a new initiative, "Label-Assemble", that aims to unleash the full potential of partial labels from an assembly of public datasets. We discovered that learning from negative examples facilitates both computer-aided disease diagnosis and detection. This discovery will be particularly crucial in novel disease diagnosis, where positive examples are hard to collect, yet negative examples are relatively easier to assemble. For example, assembling existing labels from NIH ChestX-ray14 (available since 2017) significantly improves the accuracy of COVID-19 diagnosis from 96.3% to 99.3%. In addition to diagnosis, assembling labels can also improve disease detection, e.g., the detection of pancreatic ductal adenocarcinoma (PDAC) can greatly benefit from leveraging the labels of Cysts and PanNets (two other types of pancreatic abnormalities), increasing sensitivity from 52.1% to 84.0% while maintaining a high specificity of 98.0%.

연구 동기 및 목표

• 부분적으로 레이블이 부여된 공공 데이터셋을 활용하여 의료 영상에서 레이블이 부족한 문제를 해결한다.
• 기존 방법이 데이터의 레이블이 부여된 부분만 사용하고, 레이블이 없는 또는 부분적으로 레이블이 부여된 예제의 잠재력을 간과하는 한계를 극복한다.
• 희귀 또는 신규 질환을 위한 모델 성능 향상에 도움이 되는 부정적 또는 관련 질환 예제로부터의 학습이 가능한지 조사한다.
• 다양한, 이질적이고 때로는 상충되는 레이블을 여러 데이터셋 간에 동적으로 통합할 수 있는 통합 프레임워크를 개발한다.
• 병리학적 또는 공간적으로 관련된 질환의 레이블을 조립함으로써 진단 및 검출 과제에서 성능 향상이 뚜렷하게 이루어지는지 입증한다.

제안 방법

• 다양한 카테고리의 데이터셋 간에 부분 레이블을 동적으로 통합할 수 있도록 가속 가능한 클래스 쿼리 메커니즘을 도입하여 시각적 작업을 인코딩한다.
• 공유된 컨volutional 신경망(CNN)을 사용해 이미지 특징을 추출하고, 특징과 클래스 쿼리 간의 내적 곱을 통해 예측을 계산하는 예측 헤드를 적용한다.
• 레이블이 존재할 경우 이진 교차 엔트로피 손실(ℒ_bce)을 적용하고, 레이블이 누락된 경우 의사 레이블링 및 일致성 정규화를 활용한 비지도 손실(ℒ_pseudo 및 ℒ_consist)을 적용한다.
• 의사 레이블링을 통해 레이블이 없는 샘플에 대한 예측을 생성하고, 변환에 대한 일관성을 강제하여 데이터셋 간 도메인 이동 문제를 완화한다.
• 지도 학습 및 자기지도 학습 목표를 조합하여 종합적으로 훈련함으로써 이질적 원천에서 온 레이블이 있는 및 없는 데이터를 모두 활용한다.
• 고정된 카테고리 분류와 달리 클래스 쿼리 메커니즘을 통해 임의의, 가변적인 수의 클래스를 처리할 수 있도록 다중 작업 학습에 프레임워크를 적응시킨다.

실험 결과

연구 질문

• RQ1 부분적으로 레이블이 부여된 데이터셋에서 부정 예외를 학습함으로써, 코로나19와 같은 목표 질환에 대한 모델 성능 향상이 가능한가?
• RQ2 병리학적으로 유사한 질환(예: 코로나19에 대한 폐렴)의 레이블을 조립하면, 양성 예외에만 의존하는 것보다 더 나은 진단 성능을 달성할 수 있는가?
• RQ3 공간적으로 관련된 이상소견(예: PDAC에 대한 낭종 및 PanNETs)의 레이블을 활용하면 질환 검출 민감도가 향상되는가?
• RQ4 다양한 이질적 데이터셋에서 부분 레이블만을 사용하여 전체 레이블 학습 수준의 성능을 달성할 수 있는가?
• RQ5 기존의 레이블이 있는 데이터만 사용하거나 전체 레이블이 필요한 방법과 비교해 볼 때, 제안된 Label-Assemble 프레임워크는 어떻게 성능을 내는가?

주요 결과

• Label-Assemble는 NIH ChestX-ray14에서 레이블을 조립함으로써 코로나19 진단 정확도를 96.3%에서 99.3%로 향상시켰다. 이는 부정적 및 관련 질환 예제로부터의 학습이 결정 경계 학습을 향상시킨다는 것을 보여준다.
• 췌장 도파일 아데노카르시노마(PDAC) 검출 민감도를 낭종 및 PanNETs와 같은 관련 췌장 이상소견의 레이블을 통합함으로써 52.1%에서 84.0%로 높였고, 높은 특이도(98.0%)를 유지했다.
• 성능 향상은 클래스 간 유사도와 정적 상관관계가 있다: 폐 결절 분류 과제에서 향상 정도와 조립된 질환의 유사도 간 피어슨 상관계수는 r=0.83 (p=4.93e-4)이다.
• Label-Assemble는 NIH ChestX-ray14에서 기존 방법을 초월하는 최신 기술 수준의 mAUC 0.832를 달성했으며, 14개의 흉부 질환 중 13개에서 기존 방법을 뛰어넘었다. 또한 CheXpert 성능은 기준 모델 대비 1.8% 향상되었다.
• 완전한 레이블 학습 성능과 유사한 성능을 달성하면서도 레이블링 비용을 40% 절감(75,310개의 부분 레이블 대비 105,434개의 전체 레이블)함으로써, 전체 레이블 완성은 필수적이지 않다는 것을 입증했다.
• 기존의 부분 레이블 기반 기준 모델인 DoDNet 및 Med3D보다도 성능이 뛰어나며, 특히 다중 작업 및 준지도 학습 환경에서 그 성능이 뛰어나지, 적응형 클래스 쿼리 및 일치성 정규화 구성 요소 덕분이다.