Proxy-bridged Image Reconstruction Network for Anomaly Detection in Medical Images

In the reconstruction based anomaly detection framework, it is expected that the network can well reconstruct the normal input image, and the reconstructed image of the abnormal ones would be with a large reconstruction error. In this way, the abnormal images can be distinguished from the normal images. Towards this end, in this paper, we introduce a Proxy-bridged Image Reconstruction Network (ProxyAno) for anomaly detection in medical images. The proposed ProxyAno consists of two modules, a Proxy Extraction Module ${\mathcal{F}}_{p}(\cdot)$ and an Image Reconstruction Module ${\mathcal{F}}_{g}(\cdot)$. Specifically, the proposed ProxyAno can be formulated as:

where $\hat{{\mathbf{P}}}$ denotes the predicted SI, ${\mathbf{I}}$ and ${\hat{\mathbf{I}}}$ denote original image and reconstructed image, respectively. We use ${\textbf{{Enc}}}_{p}$/${\textbf{{Dec}}}_{p}$ to denote the encoder/decoder in the Proxy Extraction Module, and use ${\textbf{{Enc}}}_{g}$/${\textbf{{Dec}}}_{g}$ to denote the encoder/decoder in the Image Reconstruction Module, respectively.

Fig. 2 illustrates the overview of our ProxyAno. To enlarge the reconstruction error for the abnormal image while maintaining a small reconstruction error for the normal images, we propose to: 1) memorize the mapping pattern between the normal input image and its corresponding SI with a memory in the Proxy Extraction Module. Once this module is well trained, the weights of this module are fixed when training the Image Reconstruction Module; 2) reconstruct the image from its SI, meanwhile repairing the abnormal SI in the Image Reconstruction Module. To achieve this, we first create a pseudo abnormal proxy ${\mathbf{P}^{{}^{\prime}}}$ and feed it into the Image Reconstruction Module with ${\hat{\mathbf{I}^{{}^{\prime}}}}={\mathcal{F}}_{g}({\mathbf{P}^{{}^{\prime}}})$. Then we propose a repairing loss to enforce the similarity between the reconstructed image ${\hat{\mathbf{I}^{{}^{\prime}}}}$ of the pseudo abnormal proxy and the normal image ${\mathbf{I}}$. Benefiting from these two improvements, in the test phase, both $\hat{{\mathbf{P}}}$ and ${\hat{\mathbf{I}}}$ tend to be close to the normal ones, leading to a small reconstruction error for the normal samples and a large reconstruction error on the abnormal ones. Then the reconstruction error is used as a measurement to detect the anomalies.

In our work, we use the SLIC [48] algorithm to obtain the superpixels. The SLIC algorithm groups meaningful pixels into a superpixel by composing spatially adjacent pixels. Based on the superpixels, we can get an SI that consists of the superpixel colored with the average intensity for pixels within each superpixel. Then we use the SI extracted with the SLIC algorithm as supervision to train a Proxy Extraction Module. It is worth noting that directly applying the SLIC algorithm on the abnormal image, the superpixel structure of lesion can also be extracted. Thus, the abnormal SI contains some lesion information; consequently, the lesion in the abnormal images may also be reconstructed. To increase the fidelity of the reconstruction of the normal images and increase the reconstruction error for the abnormal ones, on the one hand, we introduce a memory to map the abnormal image to a normal SI, on the other hand, we also propose to let the algorithm automatically repairing an SI of an abnormal image, and map it to a normal input image.

In the Proxy Extraction Module, we propose to augment the network with a memory to explicitly memorize the correspondence between the normal input and its SI. The motivations behind using a memory to augment the network are three aspects: 1) to recognize disease in medical images, the clinicians need to memorize the characteristics of normal samples, and the superpixel structure is an important characteristic in the image; 2) the memory is a type of location-independent long-term attention, which augments the ability of convolutional structure; 3) the correspondence between the normal input and its SI is very regular. Using the correspondence patterns is probably enough to generate all normal SI’s from its input. Therefore, we introduce a memory that can memorize the normal patterns. The feature is first extracted with a encoder. Then, instead directly feed the feature into the decoder, we use the feature as a query to retrieve the most relevant item in the memory. The features feed into the decoder are obtained from a selected memory item of the normal data.

