Artifact-Tolerant Clustering-Guided Contrastive Embedding Learning for Ophthalmic Images

Representation learning in the medical domain is concerned with extracting useful information which help to guide diagnostic decision-making. Common medical image modalities include X-ray, magnetic resonance imaging (MRI), functional MRI, computed tomography, ultrasound scanning, optical coherence tomography, etc. [18]. There are two broad paths of research on this topic: self-supervised or unsupervised visual representation learning and supervised visual representation learning [19]. For self-supervised methods, pretext tasks such as image semantic reconstruction and pseudo label prediction are often used in order to achieve meaningful representations [20]. For example, Bai et al. [21] proposed anatomical position prediction as a pretext task for cardiac MRI segmentation. Prakash et al. [22] adopted image denoising as a pretext task for nuclei images’ segmentation. Li et al. [23] combined two colorization based pretext tasks into a single multi-tasking framework called ColorMe to learn useful representations from scopy images. These pretext tasks can be roughly categorized into predictive, generative, contrastive and multi-tasking according to their working mechanisms [20]. For example, generative methods normally learn to generate medical images using a generative adversarial network [19], while multi-tasking methods seek to integrate various pretext tasks to co-optimize the representation learning [23]. For supervised methods, representations are trained to explain the respective labels such as the binary class of benign and malignant tumors using the end-to-end task training. Popular tasks [24] include detection and severity classification of tumors, skin lesions, colon cancer, blood cancer, etc., achieved by various convolutional neural networks (CNN) architectures such as VGGNet [25] and U-Net [26].

Because the acquisition of quality annotations of medical images is expensive and depends heavily on clinicians’ experience and knowledge, unsupervised representation learning has gained growing attention, especially in the current wave of research on generative and contrastive representation learning [27]. Therefore, we adopt an unsupervised representation learning method in this work.

A recent breakthrough in contrastive learning has shed light on the potential of discriminative model for representation learning [28]. Contrastive learning aims at learning to compare diverse augmented images by optimizing the Noise-Contrastive Estimation loss function, where the positive pair is usually formed with two augmented views of the same image, while negative ones are formed with different images [20]. SimCLR [28] learns representations by maximizing agreement between differently augmented views of the same data example at the instance level. Contrastive clustering [29] extends SimCLR to perform both instance and clustering level contrastive learning, where soft labels of instances are utilized to regularize affinities between samples. To address the reliance of most contrastive learning methods on a large number of explicit pairwise comparisons, SwAV [30] simultaneously clusters the data while enforcing consistency between cluster assignments produced for different augmentations. In a different way, InterCLR [31] uses a memory bank which stores running average features of all samples in the dataset computed in previous steps to reduce the limitation of small batch size. InterCLR additionally captures the inter-image invariance using k-means clustering labels updated in mini-batch, where a positive sample for an input image is selected within the same cluster, while the negative samples are obtained from other clusters.

We adopt a similar learning paradigm as InterCLR to capture both the inter- and intra-invariance of ophthalmic images. However, we update clusters of images after every epoch which are thereafter used for guiding the contrastive learning in the following epoch. Since we use an artifact correction module as the backbone of EyeLearn to constrain the representation learning to be gradually optimized over the epochs, the clustering will become stable with the training to provide reliable guidance for contrastive learning.

