[논문 리뷰] Imbalanced Domain Generalization for Robust Single Cell Classification in Hematological Cytomorphology

Rao Muhammad Umer, Armin Gruber외 3명|arXiv (Cornell University)|2023. 03. 14.

Digital Imaging for Blood Diseases인용 수 4

한 줄 요약

이 논문은 혈액학 세포형태학에서 단일세포 백혈구(WBC) 분류의 강건한 성능을 확보하기 위해 교차 도메인 데이터 불균형과 분포 이탈 문제를 해결하기 위한 새로운 비균형 도메인 일반화(IDG) 프레임워크를 제안한다. 도메인 불변 표현 학습과 클래스 균형 손실 최적화를 결합한 방법으로, 특히 희귀 세포 유형인 lymphocyte_atypical과 같은 소수의 세포 유형에 대한 인식 성능를 크게 향상시키며, 기존의 최고 성능(SOTA) 방법보다 외부 도메인 일반화 성능에서 뛰어난 성과를 보였다.

ABSTRACT

Accurate morphological classification of white blood cells (WBCs) is an important step in the diagnosis of leukemia, a disease in which nonfunctional blast cells accumulate in the bone marrow. Recently, deep convolutional neural networks (CNNs) have been successfully used to classify leukocytes by training them on single-cell images from a specific domain. Most CNN models assume that the distributions of the training and test data are similar, i.e., the data are independently and identically distributed. Therefore, they are not robust to different staining procedures, magnifications, resolutions, scanners, or imaging protocols, as well as variations in clinical centers or patient cohorts. In addition, domain-specific data imbalances affect the generalization performance of classifiers. Here, we train a robust CNN for WBC classification by addressing cross-domain data imbalance and domain shifts. To this end, we use two loss functions and demonstrate their effectiveness in out-of-distribution (OOD) generalization. Our approach achieves the best F1 macro score compared to other existing methods and is able to consider rare cell types. This is the first demonstration of imbalanced domain generalization in hematological cytomorphology and paves the way for robust single cell classification methods for the application in laboratories and clinics.

연구 동기 및 목표

• 다양한 임상 영상 자료에서 단일세포 WBC 분류 시 교차 도메인 데이터 불균형과 도메인 이탈 문제를 해결하기 위해.
• 특히 부족하게 표현된 세포 유형과 희귀 세포 유형에 대해 외부 도메인(OOD) 테스트 도메인으로의 일반화 성능를 향상시키기 위해.
• 다른 병원이나 영상 프로토콜 간에 불균형한 학습 데이터가 존재하더라도 소수의 클래스에서 높은 성능를 유지할 수 있는 딥 러닝 프레임워크를 개발하기 위해.
• 데이터 분포가 상당히 다를 수 있는 실제 임상 진단 환경에서 강건하고 임상적으로 신뢰할 수 있는 AI 기반 분류를 가능하게 하기 위해.

제안 방법

• 이 방법은 이중 손실 학습 전략을 사용한다: 분류를 위한 표준 크로스 엔트로피 손실과 잠재 특징 공간에서의 도메인 불변 대비 손실을 통해 도메인 이탈을 줄인다.
• 각 도메인의 클래스별 특징 분포를 모델링하기 위해 $ d(\mathbf{z}, \{\boldsymbol{\psi}_{d,c}, \boldsymbol{\Sigma}_{d,c}\}) = \frac{1}{N_{d_i,c_i}} \sqrt{(\mathbf{z}-\boldsymbol{\psi}_{d,c})^\top \boldsymbol{\Sigma}_{d,c}^{-1}(\mathbf{z}-\boldsymbol{\psi}_{d,c})} $와 같은 마할라노비스 유사도 거리 측도를 사용한다.
• 에코더와 분류기 헤드의 결합 최적화와 분離 최적화를 동시에 수행하며, 후자의 경우 클래스 균형 샘플링을 사용하여 데이터 불균형 문제를 완화한다.
• ImageNet에서 미리 훈련된 ResNet50 백본을 사용하여 소스 도메인(Matek_19, Acevedo_20)에서 5겹 교차 검증을 수행하고, 예측은 알려지지 않은 타겟 도메인(INT_20)에서 수행한다.
• 특정 도메인에서 클래스가 누락되는 경우를 명시적으로 처리하기 위해, 데이터가 부족한 상황에서도 일반화 가능한 도메인 불변 표현을 학습한다.

실험 결과

연구 질문

• RQ1 학습 데이터에 심한 교차 도메인 불균형이 존재할 경우, 딥 러닝 모델이 단일세포 WBC 분류에서 외부 도메인 일반화 성능를 강건하게 확보할 수 있는가?
• RQ2 표현 학습과 분류 손실을 함께 최적화함으로써, 다양한 영상 자료에서 희귀하고 부족하게 표현된 WBC 아형에 대한 성능 향상 정도는 어떠한가?
• RQ3 도메인 불변 특징 학습이 임상 영상 프로토콜의 도메인 이탈로 인한 성능 저하를 어느 정도 완화할 수 있는가?
• RQ4 클래스 균형 샘플링을 적용한 분리 학습 방식은 기존 도메인 일반화 방법에 비해 소수 클래스의 매크로-F1 점수를 향상시키는가?

주요 결과

• 제안된 IDG 방법은 알려지지 않은 INT_20 테스트 세트에서 최고의 F1-macro 점수 0.78 ± 0.05를 기록했으며, ERM(0.40 ± 0.05), DANN(0.35 ± 0.06), CORAL(0.43 ± 0.03)을 모두 앞서나갔다.
• 소스 도메인 검증 세트(Acevedo_20 및 Matek_19)에서 이 방법은 F1-macro 점수 0.78 ± 0.05를 기록했으며, 이는 다음으로 우수한 성능를 보인 CORAL보다 0.02점 높았다.
• Confusion matrix를 통해 ERM에 비해 소수의 클래스—특히 lymphocyte_atypical—의 인식 성능 향상이 뚜렷하게 나타났다.
• 분리 학습 변형(Ours +)은 높은 변동성을 보였지만도 경쟁 가능한 성능를 유지하여, 데이터 불균형 상황에서도 강건함을 입증했다.
• DANN과 같은 도메인 일반화 방법은 데이터 불균형 상황에서 성능가 낮아지며, 이는 DG 프레임워크에 클래스 균형 기법을 통합할 필요성이 있음을 시사한다.
• 잠재 공간에서의 마할라노비스 유사도 거리 측도를 활용함으로써, 다양한 도메인 간 클래스별 특징 분포를 더 잘 모델링할 수 있었으며, 이는 OOD 일반화 성능 향상에 기여했다.