Self-supervised Pseudo Multi-class Pre-training for
Unsupervised Anomaly Detection and Segmentation in Medical Images

Our MedMix augmentation is designed to augment medical images to simulate multiple lesions. We target a more effective data augmentation for MIA applications than the computer vision augmentations in [60] (e.g., permutations, rotations) that do not simulate medical image anomalies and may yield poor detection performance by downstream UAD methods. We realise that anomalies in different medical domains (e.g., glaucoma and colon polyps) can be visually different, but a commonality among anomalies is that they are usually represented by an unusual growth of abnormal tissue. Hence, we propose the MedMix augmentation to simulate abnormal tissue with a strong augmentation that “constructs” abnormal lesions by the cutting and pasting (from and to normal images) of small and visually deformed patches. This visual deformation is achieved by applying other transformations to patches, such as colour jittering, Gaussian noise and non-linear intensity transformations. This approach is inspired by cutmix [67] and CutPaste [27], where our contribution over those approaches is the intensification of the change present in the cropped patches by the appearance transformations above. These transformations are designed to encourage the model to learn abnormalities in terms of localised image appearance, structure, texture, and colour.

In practice, we design $|\mathcal{A}|=4$ strong augmentation distributions, where $\mathcal{A}_{n}$ includes $n\in\{0,...,3\}$ abnormalities in the image, which means that $\mathcal{A}_{0}$ denotes the normal image distribution and $\mathcal{A}_{n\in\{1,2,3\}}$ represent the abnormal image distributions, containing $\{1,2,3\}$ anomalous regions. Therefore, our loss targets the classification of MedMix augmentations, as shown in Fig. 2.

The gist of our proposed PMSACL lies in the idea of discriminating the distribution of weakly augmented samples (simulating normal images) from the distributions of different types of strongly augmented samples (simulating multiple classes of abnormal images). Instead of attracting and repelling samples within and between a large number of image classes [53, 60, 52], we propose a new contrastive loss to separate samples from the normal class and samples from pseudo abnormal sub-classes, and to enforce the clusters representing the normal and abnormal sub-classes to be dense. To this end, our proposed loss is defined as:

$\begin{split}\ell(\mathcal{D};\theta,\beta,\gamma)=\ell_{ctr}(\mathcal{D};\theta)+\ell_{PMSACL}(\mathcal{D};\theta)+\ell_{aug}(\mathcal{D};\beta)+\ell_{pos}(\mathcal{D};\gamma),\end{split}$

(1)

where $\ell_{ctr}(.)$ denotes the new distribution multi-centring loss, $\ell_{PMSACL}(.)$ represents the new PMSACL contrastive loss, $\ell_{aug}(.)$ and $\ell_{pos}(.)$ are the pretext learning losses to regularise the optimisation [60], and $\theta$, $\beta$ and $\gamma$ are trainable parameters. The loss terms in (1) rely on weak data augmentation distribution, denoted by $\mathcal{A}_{0}$, and $|\mathcal{A}|$ strong data augmentation distributions, represented by $\{\mathcal{A}_{n}\}_{n=1}^{|\mathcal{A}|}$, each denoting a different type of augmentation. From each of these distributions, we can sample augmentation functions $a:\mathcal{X}\to\mathcal{X}$.

The multi-centring loss in (1) depends on the estimation of the mean representation for each augmentation distribution, computed as

$$\mathbf{c}_{n}=\mathbb{E}_{\mathbf{x}\in\mathcal{D},a\sim\mathcal{A}_{n}}[f_{\theta}(a(\mathbf{x}))],$$

(2)

where $n\in\{0,...,|\mathcal{A}|\}$, with $\mathbf{c}_{n}$ being the mean representation of the training data augmented by the functions sampled from $\mathcal{A}_{n}$. Note that these mean representations are computed at the beginning of the training and frozen for the rest of the training. The distribution multi-centring loss is then defined as:

$$\ell_{ctr}(\mathcal{D};\theta)=\mathbb{E}_{\mathbf{x}\in\mathcal{D},n\in\{0,...,|\mathcal{A}|\},a\sim\mathcal{A}_{n}}\|f_{\theta}(a(\mathbf{x}))-\mathbf{c}_{n}\|^{2},$$

(3)

which pulls the representations of augmented samples toward their mean representations in (2), making the augmentation clusters dense in Euclidean space.

