SatFormer: Saliency-Guided Abnormality-Aware Transformer for Retinal Disease Classification in Fundus Image

1 Introduction

Retinal diseases are the leading cause of vision impairment and irreversible blindness. Several typical retinal diseases, such as diabetic retinopathy (DR), glaucoma, agerelated macular degeneration (AMD) and retinal vein occlusion (RVO), are becoming common in the working-age population worldwide [Bourne et al., 2017]. Early screening and timely proper therapy is crucial to prevent disease progression and potential vision loss. In clinical routines, ophthalmologists identify retinal diseases in fundus images based on the type and number of associated lesion symptoms, such as microaneurysms, haemorrhages, soft exudates and hard exudates [Li et al., 2019; Sun et al., 2021]. However, these pathological abnormal regions in fundus images are usually small in size and scattered over the entire retina, which makes the diagnosis difficult. Therefore, there is a strong demand for fully automatic and accurate retinal disease recognition in fundus images.

∗The corresponding author.

Convolutional neural networks (CNNs) based methods have achieved significant success in retinal disease diagnosis tasks [Wang et al., 2017; Bi et al., 2021]. However, the repeated combination of pooling and down-sampling layers in CNN models inherently results in the loss of detailed information of subtle lesions, which limits the representation of diagnostic features in scattered small pathological regions. In the ophthalmology field, the occurred lesions always determine the diagnostic results of the specific diseases. Therefore, the relatively weak ability of CNNs to retain finegrained lesion semantics may impose more difficulties for discovering diagnostic visual clues related to early symptoms.

To alleviate this limitation, several recent methods [Li et al., 2019; Sun et al., 2021] resort to adopt attention mechanisms or even transformers to explore long-range dependencies and keep more detailed information. But most of them still adopt convnets as main bodies and treat transformers as assisted modules to help encode global context into convolutional representations. As a result, these models still prefer large continuous areas and are difficult to extract multiple diversified discriminative small parts. In order to fully utilize the inherent encoding capabilities of transformers, Yu et al. [2021] explore the applicability of Vision Transformer for the retinal disease classification task. However, this method takes all individual patches into consideration without highlighting semantics of small lesions, which results in a strong bias on the most salient lesion regions while ignoring trivial detail information contained in the subtle phenotypes. Unfortunately, such characteristic may impair the performance of retinal diseases recognition. More specifically, the challenge that limits the adoption of a computer-aided diagnosis tool by the ophthalmologist is, some disease related subtle pathologies such as microaneurysms and exudates are usually ignored. Thus, it is vital to enhance the feature representation of small lesions and gain a comprehensive understanding of features from diverse lesion regions for retinal diseases recognition.

Based on the above discussions, the limitations of existing methods motivate us to design specific architectures to tackle the above challenges. The main contributions of this paper are summarized as follows.

• We are the first to explore an efficient Vision Transformer which adaptively exploits diverse lesion features from a global perspective with enhanced pathological semantics for retinal disease classification, where a comprehensive understanding of different lesions is particularly important.
• We design a saliency enhancement module to mine discriminative lesion semantics and enhance the feature representation of small lesion regions.
• We propose an abnormality-aware attention which explicitly facilitates the dependencies between different lesion regions to drive the model to distinguish the main abnormal regions as well as the small subtle lesions such as microaneurysms and exudates with remarkably reduced computation complexity.

We conduct comprehensive experiments to demonstrate the effectiveness of each proposed components. Our method significantly outperforms prior state-of-the-art methods on multiple datasets.

2 RelatedWorks

In this section, we give a brief review of retinal disease assessment based on CNNs and Vision Transformers.

CNN-based methods. In the past decade, convolutional neural networks serve as the main design paradigm for retinal disease screening tasks. For instance, Lin et al. [2018] and Li et al. [2019] utilized CNN architectures for DR grading. Phene et al. [2019] developed methods based on CNNs for glaucoma diagnosis. Though these methods have made some progress, they are still limited by insufficient representation for the tiny and subtle lesion regions in deep layers due to severe loss of spatial details caused by down-sampling operations (pooling). As a result, these methods ignore trivial lesion information. Moreover, CNN-based models inherently lack the ability to explicitly capture long-range feature dependencies between different abnormal symptoms in a global scope, which would impair the performance of disease diagnosis.

