SSiT: Saliency-guided Self-supervised Image Transformer for Diabetic Retinopathy Grading

TL;DR

Saliency-guided Self-Supervised image Transformer (SSiT) is proposed for Diabetic Retinopathy (DR) grading from fundus images, and significantly outperforms other representative state-of-the-art SSL methods on all downstream datasets and under various evaluation settings.

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%.

Figures (4)

• Figure 1: The proposed SSiT framework. There are two learning objectives, a saliency-guided contrastive loss and a saliency segmentation loss. First, different augmentations of the same input image are fed into two encoders, a query encoder and a key encoder. The two encoders have the same architecture, and the parameters of the key encoder are updated by the moving average of the learnable query encoder. Salient patches selection is performed for the input sequence of the key encoder to remove trivial patches, thus constraining the key encoder to provide target representations focusing on salient regions. The query encoder is optimized by maximizing the similarity between the output representation and the target. In addition, the query encoder is also trained to directly predict saliency segmentation to learn fine-grained information from fundus images.
• Figure 2: Saliency map comparison from two static detection methods.
• Figure 3: Visualization of saliency map segmentation results on representative DDR samples.
• Figure 4: Self-attention maps from different self-supervised ViTs with an input resolution of 1024$\times$1024 and a patch size of 16. $^{\dag}$ MoCo-v3 and MAE fail probably because the input resolution changes sharply.