[论文解读] Improve Model Generalization and Robustness to Dataset Bias with Bias-regularized Learning and Domain-guided Augmentation

Yundong Zhang, Hang Wu+3 more|arXiv (Cornell University)|Oct 12, 2019

COVID-19 diagnosis using AI参考文献 26被引用 7

一句话总结

本文提出了一种领域泛化框架，通过结合偏差正则化学习与领域引导的数据增强，提升了深度学习模型在生物医学影像中对数据集偏差的鲁棒性。该方法通过一种新型损失函数和多层交叉梯度训练，使模型学习领域不变特征，在无需微调的情况下，实现了在未见胸部X光数据集上的最先进泛化性能。

ABSTRACT

Deep Learning has thrived on the emergence of biomedical big data. However, medical datasets acquired at different institutions have inherent bias caused by various confounding factors such as operation policies, machine protocols, treatment preference and etc. As the result, models trained on one dataset, regardless of volume, cannot be confidently utilized for the others. In this study, we investigated model robustness to dataset bias using three large-scale Chest X-ray datasets: first, we assessed the dataset bias using vanilla training baseline; second, we proposed a novel multi-source domain generalization model by (a) designing a new bias-regularized loss function; and (b) synthesizing new data for domain augmentation. We showed that our model significantly outperformed the baseline and other approaches on data from unseen domain in terms of accuracy and various bias measures, without retraining or finetuning. Our method is generally applicable to other biomedical data, providing new algorithms for training models robust to bias for big data analysis and applications. Demo training code is publicly available.

研究动机与目标

• 为解决多机构生物医学数据集中固有数据集偏差导致的模型泛化能力差的问题。
• 开发一种领域泛化框架，提升模型鲁棒性，且在训练期间无需访问外部测试领域。
• 减轻医院特定协议、成像设备和标注偏差等混杂因素对医学影像数据集的影响。
• 创建一种可泛化的方法，适用于除胸部X光外的多种生物医学数据，包括电子健康记录（EHRs）和其他模态。
• 证明仅使用源域数据，通过有效的正则化和数据增强，即可显著提升模型在未见领域上的性能。

提出的方法

• 引入一种偏差正则化损失函数，通过惩罚领域特定的特征激活，促使模型学习领域不变的表征。
• 通过多层交叉梯度训练（MCT）实现领域引导的数据增强，利用来自多个源域的梯度合成多样化、领域鲁棒的训练样本。
• 采用多源训练设置，使用多个数据集（如NIH、CheXpert、MIMIC）以模拟领域偏移并提升泛化能力。
• 应用梯度引导的激活可视化（Grad-CAM）分析注意力机制，验证模型关注的是疾病相关区域而非虚假特征。
• 训练一个独立的偏差分类器以检测并正则化领域特定特征，确保主分类器学习到鲁棒且可泛化的表征。
• 采用两阶段训练流程：首先训练一个去偏差的视觉世界分类器；其次，利用该分类器的交叉梯度信号训练主模型，以促进特征不变性。

实验结果

研究问题

• RQ1 在未见领域上，仅在多个有偏差的生物医学数据集上进行训练的深度学习模型，能否在不进行微调的情况下实现有效泛化？
• RQ2 偏差正则化学习在减少对胸部X光图像中便携式扫描标签或胸管等虚假特征的依赖方面，效果如何？
• RQ3 领域引导的数据增强在异质医学影像领域中，对特征表征学习的提升程度如何？
• RQ4 与标准领域泛化基线相比，所提出方法是否能在未见数据集上实现更优性能？
• RQ5 使用所提出方法时，模型的注意力机制（通过Grad-CAM可视化）是否与放射科医生标注的病灶区域更一致？

主要发现

• 在Office-Caltech数据集上，该方法测试准确率达到47.8%，优于基线LeNet（41.6%）及其他SOTA方法如DANN（43.8%）和CrossGrad（44.3%）。
• 在Rotated-MNIST数据集上，该方法达到99.6%的准确率，显著优于CCSA基线（98.0%）及其他SOTA方法。
• 在Google Fonts数据集上，该方法达到73.1%的准确率，超过基线（68.5%）及其他方法，证明其在非医学领域泛化场景下的有效性。
• Grad-CAM可视化显示，模型与放射科医生标注的病灶区域具有强相关性（Spearman等级相关系数），表明其对疾病相关特征的关注得到改善。
• t-SNE可视化结果表明，所提模型在不同数据集（CheXpert、MIMIC）上的特征嵌入更加混合，分离度更低，表明领域特定偏差减少。
• 一个简单的CNN在分类胸部X光源数据集时达到接近100%的准确率，证实了机构间存在强烈且可检测的数据集偏差。