Transformer-based methods. ViT [Dosovitskiy et al., 2020] first proves that a pure transformer can achieve stateof-the-art performance in image classification with sufficient training data. For retinal disease classification, Yu et al. [2021] introduce a multiple instance learning head on ViT to fully exploit the features extracted from individual patches. However, this method directly utilizes ViT without explicitly exploiting the complex dependencies among diverse lesions. Different from natural images, the diagnostic features in a fundus image, usually only occupy a small part of the whole image. Thus modeling long-range dependencies across all patches is inefficient and results in less discriminative feature representations for subtle lesion regions. LAT [Sun et al., 2021] alleviates this limitation by introducing an encoderdecoder transformer framework to learn lesion-aware filters for DR grading. Nonetheless, this method still employs convnets as main feature extractors, on top of which a transformer is further applied to exploit pixel relation. Since spatial information and fine-grained details may have been lost in deep convolution layers due to down-sampling, the advantages of a transformer is not fully exploited. So far, how to exploit pathological features of retinal diseases distributed at different positions and build an optimal network structure still remains an open question.

3 Methodology

The architecture of SatFormer is illustrated in Figure 1. Similar to [Liu et al., 2021; Wang et al., 2021b], SatFormer also produces a hierarchical representation, which has four stages. Each stage consists of a saliency enhancement module (SEM) and Li sequential SatFormer blocks. A SEM gradually aggregates lesion pixels and highlights potentially salient lesion regions to generate expressive multi-scale embeddings. Then, several SatFormer blocks, each of which involves abnormality-aware attention, are set up after SEM. The image classification is performed by applying a global average pooling layer on the output feature maps of the last stage, followed by a linear classifier.

3.1 Saliency Enhancement Module

For retinal diseases, it is challenging to distinguish the scattered subtle lesion regions since pure transformer architectures lack the inductive bias, such as locality and translation equivariance. Moreover, the coarse splitting of patches limits the ability to model details within each patch [Xu et al., 2021]. Such deteriorated local details compromise the discovery of inconspicuous lesion information. Some recent methods [Li et al., 2021; Wu et al., 2021a] incorporate with convolution operations to bring locality to Vision Transformers. However, without highlighting potentially salient regions from a global perspective, these hybrid models tend to be biased towards the most discriminative lesion regions, while ignoring the diversity of lesion information. Unfortunately, the trivial or less discriminative lesion regions contained in a fundus image may be important for disease recognition.

To overcome the above limitations, we propose a saliency enhancement module (SEM) to mine more explicit lesion semantics and enhance feature activations corresponding to scattered small abnormal regions at each stage. In this process, SEM reduces the number of embeddings and increases their dimensions. The generated feature maps have H 2i+1 × W 2i+1 patch tokens for the ith stage.

Locality Enhanced Multi-Scale Embedding. Figure 2 illustrates how the SEM generates patch embeddings in Stage 1. We first use four convolution kernels with different sizes to sample patches in an input image. Considering that lesion regions usually exhibit significant variations in shape, size, and structure, distinguishing these abnormalities from a variety of different scales helps to enhance the ability for mining the discriminative visual clues related to each lesion. We keep the stride (4×4) of four kernels the same to make sure they generate the same number of embeddings. Then the four patches will be projected and concatenated as one embedding.

To control the total computational budget, we use a lower dimension d for large kernels and successively doubles it for small kernels. This scheme saves a lot of computational cost without significant performance drop. In Stage 2 - 4, the strides are set as 2×2 while kernels are set as 3×3 and 5×5. Since each token embedding already learns information from regions of different scales, our strategy empowers the model with enriched multi-scale local context modeling capabilities.