Eye diseases such as glaucoma and diabetic retinopathy are chronic disorders that may result in irreversible blindness [1]. The past decade has seen the rise of AI methods to support diagnosis of eye diseases [12, 5]. For example, Wang al et. [3] used RNFLT maps (221 samples) to predict glaucoma with a CNN model. Nayak et al. [32] propose an evolutionary convolutional network model to predict glaucoma in fundus images (1,426 samples). An et al. [2] used both the RNFLT maps and fundus images (357 samples) to predict glaucoma with a transfer learning-based CNN model. Since some ocular diseases are associated with functional vision loss, some work has attempted to predict the respective visual field damage using the ophthalmic image as input. For example, Lazaridis et al. [7] used RNFLT maps (954 samples) to predict the visual fields with a multi-input CNN and a multi-channel variational autoencoder, respectively. Christopher et al. [6] used both RNFLT maps and optic nerve head en face images (1,909 samples) to predict visual field damage with a transfer learning-based ResNet model. However, these methods use relatively small sample sizes for model training and evaluation, which tend to generate over-fitting models and biased evaluation outcomes. In addition, they assume that the ophthalmic images are noise-free or simply exclude the noisy samples before training the models, which therefore cannot handle images containing artifacts.

The essence of AI-based diagnosis is to learn good feature representation revealing potential disease biomarkers in the ophthalmic image [14, 15]. However, large image variations due to varying image artifact and patient-specific retinal anatomy make many standard deep learning methods suboptimal at learning quality representations. To mitigate this deficiency, we propose to correct the inclusive image artifact while learning its representation. We also introduce memory bank-supported contrastive learning to encourage distinguishable representations between images. Unlike many supervised methods which require a massive quantity of image annotations, our method is fully self-supervised which is more generalizable for extracting features for various types of ophthalmic images.

Without loss of generality, we consider that each ophthalmic image is represented as a two-dimensional (2D) image, e.g., as a RNFLT map. Let $\mathbf{D}\in\mathbf{R}^{|\mathbf{D}|\times 0pt\times 0pt}$ be a collection of $|\mathbf{D}|$ 2D ophthalmic images, where $0pt$ and $0pt$ are the width and height of the image, respectively. The embedding learning aims to project each image $\mathbf{X}\in\mathbf{D}$ in a $d$-dimensional latent space $\mathcal{H}\in\mathbf{R}^{d}$, i.e., $\mathbf{h}=f(\mathbf{X})$, where $f$ is the projection function defined as the embedding learning model EyeLearn (refer to Fig. 2 for illustration) and $\mathbf{h}\in\mathcal{H}$ is the vector representation of $\mathbf{X}$ with preserved clinical information conveyed by the image. In addition, we seek to improve embedding learning performance by preserving the relative affinities between ophthalmic images through the contrastive learning-based regularization in EyeLearn.

Partial convolution [33] is proposed to handle images which include missing pixels (e.g., holes). Let $\mathbf{X}_{i,j}$ be the pixels within a convolution window at the location $(i,j)$ and $\mathbf{M}_{i,j}$ be the binary mask with valid pixels being all ones and invalid pixels being all zeros. The partial convolution at location $(i,j)$ is defined as:

with

where $\mathinner{\!\left\lVert\cdot\right\rVert}_{1}$ is the L1 norm, $\odot$ is the element-wise multiplication, $\mathbf{W}$ is the weight matrix and $b$ is the bias. $r_{i,j}$ is a scaling factor to adjust for the varying number of valid input pixels within the convolution window, where $\mathbf{I}_{i,j}$ is a one-like matrix with the same shape as $\mathbf{M}_{i,j}$. We can observe from Eq. 1 that the output at every location depends only on the valid input pixels within the respective convolution window. This is achieved by maintaining a binary mask $\mathbf{M}_{i,j}$ for the feature map at each convolution layer. After a partial convolution layer $l$, $\mathbf{M}_{i,j}$ at location ($i,j$) for next layer ($l+1$) is updated as:

meaning that if there exists at least one valid input pixel within the convolution window that contributes to the output $x^{\prime}_{i,j}$, the location ($i,j$) is a valid pixel for the next convolution layer; otherwise, it is an invalid pixel.

