SISE-PC: Semi-supervised Image Subsampling for Explainable Pathology Classification

SimCLR [3] is a recent model proposed for unsupervised visual representation learning. A minibatch of $n$ images is sampled and applied with pairs of augmentation to produce $2n$ images. These augmented pairs are fed into a feature encoder (Resnet model) followed by a projection layer (Multilinear Perceptron) to obtain an latent vector ${\bf z}$. The training objective is to maximize the agreement of the latent vectors from the same image with augmented views (positive pair), while repulsing from other different images ($2n-1$ negatives). In SimCLR, the similarity between two views is defined using cosine similarity, i.e. $\rm{sim({\bf u,v})}={\bf u}^{T}{\bf v}/||{\bf u}||~||{\bf v}||$. Thus, the loss function for each positive pair of sample $(a,b)$ is defined to be

where $\tau$ refers to the temperature parameter.

The SimCLR model is trained with batch size of 128 for 120 epochs. A batch size of 128 provides 255 negative samples per positive pair from two augmented views. We apply the LAMB optimizer since training with a large batch size may lead to instability using SGD with momentum. Similar to the original SimCLR implementation in [3], we use a linear warmup for the first 10 epochs and decay the learning rate with the cosine decay schedule. Here, the default augmentation strategy from [3] is adopted. Sample augmentations are shown in Figure 2. The model training is stopped at epoch 75 with a contrastive accuracy of 99.77%.

Each OCT image $I$ is resized to [224x224] followed by SimCLR self-supervised latent vector encoding process. The output of the Resnet encoder for each image of size [7x7x512] is subjected to global averaging to extract $J$, which is a [1x512] dimensional encoded vector representation for each image. We further reduce the dimensions of the latent vector representation by subjecting $J$ to PCA-based decomposition, thereby retaining 8 top principal components per image that are found to be most representative. Thus, the data set used for subsequent selective sub-sampling is represented by {$X\in\mathcal{R}^{[n\times d]},Y\in\mathcal{R}^{[n\times 1]}$}, where $X$ depicts sample features with $d=8$ dimensions for $n=83,484$ samples, and $Y$ represents the pathological class labels. Further, cluster-specific sub-sampling methods described below.

With our intention to train a classification model with minimal training labels, we use k-means to cluster $X$ into four categories/clusters with means and standard deviations, respectively, as $[\mu_{k},\Sigma_{k}]\forall k=[1:4]$. Next, we identify uncertain samples as the ones that lie farther away from the cluster centers and are therefore more likely to lie along the classifier decision boundaries. In (2), we compute the probability for a sample $x$ to belong to cluster $k$.

Samples that have normalized probabilities of belonging to any cluster distribution in the uncertain range of [0.4-0.6] can be considered to lie at cluster peripheries. These samples are collected as $(X_{S}\subset X),X_{S}\in\mathcal{R}^{[mxd]}$, where $m<n$. Other methods such as relative z-scores and Silhouette Scores per sample were also evaluated to identify uncertain samples, but the Gaussian distribution-based sample isolation proved most effective in identifying the corner cases per cluster. Next, we use a small set of labelled data ($L=\{X_{l},Y_{l}\}$) with atmost $n_{l}=80$ labelled samples as seed to apply label propagation [6] to all unlabelled samples in $X_{S}$. The goal here is identification of samples that demonstrate high variances for transductive labelling. Thus, uncertain samples that acquire different class labels at different sample runs can be indicative of the uncertain space around decision boundaries. Using the semi-supervised label propagation method shown in Algorithm 1, we identify the most uncertain minimum sub-sampled data set $(X_{T}\subset X_{S}),X_{T}\in\mathcal{R}^{[r\times d]}$, $r<<n$, that needs to be labelled to adequately train a deep learning classifier [7].

In Algorithm 1, label-spreading using rbf-kernel is applied with $\alpha=0.01$ implying that once a label is acquired by a sample in $X_{S}$ it is less likely to be modified again. Thus, we identify the samples that acquire highly variable labels across 10 label spread runs. The minimal sub-sampled data set $X_{T}$ is returned for labelling and classifier training thereafter.

Once the labels for the minimal sub-sampled data set are collected, we use this minimal training data to further fine-tune the SimCLR pre-trained Resnet classifier. Next we apply the Gradcam library and tf_explain libraries in Python [8] to identify primary and secondary features corresponding to ROIs that explain the underlying pathology as shown in Fig. 3. Here, we observe that the feature extraction layers of Resnet focus on the regions between the inner segment layer and retinal pigment epithelium layer in the OCT images to classify pathological vs. normal images. In case this ROI is occluded, the secondary regions that are analyzed for each OCT image include the choroidal regions for macular OCT images.