Spatial Saliency Enhancement. After the embedding process, SEM performs the following procedures. First, given patch embeddings Xp e ∈ RHi×Wi×Ci , a 1 × 1 convolution is conducted to expand the embeddings to a higher dimension of Xd e ∈ RHi×Wi×(2×α×Ci) , where α is the expand ratio, Hi, Wi and Ci denote the height, width and channel dimensions of the input in stage i, respectively. A key ingredient in SEM is the spatial perception layer which adaptively enhances the response of lesions and suppresses cluttered background features. Specifically, we split Xd e into two independent parts (Xd1 e ,Xd2 e ) along the channel dimension. A spatial saliency distribution map S ∈ RHi×Wi×(α×Ci) is generated through feeding Xd2 e to a depthwise convolution layer followed by one GELU and Softmax. In experiments, we found that adding GELU before Softmax, especially in shallow layers, helps stabilize the training process. Our purpose is to generate a spatial weight distribution matrix to highlight the potential salient regions and suppress noisy backgrounds. Therefore, the Softmax normalizes the feature maps along the spatial dimension. For each small lesion, its feature activation on the spatial saliency distribution map S is expected to be higher than that of its surrounding backgrounds. Guided by S, the model learns the context of pathological semantics from a global perspective and highlights potentially salient lesion regions on the entire feature map scale. Furthermore, Xd1 e and S are combined by element-wise multiplication, generating Xs e ∈ RHi×Wi×(α×Ci) , which encodes salient semantics of abnormalities scattered over the entire retina. This procedure can be noted as:

Finally, we project Xs e to the initial dimension and then add the output with Xp e through a residual connection to obtain the output patch embeddings Xo e ∈ RHi×Wi×Ci . A LayerNorm (LN) layer and a GELU is applied following each linear projection and convolution. In this way, SEM empowers the model with enriched feature representations for small lesions and captures diverse pathological semantics.

3.2 Abnormality-Aware AttentionMechanism

Self-attention [Vaswani et al., 2017] has advantages in capturing long-range dependencies, but it incurs huge memory and computation costs. Some methods, e.g., PVT [Wang et al., 2021a], PVTv2 [Wang et al., 2022] and P2T [Wu et al., 2021b] adopt spatial-reduction (pooling operation) to downsample the feature maps when computing keys and values in the multi-head self-attention (MHSA). With the pooled features, they reduce the number of tokens but at the cost of losing fine-grained details. Other methods [Liu et al., 2021; Wang et al., 2021b] resort to computing self-attention within dilated partitions of feature maps but sacrifice the direct global relationship modeling.

Though the aforementioned methods have made some progress, they still suffer from two drawbacks. (1) They fail to adaptively perceive and group specific tokens that contain the discriminative features. (2) They lose the global receptive field and fail to gather complementary information from the less discriminative groups of tokens. As a result, they sacrifice accuracy as a trade-off with efficiency. Moreover, models without comprehensive information on diverse scattered lesions and the dependencies between different pathological features tend to be biased on the most salient lesion but ignore the less discriminating lesions that may be important for diagnosis.

Motivated by the above limitations, we propose a novel abnormality-aware attention mechanism. Specifically, apart from the position embedding, we attach an additional group embedding to each patch token. This group embedding indicates whether the corresponding token encodes useful lesion semantics. Generating group embeddings from a global view is straightforward. As shown in Figure 3, the patch embeddings Xo e ∈ RHi×Wi×Ci from SEM are fed into a 1 × 1 convolutional layer with single channel output. In this way, a token relative importance matrix D is generated via:

where θ and b represent the weights and bias of the convolutional layer, respectively. Dn,m helps to describe the probability that the patch at row n and column m contains lesion information or pathological semantics associated with abnormalities. The group embeddings G are obtained by setting an threshold ε for D. For each patch, if its corresponding Dn.m ≤ ε, then this patch will be assigned to the background group, otherwise it will be assigned to the lesion group. In this way, G enables the model to adaptively perceive and group specific patch tokens that contain the discriminative features.