During the contrastive learning process, the negative samples are normally constructed from samples within a mini-batch. However, contrastive learning prefers a larger number of negative samples to learn quality representations [31, 28]. One approach would be to set a very large batch size, but this would require impractically large computational resources. As an alternative method, we use a memory bank of size $M$ to store the representations and the clustering labels of most recent $M$ running samples. Formally, a memory bank (i.e., a size-invariant queue) can be represented as $\mathbf{B}=\{\mathbf{V},\mathbf{C}\}$, where $\mathbf{V}=\{\mathbf{h}_{i}\}_{i=1\mathrel{\mathop{\mathchar 58\relax}}M}$ is the up-to-date representations of $M$ ophthalmic images and $\mathbf{C}=\{c_{i}\}_{i=1\mathrel{\mathop{\mathchar 58\relax}}M}$ is the clustering labels of images. In this paper, the image features are updated at the end of each training batch, while the clustering labels are updated at the end of each training epoch.

Ophthalmic images are often affected by artifact or noise which prevents accurate interpretation and diagnosis of eye diseases. We propose to correct artifacts while learning the feature representation of images to address this problem. As shown in Fig. 2, the artifact correlation is implemented by an U-Net like structure [26] comprised of an encoder to learn the representation and a decoder to reconstruct the image from the learned representation. The encoder-decoder learning scheme can provide EyeLearn the potential to learn representation reflecting the corrected version of image without artifacts.

Specifically, each encoder layer is a stack of three successive operation layers, including a partial convolution layer (refer to Section III-B for detail) to extract features from the masked image, a normalization layer to normalize the convolution feature map, and a non-linear activation layer (e.g., ReLU) to transform the feature map. In comparison, each decoder layer additionally includes an up-sampling layer and a concatenation layer before the partial convolution, normalization and non-linear activation layers. The concatenation operation concatenates two feature maps from the respective encoder and decoder layers at the same level. Notably, the last decoder layer’s input contains the concatenation of the original masked image so the model can reuse valid pixels for image reconstruction. Formally, for each image $\mathbf{X}$, the model takes the ground truth image $\mathbf{X}_{gt}$, masked image $\mathbf{X}_{in}$, and binary mask $\mathbf{M}$ (invalid pixels are zeros) as inputs. It then trains to minimize the differences with respect to both valid and invalid locations between the growth truth image $\mathbf{X}_{gt}$ and the decoder-corrected image $\mathbf{X}_{out}$ as:

where $N_{\mathbf{X}_{gt}}$ denotes the number of pixel locations in $\mathbf{X}_{gt}$. In addition, several other losses have been proven useful in characterizing the similarity between the ground truth and predicted images [33] in terms of different aspects of high-level features, which include the perceptual loss $\mathcal{L}_{perceptual}$, the style loss $\mathcal{L}_{style_{out}}$, the style loss $\mathcal{L}_{style_{comp}}$ and the total variation loss $\mathcal{L}_{tv}$. These losses are calculated from the outputs of an ImageNet-pretrained VGG-16 [25] which takes $\mathbf{X}_{gt}$ and $\mathbf{X}_{out}$ as inputs. Formal definitions of these losses have been established in the literature [33]. Taken together, the autoencoder training aims to optimize a collective reconstruction loss as:

Five weight parameters ($\alpha$, $\beta$, $\gamma$, $\delta$ and $\eta$) are used to balance the importance of different parts. Finally, the encoder output through an average pooling is taken as vector representation $\mathbf{h}\in\mathbf{R}^{d}$ for the input image $\mathbf{X}_{in}$, where $d$ is the representation dimension.

Intra-contrastive learning seeks to capture the intra-image invariance through data augmentations. Each image $\mathbf{X}$ is augmented to generate two versions of variant images $\mathbf{X}_{i}$ and $\mathbf{X}_{j}$ which are then enforced to learn similar representations compared to other augmented images, i.e., they form a positive pair. In this paper, we adopt five different augmentation methods to randomly transform the image, including center crop, rotation, zooming, width shift and height shift. We design these augmentation types to account for the fact that ophthalmic image variances may be generated from varying conditions (e.g., distance of the eye to the lens and rotation of eyes) of patients when the images were taken.