To further enforce the separation between different clusters and the tightness within each cluster, we introduce a novel contrastive learning loss function. In our contrastive learning, we maximisise the cosine similarity of samples that belong to the same class (i.e., normal or one of the abnormal sub-classes) and minimise the cosine similarity of samples belonging to different classes. An interesting aspect of this optimisation is that samples are centred by their own cluster mean representation $\mathbf{c}_{n}$ from (2), so our contrastive learning, combined with the multi-centred loss in (3) will cluster samples of the same class not only in Euclidean space, but also in inner product space (with cosine measuring similarity between samples). Such re-formulated constrastive learning, combined with the multi-centring loss (3), results in a loss that produces multiple clusters, where cluster $n=0$ contains the normal images and the others, denoted by $n\in\{1,...,|\mathcal{A}|\}$), have the synthetised abnormal images. Our PMSACL loss is defined as:

$\begin{split}\ell_{PMSACL}(\mathcal{D};\theta)=\mathbb{E}_{\mathbf{x}\in\mathcal{D},n\in\{0,...,|\mathcal{A}|\},l\in\{0,1\}}\left[\ell^{x}_{PMSACL}(\mathbf{x}^{(n,l)},\mathcal{D};\theta)\right]\end{split}$

(4)

where $\mathbf{x}^{(n,l)}=a(\mathbf{x}^{(n)})$ represents one of two (indexed by $l\in\{0,1\}$) augmented data obtained from the application of a weak augmentation $a\sim\mathcal{A}_{0}$ on a strongly augmented data denoted by $\mathbf{x}^{(n)}=a(\mathbf{x})$ with $a\sim\mathcal{A}_{n}$. In (4), we have:

$\begin{split}\ell^{x}_{PMSACL}(\mathbf{x}^{(n,l)},\mathcal{D};\theta)=-\log\frac{\exp\left[\frac{1}{\tau}(\mathbf{f}^{(n,l)})^{\top}\mathbf{f}^{(n,(l+1)\,\text{mod}\,{2})}\right]}{\sum\limits_{\begin{subarray}{c}\mathbf{x}_{j}\in\mathcal{D}\\ m\in\{0,...,|\mathcal{A}|\}\\ k\in\{0,1\}\end{subarray}}\mathbb{I}(\mathbf{x}_{j}^{(m,k)}\neq\mathbf{x}^{(n,l)})\exp\left[\kappa(n,m)(\mathbf{f}^{(n,l)})^{\top}\mathbf{f}^{(m,k)})\right]},\end{split}$

(5)

where $\mathbb{I}(.)$ denotes an indicator function, $\mathbf{x}_{j}^{(m,k)}$ is defined similarly as $\mathbf{x}^{(n,l)}$ in (4), $m\in\{0,...,|\mathcal{A}|\}$ indexes the set of strong augmentations, and $k\in\{0,1\}$ indexes one of the two weak augmentations applied to the strongly augmented image. Lastly, to further constrain the normal and strongly augmented data representations in (4), our PMSACL loss minimises the distance between samples centred by their representation means computed as:

$$\mathbf{f}^{(n,l)}=\frac{f_{\theta}(\mathbf{x}^{(n,l)})-\mathbf{c}_{n}}{\|f_{\theta}(\mathbf{x}^{(n,l)})-\mathbf{c}_{n}\|_{2}},$$

(6)

where $\mathbf{c}_{n}$ is defined in (2). Also in (4) to map the representations from the same distribution into a denser region of the hyper-sphere [14], we propose a temperature scaling strategy defined as:

$$\kappa(n,m)=\begin{cases}1/(\alpha\tau)&,\text{if }n=m\\ 1/\tau&,\text{otherwise}\end{cases},$$

(7)

where $\alpha$ is a scaling factor that controls the shrinkage level of the temperature $\tau$. As a result, Eq. (7) alters the temperature for the samples that belong to the same strong augmentation distributions (i.e., when $n=m$) to a smaller value $\alpha$, which allows smaller amount of repelling strength compared to samples that belongs to strong augmentation distributions (i.e., $n\neq m$). Putting all together, the loss in (4) clusters the image representations into hyper-spheres and regions within the hyper-spheres, where each hyper-sphere and region represent a different type of augmentation.

Inspired by [53, 60], we further constrain the training in (1) with a self-supervised classification constraint $\ell_{aug}(\cdot)$ that enforces the model to classify the strong augmentation function (Fig. 1):

$$\ell_{aug}(\mathcal{D};\beta)=-\mathbb{E}_{\mathbf{x}\in\mathcal{D},n\in\{0,...,|\mathcal{A}|\},a\sim\mathcal{A}_{n}}\left[\log\mathbf{a}_{n}^{\top}f_{\beta}(f_{\theta}(a(\mathbf{x})))\right],$$

(8)