Instead of focusing on all input tokens, we propose to compute self-attention within the lesion group for efficient modeling. However, the patches in the background group may contain important complementary information that contributes to a thorough understanding of lesions and their surroundings. Thus we propose to add an extra cross-group attention to enable cross-group information exchange. As shown in Figure 3, we flatten all input patch tokens and compute selfattention within the lesion group. In the lesion group, patch tokens Xl ∈ RN2×Ci are linearly transformed to three parts, i.e., queries Qs ∈ RN2×Ck , keys Ks ∈ RN2×Ck and values Vs ∈ RN2×Cv where N2 is the sequence length of the lesion group, C, Ck, Cv are the dimensions of inputs, queries (keys) and values, respectively. The scaled dot-product attention is applied via:

As for the cross-group attention, the queries Qc ∈ RN2×Ck come from the lesion group while keys Kc ∈ RN1×Ck and values Vc ∈ RN1×Cv come from the background group. This allows each patch in the lesion group to focus on the features of all patches Xb ∈ RN1×Ci in the background group to gather complementary information that is highly relevant to lesion diagnosis. The cross-group attention is applied via:

We then add the results of the two applied attention through an element-wise sum operation and employ a linear layer to produce the output ˆXl ∈ RN2×Ci of the lesion group. Finally, patches in the background group are concatenated with ˆXl through a residual connection. We remove the group embeddings and restore the order of the patches according to their original positions. Once we obtain the output of the abnormality-aware attention module, we send it to the MLP block for proceeding computation as usual. At the end of each stage, we restore patch tokens to "images" in the spatial dimension based on the original positions. The computational cost of the abnormality-aware attention is reduced from O(N 2 ) to O(N 2 2 + N1N2), and N2 ≪ N in most cases. Multi-head version of abnormality-aware attention can be easily derived. Just as MHSA, we split the queries, keys and values to h parts and perform the abnormality-aware attention mechanism in parallel. Then we concatenate the output values of each head and linearly project them to form the final output.

Notably, each abnormality-aware attention module computes its own group embeddings. This gives the model an opportunity to rethink how to group and make full use of patches containing lesion semantics to gather as much information as possible. The learning process of group embeddings also benefits from SEM. SEM highlights potentially salient regions of small lesions and provides sufficient context to help the model perceive diverse lesions, which makes it easier and more meaningful to predict whether a patch contains pathological semantics.

Our abnormality-aware attention introduces the inductive bias that the spatial interactions should be dynamically parameterized based on the pathological semantics and the correlations between different lesion features. This characteristic is conducive to the diagnosis of diseases, especially for retinal diseases, where pathological regions are often small and scattered in distribution.

4 Experiments

4.1 Dataset and Implementation

We conduct experiments on a large dataset of fundus images collected from a regional hospital and two public benchmarks including EyePACS [Cuadros and Bresnick, 2009] and RFMiD [Pachade et al., 2021]. The collected dataset contains 28, 360 DR images, 5, 816 glaucoma images, 4, 805 AMD images, 25, 748 RVO images and 35, 369 normal images. We randomly split 70% of the dataset for training, 10% for validation and the rest 20% for testing. We adopt the ten-fold cross-validation method on our dataset. EyePACS [Cuadros and Bresnick, 2009] contains 35, 126 training, 10, 906 validation and 42, 670 testing DR images. Each image is divided into one of five DR grades. Retinal Fundus Multidisease Image Dataset (RFMiD) [Pachade et al., 2021] contains 1, 920 training, 640 validation and 640 testing images, which screens retinal images into normal and abnormal (comprising of 45 different types of diseases) categories.

All experiments are performed on 8 V100 GPUs. We employ an AdamW optimizer for 500 epochs using a cosine decay learning rate scheduler and 20 epochs of linear warm-up. A mini-batch size of 128, an initial learning rate of 0.001, and a weight decay of 0.05 are used. We use gradient clipping with a max norm of 1.0 to stabilize the training process. Augmentations such as rotation, scaling, gaussian blur, color jitting and mirroring are utilized during the training process.

For the collected dataset and EyePACS, as multi-class classification problems, we utilize evaluation metrics of accuracy, area under the curve (AUC), weighted F1 and weighted Kappa following [Yu et al., 2021; Sun et al., 2021].

For RFMiD, which contains only two classes of either disease or normal, the AUC, accuracy and F1 metric are used.