In addition, for each augmented image (e.g., $\mathbf{X}_{i}$ or $\mathbf{X}_{j}$), $N$ images are chosen to form negative pairs. Instead of constructing negative samples within each mini-batch as many previous works do [28, 29], we manage a memory bank (Fig. 2) implemented as a queue which stores the most recent representations of $M$ ($M>N$) images computed from the latest mini-batches. Then, we randomly sample $N$ images from the memory bank to form negative pairs for each augmented image. Following previous work [28], we compute the normalized temperature-scaled cross entropy (NT-Xent) as the contrastive loss for each pair of positive examples ($i,j$) after a projection head:

where $\tau$ is a temperature parameter. $Pro(\cdot)$ is the projection head implemented as a three-layer multilayer perceptron (MLP).

While intra-contrastive learning can capture the feature invariance between multiple augmented versions of an image, inter-contrastive learning seeks to explicitly model the relative affinities between different images. To this end, we adopt a clustering-guided image sampling strategy for constructing positive and negative samples. As shown in Fig. 2, we organize samples in clusters by performing the clustering (i.e., k-means) after every epoch. Therefore, images will be assigned with respective clustering labels dynamically. Then, the newly updated labels after each epoch will be immediately reflected in the memory bank for the next epoch. For image $\mathbf{X}$ in cluster $c_{i}$, we randomly sample an image $\mathbf{X}^{\prime}$ with label $c_{i}$ from the memory bank to form a positive pair with $c_{i}$. Then, we randomly sample $N$ images from other clusters to form negative pairs. Therefore, inter-contrastive learning tries to minimize the difference between images within the same clusters while encouraging different representation learning for images in different clusters. Finally, the inter-contrastive loss based on NT-Xent for each pair of positive ($u,v$) samples and $N$ pairs of negative samples $\{(u,n)\}_{n=1\mathrel{\mathop{\mathchar 58\relax}}N}$ is calculated as:

Since we use the artifact correction module in EyeLearn as the backbone for learning representations of images, the clustering based on the learned representations will become stable with training, thus providing reliable guidance for contrastive sampling purpose.

In EyeLearn, the artifact correction-based embedding learning and the contrastive learning-based regularization share the same learning core (i.e., the encoder), which collectively optimize the model to learn optimal representations of ophthalmic images under the combined objective as:

where $|\mathbf{D}|$ denotes the number of training samples, $w_{1}$ and $w_{2}$ are weight parameters used to balance the contributions of intra-contrastive and inter-contrastive learning, respectively. The entire training and optimizing procedure of EyeLearn is summarized in Algorithm 1.

The dataset includes 30,953 RNFLT maps (each containing $225\times 225$ pixels) from 11,284 unique patients who received a clinical diagnosis between 2010 and 2021. There are numbers of 14,730 left-eye RNFLT maps and 16,223 right-eye RNFLT maps. All left-eye RNFLT maps were horizontally flipped to the right-eye format. Pretraining of the EyeLearn model was done using 22,953 RNFLT maps; another 2,000 RNFLT maps were used for parameter selection. The remaining 6,000 RNFLT maps were used for evaluating the glaucoma detection and visual field mean deviation (MD) prediction.

Glaucoma detection and visual field prediction with deep learning methods have received increasing attention [34, 12], but few of existing methods focus on the unsupervised feature representation learning. We choose to compare the following high-performing methods that are relevant to our method for unsupervised representation learning of RNFLT maps:

• •

  RPCA [35]. Robust principal component analysis (PCA) is a modified version of the widely used PCA method which works well with respect to noisy or corrupted data.

• •

  DAE [36]. Denoising autoencoder (AE) extends from the AE by changing the reconstruction criterion. It is effective at learning representations of noisy or corrupted data.

