[論文レビュー] Imbalanced Domain Generalization for Robust Single Cell Classification in Hematological Cytomorphology

Rao Muhammad Umer, Armin Gruber+3 more|arXiv (Cornell University)|Mar 14, 2023

Digital Imaging for Blood Diseases被引用数 4

ひとこと要約

本論文は、血液病理学的細胞像診断における一細胞単位の白血球（WBC）分類のための、交差ドメインにおけるデータ不均衡と分布シフトに耐性のある、新規な不均衡ドメイン一般化（IDG）フレームワークを提案する。ドメイン不変表現学習とクラスバランス損失最適化を組み合わせることで、不均衡なトレーニングデータとドメインシフトの両方に対処する。本手法は、未知のテストドメインで最先端のF1-macroスコアを達成し、とくにレアな細胞タイプ（例：lymphocyte_atypical）の識別性能が顕著に向上し、分布外一般化において既存の最先端手法を上回る。

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）で五分割交差検証を用いて微調整し、未知のターゲットドメイン（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ポイント上回った。
• 誤差行列の結果から、標準的なERMと比較して、少数クラス（例：lymphocyte_atypical）の識別性能が顕著に向上していることが示された。
• デカップリングトレーニングのバリエーション（Ours +）は、やや高い分散を示したが、依然として競争力のある性能を維持しており、データ不均衡下でも頑健性を示した。
• DANNのようなドメイン一般化手法は、データが不均衡な場合に性能を発揮しないため、DGフレームワークにクラスバランスの統合が不可欠であることが示された。
• 潜在空間におけるマハラノビス距離に類似した距離の使用により、ドメイン間でクラス固有の分布をより適切にモデル化でき、分布外一般化性能の向上に寄与した。