[논문 리뷰] DRG-Net: Interactive Joint Learning of Multi-lesion Segmentation and Classification for Diabetic Retinopathy Grading

Hasan Tusfiqur, Duy M. H. Nguyen외 9명|arXiv (Cornell University)|2022. 12. 30.

Retinal Imaging and Analysis인용 수 6

한 줄 요약

DRG-Net는 도메인 불변 특징 학습과 주의 메커니즘을 사용하여 다중 병변 세분화와 질병 분류를 동시에 수행하는 상호작용적이고 통합된 학습 프레임워크를 제안한다. 사용자 피드백과 적대적 훈련을 통합함으로써 모델의 강건성과 해석 가능성을 향상시켜, 약한 지도 학습 피드백 하에 IDRID 및 FGADR 데이터셋에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Diabetic Retinopathy (DR) is a leading cause of vision loss in the world, and early DR detection is necessary to prevent vision loss and support an appropriate treatment. In this work, we leverage interactive machine learning and introduce a joint learning framework, termed DRG-Net, to effectively learn both disease grading and multi-lesion segmentation. Our DRG-Net consists of two modules: (i) DRG-AI-System to classify DR Grading, localize lesion areas, and provide visual explanations; (ii) DRG-Expert-Interaction to receive feedback from user-expert and improve the DRG-AI-System. To deal with sparse data, we utilize transfer learning mechanisms to extract invariant feature representations by using Wasserstein distance and adversarial learning-based entropy minimization. Besides, we propose a novel attention strategy at both low- and high-level features to automatically select the most significant lesion information and provide explainable properties. In terms of human interaction, we further develop DRG-Net as a tool that enables expert users to correct the system's predictions, which may then be used to update the system as a whole. Moreover, thanks to the attention mechanism and loss functions constraint between lesion features and classification features, our approach can be robust given a certain level of noise in the feedback of users. We have benchmarked DRG-Net on the two largest DR datasets, i.e., IDRID and FGADR, and compared it to various state-of-the-art deep learning networks. In addition to outperforming other SOTA approaches, DRG-Net is effectively updated using user feedback, even in a weakly-supervised manner.

연구 동기 및 목표

• 임상의와 AI 간의 상호작용적 학습을 가능하게 하여 당뇨병 망막병증(DR) 진단에서 검정 상자 모델의 한계를 해결하고 해석 가능한 학습을 실현한다.
• 공유된 주의 기반 특징을 활용하여 다중 병변 세분화와 DR 등급 분류를 동시에 최적화하여 모델의 일반화 능력과 강건성을 향상시킨다.
• 정밀한 애너테이션에 대한 의존도를 줄이기 위해 사용자가 보조정할 수 있도록 약한 지도 학습 방식(예: 경계 상자)을 허용함으로써, 픽셀 단위의 마스크가 아닌 보다 간편한 피드백 방식을 도입한다.
• 병변 특징과 분류 특징 간의 손실 제약 조건 및 주의 메커니즘을 통해 노이즈가 많거나 정확하지 않은 사용자 피드백에 대한 모델의 강건성을 향상시킨다.
• 전문가의 상호작용을 통해 지속적으로 향상되는 협업형 AI 시스템을 개발하여 최소한의 애너테이션 노력으로도 성능 향상을 이룬다.

제안 방법

• DRG-Net는 이중 모듈 아키텍처를 사용한다: 분류, 병변 지역화 및 해석을 위한 DRG-AI-System과 사용자 피드백 통합을 위한 DRG-Expert-Interaction.
• 도메인 불변 병변 특징 생성기는 전이 학습과 와서슈타인 거리 기반 적대적 학습을 활용하여 도메인 간 차이를 최소화하고 다양한 데이터셋 간 특징 일반화 능력을 향상시킨다.
• 저수준 및 고수준 특징에 모두 새로운 다중 척도 주의 메커니즘이 적용되어 병변 영역의 주목적 특징을 동적으로 강조함으로써 설명 가능성 향상에 기여한다.
• 적대적 학습 기반 엔트로피 최소화 기법을 사용하여 레이블 노이즈에 강건하고 데이터 부족 상황에서도 일반화 능력을 향상시킨다.
• 사용자 피드백은 약한 지도 학습 기반 미세조정을 통해 통합되며, 픽셀 단위의 세분화가 아닌 경계 상자 수정을 통해 모델을 업데이트한다.
• 손실 함수는 병변 특징과 분류 예측 간의 일치를 강제하여 세분화 품질 향상이 등급 분류 성능 향상으로 직접 이어지도록 설계된다.

실험 결과

연구 질문

• RQ1 병변 세분화와 DR 등급 분류를 동시에 학습하는 통합 학습 방식이 별도로 학습하는 것보다 모델 성능 향상에 기여하는가?
• RQ2 예를 들어 경계 상자와 같은 원시적인 애너테이션 형태의 사용자 피드백을 어떻게 효과적으로 약한 지도 학습 방식으로 활용할 수 있는가?
• RQ3 병변과 관련된 특징을 강조하는 주의 메커니즘이 모델의 해석 가능성과 노이즈가 많은 피드백에 대한 강건성에 얼마나 기여하는가?
• RQ4 도메인 불변 특징 학습이 망막 영상에서 도메인 이동 및 제한된 레이블 데이터로 인한 성능 저하를 완화하는 데 얼마나 효과적인가?
• RQ5 활성화 맵, 병변 마스크 등의 임상적 증거 통합이 실제 진료 환경에서 전문가의 신뢰도 및 의사결정 지원에 어떻게 기여하는가?

주요 결과

• DRG-Net는 IDRID 및 FGADR 데이터셋 모두에서 세분화 및 분류 성능 측면에서 최신 기술 수준의 방법들을 능가하며, 데이터 부족 상황에서도 뛰어난 일반화 능력을 입증한다.
• 주의 기반 특징 선택과 손실 제약 조건 덕분에 사용자 피드백 노이즈에 매우 강건하여, 수정이 정확하지 않거나 불완전한 경우에도 성능 유지가 가능하다.
• 병변 주의 기반 통합은 활성화 맵과 진짜 병변 영역 간의 겹침을 크게 향상시켰으며, 그 결과 시각적 비교(그림 17)를 통해 해석 가능성 향상이 확인되었다.
• 경계 상자 수정을 통한 사용자 피드백이 픽셀 수준의 애너테이션 없이도 세분화 및 분류 작업에서 측정 가능한 성능 향상을 이끌어내었다.
• 반복적인 전문가 상호작용을 통해 지속적인 모델 향상이 가능함을 입증하여, 임상 워크플로우에서 협업형 AI 시스템으로서의 실현 가능성을 확인했다.
• 절단 실험 결과, 적대적 학습, 도메인 적응, 주의 메커니즘의 조합이 최적의 성능과 강건성을 달성하는 데 필수적임을 확인했다.