Since this memory is learnt based on the correspondence of normal images, it can be expected that this would cause information loss for the abnormal ones when generating its SI. Consequently, the image reconstruction form this SI would be with a large reconstruction error, which is a desirable property for anomaly detection. To achieve this goal, we introduce a memory $m\in{\mathbb{R}}^{k\times d}$ to memorize the latent features in the Proxy Extraction Module, where $k$ is the size (the number of items) of the memory , and $d$ is the dimensionality of the latent vector $m_{j}\in{\mathbb{R}}^{d},j\in\{1,2,\cdots,k\}$. As shown in Fig. 2, the framework takes an image ${\mathbf{I}}$ as an input, and then ${\mathbf{I}}$ is passed through the encoder ${\textbf{{Enc}}}_{p}$ to extract a latent feature $z\in{\mathbb{R}}^{h\times w\times d}$, where $h$ and $w$ are the height and the width of the feature map. In this module, $z$ is used to search for the nearest neighbor item in the memory $m$, and we denote the nearest item of $z$ as $\tilde{z}\in{\mathbb{R}}^{s\times d}$, here $\tilde{z}$ indicates the nearest item in the memory of a feature.

Specifically, as illustrated in Fig. 3, we first flatten the latent feature as $z\in{\mathbb{R}}^{s\times d}$, where $s=h\times w$. Then $\forall i\in\{1,\cdots,s\}$, $z_{i}\in{\mathbb{R}}^{d}$ is replaced by its nearest neighbor item $m_{\text{J}}\in{\mathbb{R}}^{d}$ in the memory as follows:

Then the input to the decoder ${\textbf{{Dec}}}_{p}$ is the substituted feature $\tilde{z}$. To simplify the description, we denote ${\mathcal{G}}$ as the function of retrieving the corresponding item in the memory and using the retrieved item to replace the input. Then we arrive at the following mapping function:

In the training phase, we update the memory with exponential moving averages [49]. We denote ${z_{i,1},z_{i,1},\dots,z_{i,n_{i}}}$ as $n_{i}$ latent features that are closest to the memory item $m_{i}$. In order to make the memory item close to the set of latent features, we have the loss as:

The optimal $m_{i}$ has a closed form solution, which is the average of the latent features in the set:

In order to handle online mini-batch training, the exponential moving average [49] is used:

where $\gamma$ is a constant between 0 and 1, meaning how much history data to be kept. $n_{i}^{(t)}$ and $z_{i,j}^{(t)}$ denote $n_{i}$ and $z_{i,j}$ in the $t-th$ mini-batch, respectively.

During the training of Proxy Extraction Module, we update the encoder and decoder simultaneously with the update of the memory. The encoder and decoder are optimized as:

where ${\mathbf{P}}$ is the SI extracted from the original image ${\mathbf{I}}$ with the SLIC algorithm [48]. It is worth noting that there is no gradient defined for the function ${\mathcal{G}}$. However, it is possible to approximate the gradient with the straight-through estimator [49]. We copy the gradient of $\tilde{z}$ as the gradient of $z$. Once the Proxy Extraction Module is trained, the weights of ${\textbf{{Dec}}}_{p}$, ${\textbf{{Enc}}}_{p}$ and the items in memory are fixed when optimizing the Image Reconstruction Module.

Our method vs. MemAE [50]. Recently, Gong et al. [50] proposed to augment an Auto-Encoder with the memory module, termed MemAE. Given an input, the encoder in MemAE extracts an encoded representation, which is used as a query to retrieve the most relevant items in the memory. The multiple items are then averaged with an attention weight to get the substituted feature $\tilde{z}$ for the reconstruction. Both the MemAE and our method can take advantage of memory-augmented networks for anomaly detection. However, the specific task in this work is relatively simpler than [50]. Thus, our method only retrieve the nearest item as the substituted feature rather that retrieve multiple items. The reasons why the task in this paper is relatively simpler are in two aspects: first, the patterns in the medical images are simpler than the patterns in the natural images [51, 52, 53]; second, MemAE [50] reconstructs original image while our method predicts the SI, which contains the simple pattern than that in original image. As a result, if a combination of multiple memory items is used, some anomalies may still have the chance to be well reconstructed. Therefore, we only retrieve a single item rather multiple items for SI prediction.