where $f_{\beta}:\mathcal{Z}\rightarrow[0,1]^{|\mathcal{A}|}$ is a fully-connected (FC) layer, and $\mathbf{a}_{n}\in\{0,1\}^{|\mathcal{A}|}$ is a one-hot vector representing the strong augmentation distribution (i.e., $\mathbf{a}_{n}(j)=1$ for $j=n$, and $\mathbf{a}_{n}(j)=0$ for $j\neq n$).

The final constraint in (1) is based on the relative patch location from the centre of the training image and is adapted for local patches. This constraint is added to learn positional and texture characteristics of the image in a self-supervised manner. Inspired by [19], the positional constraint predicts the relative position of the paired image patches, with its loss defined as

$$\ell_{pos}(\mathcal{D};\gamma)=-\mathbb{E}_{\{\mathbf{x}_{\omega_{1}},\mathbf{x}_{\omega_{2}}\}\sim\mathbf{x}\in\mathcal{D}}\left[\log\mathbf{p}^{\top}f_{\gamma}(f_{\theta}(\mathbf{x}_{\omega_{1}}),f_{\theta}(\mathbf{x}_{\omega_{2}}))\right],$$

(9)

where $\mathbf{x}_{\omega_{1}}$ is a randomly selected fixed-size image patch from $\mathbf{x}$, $\mathbf{x}_{\omega_{2}}$ is another image patch from one of its eight neighbouring patches, $\omega_{1},\omega_{2}\in\Omega$ represents indices to the image lattice, $f_{\gamma}:\mathcal{Z}\times\mathcal{Z}\rightarrow[0,1]^{8}$, and $\mathbf{p}=\{0,1\}^{8}$ is a one-hot encoding of the patch location. The constraints in (8) and (9) are designed to improve training regularisation.

After pre-training $f_{\theta}(\cdot)$ with PMSACL, we fine-tune it with a SOTA UAD, such as IGD [11] or PaDiM [17]. Those methods use the same training set $\mathcal{D}$ as PMSACL, containing only normal images from healthy patients.

IGD [11] combines three loss functions: 1) two reconstruction losses based on local and global multi-scale structural similarity index measure (MS-SSIM) [66] and mean absolute error (MAE) to train the encoder $f_{\theta}:\mathcal{X}\to\mathcal{Z}$ and decoder $g_{\phi}:\mathcal{Z}\to\mathcal{X}$, 2) a regularisation loss to train adversarial interpolations from the encoder [5], and 3) an anomaly classification loss to train $h_{\psi}:\mathcal{Z}\to[0,1]$. The anomaly detection score of image $\mathbf{x}$ is defined by

$$s_{IGD}(\mathbf{x})=\xi\ell_{rec}(\mathbf{x},\tilde{\mathbf{x}})+(1-\xi)(1-h_{\psi}(f_{\theta}(\mathbf{x}))),$$

(10)

where $\tilde{\mathbf{x}}=g_{\phi}(f_{\theta}(\mathbf{x}))$, $h_{\psi}(.)$ returns the likelihood that $\mathbf{x}$ is a normal image, $\xi\in[0,1]$ is a hyper-parameter, and

$\begin{split}\ell_{rec}(\mathbf{x},\tilde{\mathbf{x}})=\rho\|\mathbf{x}-\tilde{\mathbf{x}}\|_{1}+(1-\rho)\left(1-\left(\nu m_{G}(\mathbf{x},\tilde{\mathbf{x}})+(1-\nu)m_{L}(\mathbf{x},\tilde{\mathbf{x}})\right)\right),\end{split}$

(11)

with $\rho,\nu\in[0,1]$, $m_{G}(\cdot)$ and $m_{L}(\cdot)$ denoting the global and local MS-SSIM scores from the global and local models, respectively [11]. In particular, the Gaussian Anomaly classifier $h_{\psi}$ constrains the latent space of the IGD encoder with the following loss:

$$\ell_{h}(\mathbf{x})=h_{\psi}(f_{\theta}(\mathbf{x})),$$

(12)

