[论文解读] DRStageNet: Deep Learning for Diabetic Retinopathy Staging from Fundus Images

Yevgeniy Men, Jonathan Fhima+4 more|arXiv (Cornell University)|Dec 22, 2023

Retinal Imaging and Analysis被引用 5

一句话总结

DRStageNet 是一种自监督视觉变换器模型，通过在七个公开的糖尿病性视网膜病变（DR）数据集上采用多源域自适应策略进行微调，以提升在多样化临床环境中的泛化能力。该模型在 DR 分期任务中实现了最先进性能，生成高分辨率 grad-rollout 可解释性热力图，并在标签不准确和共病情况下表现出强鲁棒性。

ABSTRACT

Diabetic retinopathy (DR) is a prevalent complication of diabetes associated with a significant risk of vision loss. Timely identification is critical to curb vision impairment. Algorithms for DR staging from digital fundus images (DFIs) have been recently proposed. However, models often fail to generalize due to distribution shifts between the source domain on which the model was trained and the target domain where it is deployed. A common and particularly challenging shift is often encountered when the source- and target-domain supports do not fully overlap. In this research, we introduce DRStageNet, a deep learning model designed to mitigate this challenge. We used seven publicly available datasets, comprising a total of 93,534 DFIs that cover a variety of patient demographics, ethnicities, geographic origins and comorbidities. We fine-tune DINOv2, a pretrained model of self-supervised vision transformer, and implement a multi-source domain fine-tuning strategy to enhance generalization performance. We benchmark and demonstrate the superiority of our method to two state-of-the-art benchmarks, including a recently published foundation model. We adapted the grad-rollout method to our regression task in order to provide high-resolution explainability heatmaps. The error analysis showed that 59\% of the main errors had incorrect reference labels. DRStageNet is accessible at URL [upon acceptance of the manuscript].

研究动机与目标

• 为解决因训练数据集与部署数据集之间存在领域差异，导致糖尿病性视网膜病变（DR）分期模型泛化能力差的问题。
• 通过在具有不同人口统计特征、成像协议和标签质量的多样化数据集上应用多源域微调，提升 DR 分期性能。
• 通过将 grad-rollout 方法适配至回归型 DR 分期任务，提升模型可解释性，生成高分辨率注意力热力图。
• 研究标签不准确性和共病对模型性能及实际部署中模型鲁棒性的影响。
• 开发一种临床可用的 DR 分期系统，为专科医生和非专科医生提供准确且可解释的预测。

提出的方法

• DRStageNet 采用 DINOv2 作为特征主干网络，DINOv2 是在 ImageNet-21k 上通过自监督学习预训练的视觉变换器。
• 模型使用简单的全连接回归头预测连续的 ICDR 分期（0–4），将分期视为回归问题，以保留有序关系。
• 在七个公开数据集（如 EYEPACS、IDRiD、DDR）上应用多源域微调（MST）策略，以提升领域泛化能力。
• 微调过程使用均方误差（MSE）损失，以优化 ICDR 尺度上的回归性能。
• 将 grad-rollout 注意力可视化方法适配至回归任务，生成高分辨率、与类别无关的可解释性热力图。
• 在多个目标领域上对模型进行评估，包括使用患者级别标签的数据集（如 DRTiD），以检验其对标注级别差异的鲁棒性。

Figure 1: DRStageNet consists of a DINOv2 [ 19 ] pretrained backbone that is joined with a simple fully connected regression head. Additionally, we utilize a multi-source domain fine-tuning approach by combining seven open source DR datasets. DRStageNet is fine-tuned using the MSE loss.

实验结果

研究问题

• RQ1 多源域微调是否能提升深度学习模型在具有非重叠领域支持的多样化临床数据集上的糖尿病性视网膜病变分期泛化能力？
• RQ2 在分布偏移条件下，与监督预训练或其他预训练策略相比，使用自监督视觉变换器主干网络在 DR 分期中的表现如何？
• RQ3 公共数据集中标签不准确性和共病在多大程度上影响 DR 分期模型的性能和可靠性？
• RQ4 经适配的 grad-rollout 注意力可视化能否为基于回归的 DR 分期预测提供临床有意义的高分辨率解释？
• RQ5 与多分类分类相比，通过引入有序约束将 DR 分期建模为回归任务，是否能提升性能和鲁棒性？

主要发现

• DRStageNet 在跨数据集泛化方面优于两项最先进基准模型，包括一项近期的基础模型，尤其在存在分布偏移的目标领域表现更优。
• 与单源微调相比，多源域微调策略显著提升了泛化性能，减少了对任一单一数据集的过拟合。
• 59% 的主要预测错误可归因于训练数据中错误的参考标签，凸显了标签质量对模型性能的决定性影响。
• 40% 的误分类图像与血管阻塞等共病相关，这些病症与 DR 病变相似，但在训练数据中代表性不足。
• grad-rollout 可解释性热力图在 DDR 数据集中与真实病变分割结果表现出强空间一致性，尽管部分 DR 特征被遗漏，表明模型对特定病理具有选择性关注。
• 在 DRTiD 数据集上（使用双视野 DFI 的患者级别 ICDR 标签），所有模型（包括 DRStageNet）性能均下降，表明未来模型需整合多视野图像。

Figure 2: DRStageNet confusion matrices. For each confusion matrix but EYEPACS, the classification is reported for a model trained on all other datasets. For EYEPACS-test, the EYEPACS-train set and all other datasets are used to train the model.