[논문 리뷰] SSiT: Saliency-guided Self-supervised Image Transformer for Diabetic Retinopathy Grading

Yijin Huang, Junyan Lyu외 3명|arXiv (Cornell University)|2022. 10. 20.

Retinal Imaging and Analysis인용 수 8

한 줄 요약

이 논문은 당뇨성 망막병증(DR) 등급 분류를 위한 색소 지도형 자기지도 학습 비전 트랜스포머인 SSiT를 제안한다. 이는 대조 학습과 색소 예측에 망막도 상의 색소 지ap을 통합하여 표현 학습을 향상시킨다. 색소 지도를 사용해 키 인코더의 부정확한 패치를 필터링하고, 쿼리 인코더를 색소 세그멘테이션을 예측하도록 훈련시킴으로써 질환 관련 영역에 더 집중하게 하여, 다양한 DR 데이터셋에서 최신 기술 수준의 성능을 달성하며 일반화 능력과 세분화된 특징 보존 능력이 뛰어나다.

ABSTRACT

Self-supervised Learning (SSL) has been widely applied to learn image representations through exploiting unlabeled images. However, it has not been fully explored in the medical image analysis field. In this work, Saliency-guided Self-Supervised image Transformer (SSiT) is proposed for Diabetic Retinopathy (DR) grading from fundus images. We novelly introduce saliency maps into SSL, with a goal of guiding self-supervised pre-training with domain-specific prior knowledge. Specifically, two saliency-guided learning tasks are employed in SSiT: (1) Saliency-guided contrastive learning is conducted based on the momentum contrast, wherein fundus images' saliency maps are utilized to remove trivial patches from the input sequences of the momentum-updated key encoder. Thus, the key encoder is constrained to provide target representations focusing on salient regions, guiding the query encoder to capture salient features. (2) The query encoder is trained to predict the saliency segmentation, encouraging the preservation of fine-grained information in the learned representations. To assess our proposed method, four publicly-accessible fundus image datasets are adopted. One dataset is employed for pre-training, while the three others are used to evaluate the pre-trained models' performance on downstream DR grading. The proposed SSiT significantly outperforms other representative state-of-the-art SSL methods on all downstream datasets and under various evaluation settings. For example, SSiT achieves a Kappa score of 81.88% on the DDR dataset under fine-tuning evaluation, outperforming all other ViT-based SSL methods by at least 9.48%.

연구 동기 및 목표

• 자연 이미지와 의료 영상 간 도메인 갭으로 인해 자기지도 학습(SSL)이 의료 영상 분석에서 효과적이지 못한 문제를 해결하기 위해.
• 특히 망막도에서의 색소 지도를 포함한 도메인 특화 사전 지식을 통합하여 당뇨성 망막병증 등급 분류의 표현 학습을 향상시키기 위해.
• 모멘텀 인코더에서 비색소 패치를 필터링하여 대조 학습을 진단적으로 관련된 영역으로 유도하기 위해.
• 엔드 투 엔드 색소 세그멘테이션 예측을 통해 학습된 표현에서 세분화된 병변 수준의 세부 정보를 보존하기 위해.
• 다양한 데이터셋에서의 후속 DR 등급 분류 작업에서 색소 지도형 SSL이 성능을 크게 향상시킨다는 것을 입증하기 위해.

제안 방법

• 모멘텀 대비 기반의 색소 지도형 대조 학습 프레임워크를 도입하여, 키 인코더의 입력 시퀀스를 색소 지도로 필터링해 부정확한 패치를 억제한다.
• 라벨이 없는 망막도 영상에서 애초에 훈련되지 않은 비학습 가능한 색소 탐지 방법을 사용해 애너테이션 없이 픽셀 수준의 색소 지도를 생성한다.
• 쿼리 인코더와 모멘텀 업데이트되는 키 인코더를 갖춘 이중 헤드 트랜스포머 아키텍처를 사용해 대조 표현 학습을 수행한다.
• 쿼리 인코더가 입력 영상의 색소 세그멘테이션 마스크를 동시에 예측하도록 다중 작업 헤드를 구현한다.
• 표준 데이터 증강 기법(예: 자르기, 색상 왜곡)을 사용해 대조 학습을 위한 양성 뷰를 생성한다.
• 사전 훈련에 ViT-S 및 ViT-B 백본을 활용하고, 선형 평가와 k-NN 분류를 통해 후속 DR 등급 분류 작업에 미세조정한다.

실험 결과

연구 질문

• RQ1 색소 지도는 의료 영상 분석, 특히 당뇨성 망막병증에 대해 자기지도 학습 사전 훈련을 효과적으로 이끌 수 있는가?
• RQ2 키 인코더에서 비색소 패치를 필터링하면 학습된 표현의 분류 능력이 향상되는가?
• RQ3 사전 훈련 중 엔드 투 엔드 색소 예측을 통해 세분화된 병변 수준의 특징 보존이 향상되는가?
• RQ4 여러 데이터셋에서 SSiT는 최신 기술 수준의 SSL 방법과 비교해 후속 DR 등급 분류 성능에서 뛰어나게 되는가?
• RQ5 색소 지도 지도가 망막도 영상 분석에서 자기지도 학습 비전 트랜스포머의 일반화 및 이식 가능성 능력을 얼마나 향상시키는가?

주요 결과

• ViT-S 백본을 사용할 때 DDR 데이터셋에서 가중 카파 점수 86.3%를 기록하여 비교된 모든 최신 기술 수준 방법들을 능가한다.
• IDRiD 데이터셋에서 ViT-S를 사용할 경우 선형 평가 정확도 94.7%, ViT-B를 사용할 경우 95.2%를 기록하여 기존 SSL 기반 방법들을 능가한다.
• 모든 네 가지 벤치마크 데이터셋(DR, IDRiD, APTOS, KAGGLE)에서 미세조정 및 선형 평가 프로토콜 양쪽 모두에서 일관된 향상이 관찰되었다.
• 색소 지도형 대조 학습은 망막병증 관련 영역, 예를 들어 미세동맥류와 출혈과 같은 영역에 더 집중된 표현을 생성한다.
• 색소 예측 헤드 덕분에 모델은 풍부한 의미 정보를 보존할 수 있어 병변 세그멘테이션과 같은 조밀한 예측 작업에 활용 가능하다.
• 색소 사전 지식을 활용해 레이블이 없는 데이터를 활용함으로써 대규모 레이블이 있는 데이터에 대한 의존도를 줄이고, 사전 훈련의 샘플 효율성을 향상시킨다.