Few-Shot Anomaly Detection for Polyp Frames from Colonoscopy

The data set is obtained from 18 colonoscopy videos from 15 patients. Video frames containing blurred visual information are removed using the variance of Laplacian method [5]. We then sub-sample consecutive frames by taking one of every five frames because the correlation between them makes training ineffective. We also remove frames containing feces and water to reduce the need for a very large normal training set (we plan to handle such distractors in future work). This data set is defined by $\mathcal{D}=\{(\mathbf{x},d,y)_{i}\}_{i=1}^{|\mathcal{D}|}$, where $\mathbf{x}:\Omega\rightarrow\mathbb{R}^{3}$ denotes a colonoscopy frame ($\Omega$ represents the frame lattice), $d\in\mathbb{N}$ represents patient identification²²2Note that the data set has been de-identified, so $d$ is useful only for splitting $\mathcal{D}$ into training, testing and validation sets in a patient-wise manner., $y\in\mathcal{Y}=\{Normal,Abnormal\}$ denotes the normal (without polyp) and abnormal (with polyp) classes. The distribution of this data set is as follows: 1) Training set: a set of 13250 normal images (without polyps), denoted by $\mathcal{D}_{N}\subset\mathcal{D}$, and a set containing between 10 and 80 abnormal images, denoted by $\mathcal{D}_{A}\subset\mathcal{D}$; 2) Validation set: 100 normal images and 100 abnormal images for model selection; and 3) Testing set: 967 images, with 217 (25% of the set) abnormal images and 700 (75% of the set) normal images. The patients in the testing set do not appear in the training/validation sets and vice versa. This abnormality proportion (on the testing set) is commonly defined in other anomaly detection literature [21, 26]. These frames were obtained with the Olympus ®190 dual focus endoscope.

The training process of our proposed FSAD-NET (Fig. 3) is divided into two stages: 1) pre-training of a feature encoder $\mathbf{z}=f_{E}(\mathbf{x};\theta_{E})$ ($\theta_{E}$ is the encoder parameter and $\mathbf{z}\in\mathbb{R}^{Z}$) to learn an image embedding that maximises the mutual information (MI) between normal images $\mathbf{x}\in\mathcal{D}_{N}$ and their embeddings $\mathbf{z}$ [6]; and 2) training of the SIN $f_{S}(f_{E}(\mathbf{x};\theta_{E});\theta_{S})$ [19], parameterised by $\theta_{S}$, with a contrastive-like loss that uses $\mathcal{D}_{N}$ and $\mathcal{D}_{A}$ to achieve the goal $f_{S}(f_{E}(\mathbf{x}\in\mathcal{D}_{A};\theta_{E});\theta_{S})>f_{S}(f_{E}(\mathbf{x}\in\mathcal{D}_{N};\theta_{E});\theta_{S})$.

More specifically, the training of the encoder to maximise the MI between the normal samples $\mathbf{x}\in\mathcal{D}_{N}$ and their feature embeddings $\mathbf{z}=f_{E}(\mathbf{x}\in\mathcal{D}_{N};\theta_{E})$ [6] is achieved with

where $\alpha,\beta,\gamma$ are the model hyperparameters, the functions $\hat{I}_{G}(.)$ and $\hat{I}_{L}(.)$ denote an MI lower bound based on the Donsker-Varadhan representation of the Kullback-Leibler (KL)-divergence [6], defined by

with $\mathbb{J}$ denoting the joint distribution between images $\mathbf{x}$ and their respective embeddings $\mathbf{z}=f_{E}(\mathbf{x};\theta_{E})$, $\mathbb{M}$ representing the product of the marginals of the images and embeddings, and $f_{G}(\mathbf{x},f_{E}(\mathbf{x};\theta_{E});\theta_{G})$ being a discriminator parameterised by $\theta_{G}$. Also in (1), the function $\hat{I}_{\theta_{L}}(\mathbf{x}(i);f_{E}(\mathbf{x}(i);\theta_{E}))$, defined similarly as (2) for the discriminator $f_{L}(\mathbf{x}(\omega),f_{E}(\mathbf{x}(\omega);\theta_{E});\theta_{L})$, is the local MI between image regions $\mathbf{x}(\omega)$ ($\omega\in\mathcal{M}\subset\Omega$) and respective local embeddings $f_{E}(\mathbf{x}(\omega),\theta_{E})$. Moreover in (1),

with $\mathbb{V}$ denoting a prior distribution for the embeddings $\mathbf{z}$ ($\mathbb{V}$ is assumed to be a normal distribution $\mathcal{N}(.;\mu_{\mathbb{V}},\Sigma_{\mathbb{V}})$, with mean $\mu_{\mathbb{V}}$ and covariance $\Sigma_{\mathbb{V}}$), $\mathbb{P}$ the distribution of the embeddings $\mathbf{z}=f_{E}(\mathbf{x}\in\mathcal{N}_{N};\theta_{E})$, and $d(.;\phi)$ is a discriminator modelled with adversarial training to estimate the likelihood that the input is sampled from $\mathbb{V}$ or $\mathbb{P}$. This objective function pulls the feature embeddings of the normal images toward $\mathcal{N}(.;\mu_{\mathbb{V}},\Sigma_{\mathbb{V}})$.

The next step of the learning process consists of computing the embeddings of normal and abnormal images with $\mathbf{z}=f_{E}(\mathbf{x}\in\mathcal{D}_{A}\bigcup\mathcal{D}_{N};\theta_{E}^{*})$ to train $f_{S}(\mathbf{z};\theta_{S})$ using a contrastive-like loss to directly optimise the anomaly score [19]. More specifically, the constrastive loss for each training sample is defined as:

where $\mathbb{I}(.)$ is an indicator function that is equal to one when the condition in the parameter is true, and zero otherwise, $s(x)=\frac{x-\mu_{S}}{\sigma_{S}}$ with $\mu_{S}=0$ and $\sigma_{S}=1$ representing the mean and standard deviation of the prior distribution for the anomaly scores for normal images, and $a$ is the minimum margin between $\mu_{S}$ and the anomaly scores of abnormal images [19]. The loss in (4) pulls the scores from normal images to $\mu_{S}$ and pushes the scores of abnormal images away from $\mu_{S}$ with a margin of at least $a$.

During inference, we take a test image $\mathbf{x}$, compute the feature embedding with $f_{E}(\mathbf{x};\theta_{E})$ and then compute the score with $s=f_{S}(\mathbf{z};\theta_{S})$ – the score result $s$ is then compared to a threshold $\tau$ to determine if the test image is normal or abnormal. We considered the score $s$ as the estimation of the notion of closeness which is related to the likelihood that the embedding of a colonoscopy image is classified as belonging to the set of normal images.

The original colonoscopy images are resized from initial resolution $1072\times 1072\times 3$ to $64\times 64\times 3$ to reduce the training and inference computational costs. We found that $64\times 64\times 3$ is the minimum size that we can use without a negative impact on AUC. We note the polyps are still visible at such resolution, as shown in Fig S4. The model selection (to select optimiser, learning rate and model structure) is done using the validation set mentioned in Sec. 3.1. We use Adam [8] optimiser during training with a learning rate of 0.0001 for the encoder and SIN learning. We adopt batch normalisation for both stages. We make sure our method uses a similar backbone architecture as other competing approaches in Tab. 1. In particular, the encoder $f_{E}(.;\theta_{E})$ uses four convolution layers (with 64, 128, 256, 512 filters of size 4 $\times$ 4). The global discriminator $f_{G}(.;\theta_{G})$ has three convolutional layers (with 128, 64, 32 filters of size 3 $\times$ 3). The local discriminator $f_{G}(.;\theta_{G})$ has three convolutional layers (with 192, 512, 512 filters of size 1 $\times$ 1). The prior discriminator $d(.;\phi)$ has three linear layers with 1000, 200, 1 nodes per layer). We also use the validation set to estimate $a=6$ in (4). In (1), we follow the DIM paper for setting the hyper-parameters as follows [6]: $\alpha=0.5$, $\gamma=1$, $\beta=0.1$. For the prior distribution for the embeddings in (3), we set $\mu_{\mathbb{V}}=\mathbf{0}$ (i.e., a $Z$-dimensional vector of zeros), and $\Sigma_{\mathbb{V}}$ is a $Z\times Z$ identity matrix. To train the model, we first train the encoder, local, global and prior discriminator (representation learning stage) for 6000 epochs with a mini-batch of 64 samples. We then train SIN for 1000 epochs, with a batch size of 64, while fixing the parameters of encoder, local, global and prior discriminator. We implement our method using Pytorch [20].

The detection results are measured with the area under the receiver operating characteristic curve (AUC) on the test set [26, 21], computed by varying the inference threshold $\tau$ for the score result $s$.

The test set AUC results shown in Table 1 are divided into zero-shot and few-shot. The zero-shot rows show results obtained from the following zero-shot anomaly detection methods³³3Codes were downloaded from the authors’ Github pages and tuned for our problem.: ADGAN [14], OCGAN [21], f-anogan and its variants [26] that involve image-to-image mean square error (MSE) loss (izi), Z-to-Z MSE loss (ziz) and its hybrid version (izif). Our FSAD-NET model outperforms all zero-shot learning methods by a large margin, showing the importance of using a few abnormal samples for training. For the few-shot results, we consider the cases where we have 30 and 40 abnormal training images, and we test several variants of the FSAD-NET. We use between 30 and 40 abnormal training images because that is the range, where we observe that our FSAD-NET produces stable AUC results. As a baseline approach, we train Densenet121 [7] using high levels of data augmentation to deal with the training imbalance issue. However, our FSAD-Net outperforms Densenet121 by a large margin. The variants of FSAD-NET are designed to test the importance of each stage of our method. The methods labelled as ’Cross entropy’ and ’Focal loss’ replace the contrastive loss in (4) by the cross entropy loss (commonly used in classification problems) [4] and the focal loss (robust to imbalanced learning problems) [11], respectively. FSAD-NET shows substantially better results, indicating the importance of using a more appropriate loss function for few-shot anomaly detection. To show the importance of representation learning (RL) in FSAD-Net, we tested FSAD-Net without it, which shows much lower AUC results than competing approaches. Also, we compared our method with a few-shot learning baseline [23], which proposes a learning algorithm for highly imbalanced learning problems. When used to train FSAD-Net, it achieved 78.62% of mean AUC when training with 40 abnormal training samples. Hence our model shows more accurate results than that approach. Furthermore, we test the importance of DIM to train the encoder in (1) by replacing it by the deep auto-encoder [16] (labelled as AE network) – results show that FSAD-NET is more accurate, indicating the effectiveness of using MI and prior distribution for learning the feature embeddings in (1).

We further investigate the performance of our proposed FSAD-NET as a function of the number of abnormal training images that can vary from $10$ to $80$. For each number of abnormal training images, we train our model three times, using different training sets each time, and we compute the mean and standard deviation of the AUC results. The result of this experiment in Fig. 3 shows that: 1) the performance stabilises between 85%-90% when feeding the model 30 or more abnormal training images; and 2) our method is robust to extremely small training sets of abnormal images. We show a few true positive, true negative, false positive and false negative results produce by FSAD-NET in Fig. S4.