In anomaly detection, the abnormal images are usually not available in the training phase. However, compared with the normal SI, the input image could be regarded as a pseudo abnormal SI in appearance. This inspires us to synthesize a pseudo abnormal proxy by cutting a patch from the normal image, and paste it into the normal proxy. Then we obtain the pseudo abnormal proxy ${\mathbf{P}^{{}^{\prime}}}$, and ${\mathbf{M}}\in\{0,1\}$ denotes a binary mask indicating the location of pasting the patch into ${\mathbf{P}}^{\prime}$. The pseudo abnormal proxy is inputted into the Image Reconstruction Module and we propose a repairing loss to enforce the output of abnormal proxy as a normal image. Thus, in the test phase, the reconstruction error corresponding to the abnormal images would be boosted, and consequently, the discrepancy between the normal images and the abnormal images is enlarged, which is a desirable feature for anomaly detection.

To obtain diverse pseudo abnormal proxy ${\mathbf{P}^{{}^{\prime}}}$, we propose a Pseudo Abnormal Proxy Constructor (PAPC) that crops the patch from the image with a random size and pasts it into the proxy at a random position. As shown in Fig. 4, the pseudo abnormal proxy can mimic the abnormal images to some extent. In the training phase, we enforce the network to reconstruct the normal image even with the pseudo abnormal SI, which is like the ‘anomalies repairing’. In this way, the reconstruction error corresponding to the abnormal images would be boosted, and consequently, the discrepancy between the normal images and the abnormal images is enlarged, which is a desirable feature for anomaly detection.

Specifically, in the training of the Image Reconstruction Module, the forward pass of the normal proxy and the pseudo abnormal proxy is denoted as ${\hat{\mathbf{I}}}={\mathcal{F}}_{g}(\hat{{\mathbf{P}}})$ and ${\hat{\mathbf{I}^{{}^{\prime}}}}={\mathcal{F}}_{g}({\mathbf{P}^{{}^{\prime}}})$, respectively. The function ${\mathcal{F}}_{g}(\cdot)$ is shared.

Normal Image Reconstruction Loss. For the normal proxy, the image reconstruction loss is shown as follow:

where ${\mathbf{D}}$ denotes the discriminator [16], and the $\lambda_{\text{g}}=0.01$ is a hyper-parameter.

Abnormal Image Repairing Loss. For the pseudo abnormal proxy, we propose a repairing loss to enforce the output of abnormal proxy as a normal image. The repairing loss consists of a global regularization item and a local regularization item, which are defined as follows:

where $\odot$ is element-wise multiplication.

Total Objective Function. We arrive at the objective function for the Image Reconstruction Module:

where $\lambda_{\text{gobal}}$ and $\lambda_{\text{local}}$ are the hyper-parameters. Empirically, we set $\lambda_{\text{gobal}}=0.25$, $\lambda_{\text{local}}=0.5$ on all datasets in our experiments.

Taking an original image ${\mathbf{I}}\in{\mathbb{R}}^{w\times h}$ as an input, the proposed method can obtain an extracted proxy $\hat{{\mathbf{P}}}\in{\mathbb{R}}^{w\times h}$ and a reconstructed image ${\hat{\mathbf{I}}}\in{\mathbb{R}}^{w\times h}$. The $w$ and $h$ denote the width and the height of image, respectively. We compute the anomaly score for the image-level anomaly detection in the latent feature space, and compute the anomaly score map for the pixel-level anomaly detection in the image space.

Anomaly Score for Image-level Anomaly Detection. It has been proven that computing the anomaly score by mapping the image space to the latent space is effective [18, 36]. However, training an additional encoder is inefficient and redundant. We apply the existing encoder ${\textbf{{Enc}}}_{p}$ in Proxy Extraction Module to map the image to a latent space, and we get the latent feature as:

where $\hat{z}$ is the latent feature of the reconstructed image ${\hat{\mathbf{I}}}$.

We compute the anomaly score as:

Anomaly Score Map for Pixel-level Anomaly Detection. To get the pixel-level anomaly map ${\mathbf{A}}_{\text{pix}}\in{\mathbb{R}}^{w\times h}$ for lesion segmentation, we compute the anomaly score map in the image space as:

We use the same encoder and decoder in both the Proxy Extraction Module and the Image Reconstruction Module. The detailed architectures of the encoder and the decoder are shown in Fig. 5.

In this part, to verify the effectiveness of proxy approach, we study different proxy types and compare them with baseline that without proxy (i.e., Auto-Encoder). Specifically, besides the SI, we also consider the following five types of proxy: 1) the edge extracted by a Canny edge detector; 2) smooth image, which is a short for the image smoothed with Gaussian blur; 3) image with smooth patches, which denotes the image consists of patches colored with the average intensity in each patches; 4) edge $\oplus$ smooth image, here $\oplus$ denotes the concatenation; 5) edge $\oplus$ image with smooth patches. The experiments are conducted on the retinal OCT [54] dataset. The illustration of smooth image and image with smooth patches can be found in Fig. 8.

The quantitative results are reported in Table III, from which we can see that all proxy methods outperform the baseline without proxy, validating proxy approach is effective. We can also see that the SI-based proxy is superior to all other types of proxy. The reason of SI-based proxy outperforms the edge-based proxy is possibly because that the edge contains less texture information, and makes the image reconstruction difficult. Compared with the edges, the proposed SI contains more texture information (i.e., the average intensity within each superpixel), which makes the image reconstruction easier, and facilitates the anomaly detection. The reason of SI-based proxy outperforms the proxy based on smooth image and image with smooth patches is possibly because that the SI contains more structure information. Compare with the smooth image, in the SI the edge is relatively enhanced. Compare with the image with smooth patches, the SI contains more semantic information and preserves local structure. As a summary, compare with other types of proxy, the proposed SI contains more texture and structure information, validating that both texture and structure information are important for anomaly detection. This similar argument can also be found in [47]. Moreover, with concatenating edge, the performance of both smooth image and image with smooth patches improved. This also validate the argument that both texture and structure information are important for anomaly detection.

The qualitative results are shown in Fig. 9. As shown in the images in the $1^{st}$ row of Fig. 9, compared with the edge-based reconstruction, our SI-based solution reconstruct the normal image more accurately. Furthermore, the images in the $2^{nd}$ column of Fig. 9 show that although the self-reconstruction based method (i.e., ‘image $\rightarrow$ image’ using the AE) is trained well on the normal OCT images, when we feed the abnormal images (i.e., DME image and drusen image) to the trained AE, it can also well reconstruct the lesions in the abnormal images. On the contrast, if we use an SI as the proxy, the issue is eased. Specifically, as shown by the images in the $5^{th}$ column, since the SI is obtained by a neural network that trained on the normal samples, the network cannot well extract the SI on the DEM images and the drusen images, especial the lesion regions in the abnormal images. Based on the poorly extracted SI, the abnormal images cannot be well reconstructed. The images in the $6^{th}$ column show that the lesions in the DEM image and the drusen image tend to be repaired, thus the reconstruction error between input image and reconstructed image is larger than self-reconstruction based method.

In this part, we conduct several ablation studies on retinal OCT dataset [54] to evaluate each component of our method. We report the results of the difference components in Table IV. Let EncDec denotes the encoder-decoder architecture. The relationship between EncDen and Auto-Encoder is: if the output of EncDec is the same as the input, the EncDec is actually the Auto-Encoder; otherwise, the EncDec is not the Auto-Encoder. The mem^∗ denotes the memory implemented in [50], while the mem denotes the memory implemented in this work. The specific difference between mem^∗ and mem is described in the Section III-C, Proxy Extraction Module. “EncDec $+$ mem” (or mem^∗) denotes that there is a memory between the encoder and decoder in the Auto-Encoder. Both Proxy Extraction Module and Image Reconstruction Module contain the EncDec, thus we use “$2\times$EncDec” to simply denote Proxy Extraction Module and Image Reconstruction Module. Between these two modules, we use SI as the intermediate proxy, which is denoted as “$2\times$EncDec $+$ SI”. Based on “$2\times$EncDec $+$ SI”, “$+$ mem” denotes the Proxy Extraction Module is equipped with the memory, while “$+$ rep” denotes the Image Reconstruction Module is equipped with the proposed repairing loss. “$+$ lat” denotes that we compute the anomaly score in the latent space rather than in the image space.