• •

  SimCLR [28]. This is a contrastive learning method which performs data augmentations and forces similar or dissimilar images in order to learn similar or dissimilar representations.

• •

  EyeLearn. This is our proposed artifact-corrected representation learning combined with intra-contrastive and inter-contrastive learning-based regularization.

We perform glaucoma detection and visual field detection based on the learned representations by the above methods. We further compare EyeLearn with several state-of-the-art visual representation learning methods in two categories which are designed for end-to-end supervised classification training, including the vision transformer which has been used for relevant glaucoma detection tasks [37]:

• •

  Vision Transformer (ViT). This is a type of deep neural network mainly based on the self-attention mechanism which has achieved the state of the art in many image learning tasks. Here, we include several recent strong transformers, CCT [38], ViT [39] and CrossVit [40], for comparison.

• •

  Vision Graph Neural Networks (GNN). This is an recent vision feature learning model which models an image as a graph of small patches and uses GNN to capture the semantic relevance between neighborhood batches [41].

Because vision transformer and vision GNN are used for supervised image classification, we compare their performance through the binary glaucoma detection task in the experiment.

For glaucoma detection, we train a linear support vector classifier (SVC) based on learned representations and adopt the Accuracy (Acc) and F1 score as evaluation metrics. For visual field (i.e., MD) prediction, we train a ridge regression model and adopt the mean absolute error (MAE) and R-squared score (R²) as evaluation metrics. We repeat the evaluations 40 times with random partitions of various ratios of training data. The mean performances and standard deviations are reported.

We train the EyeLearn model 80 epochs with a batch size of 4 and adopt the following parameter settings. The learning rate for the Adam optimizer is 0.0002. In Eq. 6, $\alpha$, $\beta$, $\gamma$, $\delta$ and $\eta$ are set as 6, 0.05, 1, 1 and 0.1, respectively. In Eq. 9, $w_{1}$ and $w_{2}$ are set as 0.002 and 0.001, respectively. The representation dimension $d$ is set as 512. The number of image clusters is set as 7. The bank size $M$ and number of negative samples $N$ are set as 800 and 16, respectively. Detailed parameter settings are referred to our implementation of the EyeLearn model¹¹1https://github.com/codesharea/EyeLearn..

In this section, we first report the embedding performance of various baseline methods by the downstream visual field prediction and glaucoma detection tasks. Then, we examine the contributions of various learning components in EyeLearn through the ablation study. Finally, we compare the performance at various proportions of artifacts in the RNFLT maps.

We examine the impact of correcting image artifact with the example images in three categories $\{A,B\}$, $\{C,D\}$ and $\{E,F\}$ (Fig. 4). We can observe that EyeLearn can impute visually meaningful pixels at the artifact regions. For example, the lower thickness bundles in A and B are missing, which are roughly predicted by our model. We further obtain the pairwise Pearson correlations between images in the second row (i.e., calculated based on the raw features in the image) and third rows (i.e., calculated based on the learned representations by EyeLearn). The results are shown in Fig. 5. We can observe that images with artifacts fall in two categories $\{A,B,C,D\}$ and $\{E,F\}$, while EyeLearn tends to recover the three true categories with the artifacts corrected. In addition, after correcting the image artifacts, the pairwise correlations of $\{A,B\}$, $\{C,D\}$ and $\{E,F\}$ become closer, as can been seen by comparing the color maps between Fig. 5 (a) and (b). This means that EyeLearn is able to learn closed representations for similar images and recover their true relationships that have been distorted by artifacts.

We study the impacts of important hyperparameters including the weight parameters $w_{1}$ and $w_{2}$ used in Eq. 9, the number of clusters $C$ in the clustering-guided contrastive learning, and the size of the memory bank $M$. We can observe from Fig. 6 that the variances of these parameters have significant impacts to the model performance. For example, EyeLearn tends to achieve improved embedding performance with a larger memory bank size observed from the trend in Fig. 6 (d).