[论文解读] 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

一句话总结

本文提出了一种新颖的不平衡域泛化（IDG）框架，用于在血液病细胞形态学中实现鲁棒的单细胞白血球（WBC）分类，结合域不变表示学习与类别平衡损失优化，以解决跨域数据不平衡和分布偏移问题。该方法在未见测试域上实现了最先进（SOTA）的F1-macro得分，显著提升了少数类识别能力，特别是对罕见细胞类型（如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）上进行五折交叉验证微调，并在未见目标域（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）的识别能力，混淆矩阵显示ERM漏检了大量此类样本。
• 解耦训练变体（Ours +）尽管方差略高，但仍保持了具有竞争力的性能，表明其在数据不平衡条件下的鲁棒性。
• 域泛化方法如DANN在数据不平衡时表现欠佳，凸显了将类别平衡整合进DG框架的必要性。
• 在潜在空间中使用马氏距离类度量，有助于更准确地建模跨域的类别特定分布，从而提升分布外泛化性能。