For fair comparison, we reap the benefit of pre-trained weights of MLP layers and QKV attention on ImageNet pre-training. Concretely, we align channel numbers of SatFormer blocks to those of pre-trained models so that we load the weights of MLP layers and QKV attention. In each abnormality-aware attention layer, we reuse the pretrained weights of QsKsVs to initialize QcKcVc. Similar to ViT [Dosovitskiy et al., 2020] and Swin [Liu et al., 2021], we build two architecture variants. The hyper-parameters of these model variants are:

• SatFormer-S: C = 96, block numbers = {2, 2, 8, 2}, H = {3, 6, 12, 12}, d = {16, 64, 128, 256}, ε = 0.5, α = {2, 2, 1, 1}
• SatFormer-B: C = 128, block numbers = {2, 2, 18, 2}, H = {3, 6, 12, 24}, d = {16, 64, 128, 256}, ε = 0.5, α = {2, 2, 1, 1}

where C and H mean embedding dimensions and the number of heads in the multi-head self-attention, respectively. d and α are the kernel dimension and expansion ratio in the SEM, respectively. ε is the threshold for grouping. -S and -B for small and base, respectively.

4.2 Comparisons with State-of-the-art Methods

We compare our method with the state-of-the-arts including representative CNN-based models, transformer-based models and hybrid architectures. The results are shown in Table 1. SatFormer noticeably surpasses the other state-of-theart structures.

Compared against strong Transformer based baselines ViTL/16, PVTv2-B5, P2T-Large, Swin-L, and Twins-SVT-L, SatFormer-B outperforms them at least absolute 2.3% in the weighted Kappa on EyePACS, and 4.2% in F1 on RFMiD. Further, even our SatFormer-S achieves at least 2% improvement of F1 relative to these large models with significantly lower paramerters. Notably, as for RFMiD dataset, the AUC and F1 score of SatFormer-B surpass that of MIL-VT, which is pre-trained on a large fundus dataset, with much lower FLOPs (65.0G vs. 200.0G). This success has been attributed to the proposed abnormality-aware attention module, in which the weights are dynamically computed based on the similarity or affinity between every pair of tokens in the lesion group. As a consequence, the captured global dependencies are potentially more efficient and discriminative than interactions built on all patch tokens.

Compared with the state-of-the-art ConvNets, i.e., EfficientNet and ReXNet, the SatFormer-B achieves at least absolute 1.8% improvement of Kappa on EyePACS and 2.5% improvement of F1 on RFMiD. Moreover, our smallest model SatFormer-S with 27M parameters and 22.5G FLOPs surpasses the CANet by 3.1% of Kappa on EyePACS and 4.5% of F1 score on RFMiD, while CANet has 3 times the FLOPs of SatFormer-S. This effectiveness can be explained that our approach is capable of both perceiving locally low contrast small-sized lesions and capturing dependencies between scattered pathologies globally, while CNN-based models suffer from spatial detail loss and lack the ability to model interactions between lesion features.

Compared to other hybrid models, SatFormer significantly outperforms the other SOTAmodels for all the three datasets. SatFormer-S achieves better performance compared with CvT-21 and CrossFormer-L with fewer parameters and similar FLOPs. Our SatFormer-B surpasses other hybrid architectures by a large margin in terms of all evaluation metrics with acceptable parameters and FLOPs. Different from previous hybrid models, which either simply treat transformers as assisted modules or compute self-attention within dilated partitions of feature maps, our SatFormer adaptively perceives and groups specific tokens to explicitly model the correlation between different lesions, resulting in a more comprehensive understanding of pathological clues.

4.3 Ablation Study

Effectiveness of the Saliency Enhancement Module. Table 2 shows the ablation study on the effectiveness of the SEM. All models achieve more than 1.0% absolute performance gain on the RFMiD dataset by adopting our SEM. This shows applying our SEM yields significant improvements on various architectures without too much additional computational cost. Moreover, we compare the spatial perception layer in SEM with several other close alternatives which also compute a spatial weight distribution matrix. The results in Table 3 show that our spatial perception layer works better than Non-local and CBAM. In practice, Non-local needs to compute a matrix quadratic over the HW while our spatial saliency distribution map is linear over the input squence length. CBAM adopts average pooling to compute spatial weight matrix, which may lose important clues about distinctive object features. These experiments verify the effectiveness of our SEM and show that it generalizes well to various transformer architectures.

