[论文解读] Learning to Unlearn: Building Immunity to Dataset Bias in Medical Imaging Studies

Ahmed Ashraf, Shehroz S. Khan+3 more|arXiv (Cornell University)|Dec 3, 2018

Domain Adaptation and Few-Shot Learning参考文献 8被引用 19

一句话总结

本文提出一种框架，通过将数据映射到一个隐藏空间，使研究所属关系在该空间中无法区分，同时保持对MCI与健康等临床变量的分类性能，从而消除医学影像中的数据集特异性偏差。该方法相较于简单的留一数据集外方法平均AUC提升47个百分点，并将性能下降幅度控制在20%以内，优于同数据集基线模型。

ABSTRACT

Medical imaging machine learning algorithms are usually evaluated on a single dataset. Although training and testing are performed on different subsets of the dataset, models built on one study show limited capability to generalize to other studies. While database bias has been recognized as a serious problem in the computer vision community, it has remained largely unnoticed in medical imaging research. Transfer learning thus remains confined to the re-use of feature representations requiring re-training on the new dataset. As a result, machine learning models do not generalize even when trained on imaging datasets that were captured to study the same variable of interest. The ability to transfer knowledge gleaned from one study to another, without the need for re-training, if possible, would provide reassurance that the models are learning knowledge fundamental to the problem under study instead of latching onto the idiosyncracies of a dataset. In this paper, we situate the problem of dataset bias in the context of medical imaging studies. We show empirical evidence that such a problem exists in medical datasets. We then present a framework to unlearn study membership as a means to handle the problem of database bias. Our main idea is to take the data from the original feature space to an intermediate space where the data points are indistinguishable in terms of which study they come from, while maintaining the recognition capability with respect to the variable of interest. This will promote models which learn the more general properties of the etiology under study instead of aligning to dataset-specific peculiarities. Essentially, our proposed model learns to unlearn the dataset bias.

研究动机与目标

• 解决医学影像中的数据集偏差问题，即模型过度拟合于研究特有的特征而非学习可泛化的特征。
• 实现在不针对目标数据集重新训练的前提下实现跨数据集泛化，促进学习到根本性疾病相关模式的模型。
• 将原本用于种族或性别等敏感属性的公平表征学习技术，适配用于混淆医学影像中的数据集所属关系。
• 使用四个人脑MRI数据集进行MCI与健康对照分类的留一数据集外策略评估该方法。
• 证明消除数据集所属关系可带来更鲁棒、更具泛化能力的模型，从而最小化跨数据集的性能下降。

提出的方法

• 通过带有基于原型正则化组件的深度自编码器将输入特征映射到隐藏空间，以强制不同数据集间的数据点在表征上不可区分。
• 通过最小化分类器预测源数据集的能力，同时最大化临床目标（如MCI与健康）的分类准确率，学习领域不变表征。
• 目标函数结合了重建损失、分类损失（Softmax）和领域判别损失，以促使模型消除对数据集所属关系的记忆。
• 框架采用多头结构：一个头用于分类，另一个头用于领域判别，后者通过对抗方式训练，使其难以被隐藏表征欺骗。
• 通过嵌套的留一数据集外交叉验证调整超参数，包括权衡权重（αJ, αE, αL）和L2正则化，以优化ROC-AUC。
• 采用留一数据集外协议进行评估，确保训练过程中不使用任何目标数据，并通过AUC及与同数据集基线相比的相对性能下降来衡量性能。

实验结果

研究问题

• RQ1 能否训练一个深度学习模型，在同时分类临床结果（如MCI与健康）的同时，消除输入图像所属数据集的信息？
• RQ2 与标准的留一数据集外方法和无偏SVM基线相比，消除数据集所属关系是否能带来更好的跨数据集泛化性能？
• RQ3 在所提方法下，跨数据集模型的性能下降与同数据集基线相比有多大？
• RQ4 不同数据集（如ADNI1、ADNI2、AIBL）在其特有的特征印记上存在何种差异，这如何影响泛化能力？
• RQ5 能否将源自公平表征学习的框架有效重用于缓解医学影像中的数据集偏差？

主要发现

• 所提方法在所有留一数据集外实验中平均AUC达到0.72，显著优于简单的留一数据集外基线方法，后者在所有数据集上表现均低于随机水平（AUC < 0.5）。
• 与简单的留一数据集外基线相比，AUC平均提升47个百分点；与无偏SVM基线相比，提升18个百分点。
• 与同数据集分类相比，所有数据集的性能下降均低于20%，表明具备强大的泛化能力。
• ADNI1和ADNI2在‘识别研究’任务中更易被识别，表明其具有更高的数据集特异性特征；而AIBL与其他数据集更相似，因此泛化能力更强。
• 该方法在所有数据集上均持续优于两种基线，证明对数据集特异性差异具有鲁棒性。
• 结果证实，消除数据集所属关系可使模型学习到更具泛化能力、与病因学相关的特征，而非数据集特有的伪影。