where the classifier is defined with $h_{\psi}(f_{\theta}(\mathbf{x}))=1-\exp\left(-\frac{\|f_{\theta}(\mathbf{x})-\mu_{\gamma}\|_{2}^{2}}{\sigma_{\gamma}^{2}}\right)$, $f_{\theta}:\mathcal{X}\to\mathcal{Z}$ represents the encoder parameterised by $\theta$, with $\mathcal{Z}\subset\mathbb{R}^{Z}$ denoting the space of latent embeddings of the auto-encoder. The mean and standard deviation values above are estimated with $\mu_{\gamma}=\frac{1}{|\mathcal{D}|}\sum_{i=1}^{|\mathcal{D}|}f_{\theta}(\mathbf{x}_{i})$ and $\sigma^{2}_{\gamma}=\frac{\kappa}{|\mathcal{D}|}\sum_{i=1}^{|\mathcal{D}|}\|f_{\theta}(\mathbf{x}_{i})-\mu_{\gamma}\|_{2}^{2}$, where $\kappa\in(0,1]$ is a constant that regularises the estimation of $\sigma^{2}_{\gamma}$ to prevent numerical instabilities during training. Such an anomaly classifier optimisation has shown to be less sensitive to outliers compared with other anomaly detection approaches. Anomaly segmentation uses (10) to compute $s_{IGD}(\mathbf{x}_{\omega})$, $\forall\omega\in\Omega$ using global and local models, where $\mathbf{x}_{\omega}\in\mathbb{R}^{\hat{H}\times\hat{W}\times C}$ is an image patch. This forms a heatmap, where large values of $s_{IGD}(.)$ denote anomalous regions. The final heatmap is formed by summing up the global and local heatmaps.

PaDiM [17] utilises the multi-layer features from the pre-trained network $f_{\theta}(.)$ to learn a position dependent multi-variate Gaussian distribution of normal image patches. Training uses samples collected from the concatenation of the multi-layer features from each patch position $\omega\in\Omega$ to learn the mean and covariance of the Gaussian model denoted by $\mathcal{N}(\mu_{\omega},\Sigma_{\omega})$ [17]. Anomaly detection is based on the Mahalanobis distance between the concatenated testing patch feature $\mathbf{x}_{\omega}$ and the learned Gaussian distribution $\mathcal{N}(\mu_{\omega},\Sigma_{\omega})$ at that patch position $\omega\in\Omega$ to provide a score of each patch position [17]. In particular, anomaly segmentation is inferred using the following anomaly score map:

$$s_{PaDiM}(\mathbf{x}_{\omega})=\sqrt{(\mathbf{x}_{\omega}-\mu_{\omega})^{\top}\Sigma_{\omega}^{-1}(\mathbf{x}_{\omega}-\mu_{\omega})},$$

(13)

and the final score of the whole image $\mathbf{x}$ is defined as: $s_{PaDiM}(\mathbf{x})=\max_{\omega\in\Omega}s_{PaDiM}(\mathbf{x}_{\omega})$.

We test our self-supervised pre-training PMSACL on four health screening datasets, where we run experiments for both anomaly detection and localisation. The datasets for anomaly detection and localisation are: the colonoscopy images of Hyper-Kvasir dataset [6], and the glaucoma dataset using fundus images [28]. We also run anomaly detection without localisation experiments on the following datasets: Liu et al.’s colonoscopy dataset [36], and Covid-19 chest ray dataset [64] – these two datasets do not have lesion segmentation annotations, so we test anomaly detection only.

Hyper-Kvasir is a large multi-class public gastrointestinal imaging dataset [6]. We use a subset of the normal (i.e., healthy) images from the dataset for training. Specifically, 2,100 images from ‘cecum’, ‘ileum’ and ‘bbps-2-3’ are selected as normal, from which we use 1,600 for training and 500 for testing. We also take 1,000 abnormal images and their segmentation masks of polyps to be used exclusively for testing, where all images have size 300 $\times$ 300 pixels.

LAG is a large scale fundus image dataset for glaucoma diagnosis [28]. For the experiments, we use 2,343 normal (negative glaucoma) images for training, and 800 normal images and 1,711 abnormal images with positive glaucoma with annotated attention maps by ophthalmologists for testing, where images are 500 $\times$ 500 pixels. The annotated attention maps are based on eye tracking, in which the maps are used by the ophthalmologists to explore the region of interest for glaucoma diagnosis [28].

Liu et al.’s colonoscopy dataset is a colonoscopy image dataset with 18 colonoscopy videos from 15 patients [36]. The training set contains 13,250 normal (healthy) images without polyps, and the testing set contains 967 images, with 290 abnormal images with polyps and 677 normal (healthy) images without polyps, where all images have size 64 $\times$ 64 pixels.

Covid-X [64] has a training set with 1,670 Covid-19 positive and 13,794 Covid-19 negative CXR images. The test set contains 400 CXR images, consisting of 200 positive and 200 negative images. We train the methods with the 13,794 Covid-19 negative CXR training images and test on the 400 CXR images, where images are 299 $\times$ 299 pixels.