Significance of the Abnormality-Aware Attention. Table 4 shows the ablation study on the effectiveness of the abnormality-aware attention. Several self-attention mechanisms used in PVT, Swin, Twins and CrossFormer are compared. Our SatFormer-B outperforms them at least absolute 1.7% Kappa on the EyePACS dataset and achieves at least absolute 2.8% performance gain on the RFMiD dataset. Moreover, Table 2 shows that even without SEM, our SatFormer-B (w/o.) still outperforms other Vison Transformers. This indicates applying our abnormality-aware attention in other Vison Transformers (like Swin-B or Swin-L) also benefits the overall performance. In particular, PVT adopts spatial-reduction to reduce the sequence length and thus sacrifices the finegrained features, while Swin restricts the self-attention in a local window, sacrificing the long-distance attention. Twins and CrossFormer compute self-attention within dilated partitions of feature maps, but do not adaptively group specific tokens to explicitly capture long-range feature dependencies between lesion regions. Compared with simply splitting the feature map into multiple windows in which tokens share the same surroundings, our adaptive patch grouping strategy prevents tremendous background information from overwhelming the subtle visual clues related to pathological abnormalities. In this way, high-level semantic information from different pathological regions are combined to form a comprehensive lesion feature understanding. The results show that our abnormality-aware attention is most conducive to improving the performance.

Importance of the Cross-Group Attention. In Table 5, we show that it is crucial to make use of the cross-group attention, where removing it deteriorates the overall classification performance by over 1%. The underlying reason is that the cross-group information exchange helps the model to learn important supplementary information such as the structure of the entire retina, which is beneficial to a thorough understanding of lesions and their surroundings. Otherwise, the network will not be able to learn useful representations from the tokens in the background group.

4.4 Visualization of AttentionMaps

Several representative visualizations of attention maps in Figure 4 clearly indicate that the scattered pathological regions are well captured and perceived by our SatFormer. On the contrary, the existing methods show a strong bias on the most salient lesion region and ignore trivial lesion information contained in the subtle regions. Moreover, it is interesting to observe that the existing state-of-the-art Vision Transformers still have high responses in irrelevant regions (e.g., Swin and CrossFormer). This phenomenon is consistent with the underlying mechanism of their self-attention functions, which compute dependencies between patches from naive partitions of feature maps. Different from existing methods, we propose to adaptively perceive and group specific tokens containing discriminative features. Our abnormality-aware attention adaptively selects patch tokens and gather useful information to update the chosen tokens instead of all tokens. Only the visualization results of our SatFormer show a thorough comprehension of diverse subtle lesion regions. This phenomenon also reveals that the overall accuracy improvement of SatFormer compared to other models is brought about by capturing trivial or less discriminative lesion regions contained in a fundus image.

5 Conclusions and FutureWork

In this paper, we propose a saliency-guided abnormalityaware transformer for retinal disease classification, which explicitly captures long-range dependencies between scattered subtle lesions from a global perspective. Particularly, we design a saliency enhancement module to enhance the semantics of small pathological regions. Our abnormality-aware attention with an adaptive patch grouping strategy helps the model gain a comprehensive understand of diverse lesions in an efficient way. Experiments show our SatFormer significantly outperforms prior state-of-the-art methods.

SatFormer follows the hierarchical design of popular Vision Transformers (e.g., Swin Transformer and PVT) which generate multi-scale feature maps. In future research, it is worth exploring the capability of SatFormer as a generalpurpose backbone for dense prediction tasks (e.g., object detection and semantic segmentation).

Acknowledgments

This research was partially supported by the Key Research and Development Program of Zhejiang Province under Grant 2021C03032, the National Natural Science Foundation of China under Grant 61972343, and the National Major Scientific Research Instrument Development Project under Grant 81827804.