[论文解读] Multi-task UNet: Jointly Boosting Saliency Prediction and Disease Classification on Chest X-ray Images

Hongzhi Zhu, Robert Rohling+1 more|arXiv (Cornell University)|Feb 15, 2022

Visual Attention and Saliency Detection被引用 6

一句话总结

本文提出了一种多任务UNet架构，通过深度学习联合执行胸部X光（CXR）图像的显著性预测与疾病分类。通过利用共享特征学习并采用优化的多任务训练方案以减轻过拟合，该模型在两项任务上均达到最先进性能，优于专门用于显著性预测或分类的模型。

ABSTRACT

Human visual attention has recently shown its distinct capability in boosting machine learning models. However, studies that aim to facilitate medical tasks with human visual attention are still scarce. To support the use of visual attention, this paper describes a novel deep learning model for visual saliency prediction on chest X-ray (CXR) images. To cope with data deficiency, we exploit the multi-task learning method and tackles disease classification on CXR simultaneously. For a more robust training process, we propose a further optimized multi-task learning scheme to better handle model overfitting. Experiments show our proposed deep learning model with our new learning scheme can outperform existing methods dedicated either for saliency prediction or image classification. The code used in this paper is available at https://github.com/hz-zhu/MT-UNet.

研究动机与目标

• 通过利用多任务学习改善表征学习，以应对医学影像中的数据稀缺问题。
• 开发一种统一的深度学习框架，联合预测人类视觉显著性并分类CXR图像中的疾病。
• 通过优化多任务训练方案以减少过拟合，提升模型的鲁棒性与泛化能力。
• 在CXR数据集上，超越仅专注于显著性预测或疾病分类的现有最先进模型。
• 通过将人类视觉注意力整合到深度学习流程中，实现医学影像任务的可迁移且高效的训练。

提出的方法

• 设计一种多任务UNet架构，使显著性预测与疾病分类分支共享编码器特征。
• 采用基于U-Net架构的共享编码器，从CXR图像中提取分层特征。
• 实现两个并行解码头：一个用于预测2D显著性热力图，另一个用于多类别疾病分类。
• 引入一种优化的多任务学习方案，采用动态损失加权以平衡梯度并减少训练过程中的过拟合。
• 应用数据增强与正则化技术，以提升在有限医疗数据集上的泛化能力。
• 使用组合损失函数端到端训练模型，该损失函数结合了显著性的二元交叉熵与分类任务的分类交叉熵。

实验结果

研究问题

• RQ1 与单任务模型相比，联合学习显著性预测与疾病分类是否能提升两项任务的性能？
• RQ2 在低数据量的医学影像场景中，共享特征的多任务学习在表征学习方面有何增强作用？
• RQ3 所提出的优化多任务训练方案在小型CXR数据集上在多大程度上减少了过拟合？
• RQ4 通过显著性预测引入人类视觉注意力，是否能提升疾病分类模型的可解释性与性能？
• RQ5 统一架构是否能在CXR图像的显著性预测与疾病分类任务中均优于专用模型？

主要发现

• 所提出的多任务UNet在显著性预测与疾病分类任务上均优于现有最先进模型。
• 优化的多任务学习方案显著减少了过拟合，尤其在小型医疗数据集上，提升了泛化能力。
• 联合训练带来了更优的特征学习，表现为与单任务基线相比，显著性与分类指标性能均得到提升。
• 该模型在疾病分类任务上取得了更高的AUC分数，并且与人工标注的显著性图具有更好的相关性，优于现有方法。
• 消融实验证实，共享特征学习与动态损失加权对性能提升至关重要。
• 代码与训练好的模型已公开发布于GitHub，支持可复现性，并推动医学多任务学习的进一步研究。