[논문 리뷰] DRStageNet: Deep Learning for Diabetic Retinopathy Staging from Fundus Images

Yevgeniy Men, Jonathan Fhima외 4명|arXiv (Cornell University)|2023. 12. 22.

Retinal Imaging and Analysis인용 수 5

한 줄 요약

DRStageNet는 다양한 임상 환경에서의 일반화를 향상시키기 위해 다중 소스 도메인 적응 전략을 사용하여 일곱 개인 공개 당뇨성 망막병변(DR) 데이터셋에서 사전 학습된 자기지도 학습 비전 트랜스포머 모델을 미세조정한 모델이다. 고해상도 grad-rollout 해석 가능 히트맵과 함께 DR 단계 분류에서 최신 기술 수준의 성능을 달성하며, 레이블 정확도 부족과 공존 질환에도 불구하고 뛰어난 강건성을 보여준다.

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) 단계 분류 모델의 낮은 일반화 성능 문제를 해결하기 위해.
• 다양한 인구 통계, 영상 촬영 프rotocol, 레이블 품질을 가진 다양한 데이터셋 간의 다중 소스 도메인 미세조정을 활용하여 DR 단계 분류 성능을 향상시키기 위해.
• 회귀 기반 DR 단계 분류에 대해 grad-rollout 방법을 적응시켜 고해상도의 주의 히트맵을 생성함으로써 모델의 해석 가능성을 향상시키기 위해.
• 레이블 정확도 부족과 공존 질환의 영향이 실제 배포 환경에서 모델 성능과 강건성에 미치는 영향을 조사하기 위해.
• 전문의와 비전문가 모두가 정확하고 설명 가능한 예측을 지원할 수 있는 임상적으로 활용 가능한 DR 단계 분류 시스템을 개발하기 위해.

제안 방법

• DRStageNet는 ImageNet-21k에서 자기지도 학습을 통해 사전 학습된 비전 트랜스포머인 DINOv2를 DR 단계 분류를 위한 특징 기반으로 사용한다.
• 모델은 연속적인 ICDR 단계(0–4)를 예측하기 위해 단순한 완전 연결 회귀 헤드를 사용하며, 단계 분류 문제를 순서 관계를 유지하는 회귀 문제로 간주한다.
• 일곱 개인 공개 데이터셋(EYEPACS, IDRiD, DDR 등)을 대상으로 다중 소스 도메인 미세조정(MST) 전략을 적용하여 도메인 일반화 성능을 향상시킨다.
• ICDR 척도에서의 회귀 성능 최적화를 위해 평균 제곱 오차(MSE) 손실을 사용하여 미세조정을 수행한다.
• 고해상도의 클래스에 종속되지 않는 설명 가능 히트맵을 생성하기 위해 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 병변을 모방하고 훈련 데이터에서 희소하게 나타나는 특성이다.
• DDR 데이터셋에서 grad-rollout 설명 가능 히트맵은 실제 병변 세그멘테이션과 강한 공간적 일치를 보였지만 일부 DR 특징는 누락되어 특정 병변에 대한 선택적 주의를 보였음을 시사한다.
• 이중 필드 DFI에서 환자 수준의 ICDR 레이블을 사용하는 DRTiD 데이터셋에서, 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.