For the proposed PMSACL pre-training, we use Resnet18 [23] as the backbone architecture for the encoder $f_{\theta}(\mathbf{x})$, and similarly to previous works [7, 52], we add an MLP to this backbone as the projection head for the contrastive learning, which outputs features in $\mathcal{Z}$ of size 128. All images from the Hyper-Kvasir [6], LAG [28] and Covid-X [64] datasets are resized to 256 $\times$ 256 pixels. For the Liu et al.’s colonoscopy dataset [36], we use the original image size of 64 $\times$ 64 pixels. The batch size is set to 32 and learning rate to 0.01 for the self-supervised pre-training on all datasets. The model is trained using stochastic gradient descent (SGD) optimiser with momentum.

We investigate the impact of different strong augmentations in $\mathcal{A}_{n}$, including rotation, permutation, cutout, Gaussian noise and our proposed MedMix. For MedMix patches, we randomly apply colour jittering, Gaussian noise, and non-linear intensity transformations (i.e., fisheye and horizontal wave transformations). In particular, the probability of applying the colour jittering, Gaussian noise, fisheye, and horizontal wave transformations is 25%. For the patch selection and placement, we randomly cropped a patch from a random image within the minibatch and randomly paste the cropped patch into a random location of another random image from the minibatch. For non-linear transformation, three randomly selected control points are used to alter the MedMix cropped patches and we follow the default settings from the opencv library. The weak augmentations in $\mathcal{A}_{0}$ are the same as in SimCLR [7], namely: colour jittering, random greyscale, crop, resize, and Gaussian blur. We use the same weak augmentation hyper-parameters as the original SimCLR paper (i.e., Gaussian kernel parameters, crop resolution, etc).

The model pre-trained with PMSACL is fine-tuned with IGD [11] or PaDiM [17]. For IGD [11], we pre-train the global and local models (see Figure 1), where the patch position prediction loss in Eq. 9 is only fine-tuned for the local model. For PaDiM [17], we pre-train the global model and use it to fine-tune the anomaly detection and segmentation models. For the training of IGD [11] and PaDiM [17], we use the hyper-parameters suggested by the respective papers. In our experiments, the local map for IGD is obtained by considering each 32$\times$32-pixel patch as an instance and apply our proposed self-supervised learning to it. The global map for IGD is computed based on the whole image sized as 256 $\times$ 256 pixels for Hyper-Kvasir, LAG and Covid-X datasets. For Liu et al.’s colonoscopy dataset, we only train the model globally with the image size 64 $\times$ 64. For the auto-encoder in IGD, we use the setup suggested in [11], where the global model is trained with images of size 256$\times$256 pixels or 64 $\times$ 64 for Liu et al.’s colonoscopy dataset, and the local model is trained with image patches of size 32$\times$32. For PaDiM [17], we only use the default setup in their work and compute the segmentation mask based on the images of size 256$\times$256 pixels for Hyper-Kvasir, LAG and Covid-X datasets, and 64 $\times$ 64 for Liu et al.’s colonoscopy dataset.

The anomaly detection performance is quantitatively assessed by the area under the receiver operating characteristic curve (AUROC), specificity, sensitivity and accuracy. AUROC assesses anomaly detection by varying the classification threshold and computing the area under the ROC curve. Sensitivity and specificity reflect the percentage of positives and negatives that are correctly detected. Accuracy shows the overall performance of correctly detected samples for both positive and negative images, where the classification threshold is estimated with a small validation set that contains 50 normal and 50 abnormal images that are randomly sampled from the testing set. Note that the validation set is only used for threshold estimation. For anomaly segmentation, the performance is measured by Intersection over Union (IoU), Dice score, and Pro-score [4]. IoU is computed by dividing the intersection by the union between the predicted segmentation and the ground truth mask. Dice also takes the predicted segmentation and the ground truth mask and divides two times their intersection by their sum. Pro-score weights the ground-truth masks of different sizes equally [4] to verify if both large and small abnormal lesions are accurately segmented.

In this section, we show the anomaly detection results on all datasets.

In this section, we show the anomaly localisation results on Hyper-Kvasir and LAG.

In this section, we show examples of anomaly localisation and detection results, and t-SNE results displaying the distribution of image representations of the normal and pseudo abnormal classes in the feature space.

In this section, we study the roles played by PMSACL components. We start by investigating the loss terms in (1). Then we study the impact of using different types of strong data augmentation to generate the pseudo abnormal images and the number of pseudo abnormal classes in MedMix (i.e., the size $|\mathcal{A}|$ in (1)). We also compare our approach with other recently proposed self-supervised pre-training approaches.