Visual Anomaly Detection Via Partition Memory Bank Module and Error Estimation

For unsupervised anomaly detection, the detection model is trained by only normal samples, which is used to detect abnormal samples and localize abnormal regions [13, 2]. Traditional classifiers are introduced for anomaly detection, such as SVM [26] and OC-SVM [27]. With the popularity of deep learning, researchers proposed deep one-class methods such as DSVDD [28] and OCNN [29]. The common idea of these methods is to learn the decision boundary of normal samples. For example, DSVDD learns a spherical discriminant plane to improve discrimination efficiency. There are also other methods that model the distribution features of normal samples, such as MRF [30], MDT [31] and GMM [17, 2]. However, these methods are less effective to deal with high dimensional data.

Because of the outstanding performance of transfer learning in other computer vision tasks, the pretrained network is introduced to the field of anomaly detection [32, 6, 33]. In [6], since the student network only learns latent representations of normal samples, the degree of sample abnormality is measured by the difference between latent feature representations from multiple student and teacher networks. However, they rely on additional training datasets and huge resource consumption. Some researches have proposed the employment of GAN[34] to solve the anomaly detection problem. The generator in GAN is used to learn normal sample features and discriminator to identify anomalies [35, 36, 37]. However, GAN may be unstable and tends to derive unsatisfactory results in the actual situation.

Recently, reconstruction-based methods have become popular in anomaly detection, which expects to poorly reconstruct abnormal samples and enlarges the gap of the reconstruction error of abnormal samples and normal samples [38, 12, 13, 39, 40, 41, 14, 24]. AE [15] is used to reconstruct samples. For example, AE-SSIM [14] attempts to reconstruct normal samples directly using AE, while ARFAD [24] learns the reconstruction features of normal samples through self-supervised tasks such as rotation. However, these methods successfully reconstruct not only normal samples but also abnormal samples, which leads to limited performance.

Different from the above methods, this work proposes a new PMB module to reconstruct successfully normal samples and poorly abnormal samples.

The memory network was initially adapted to the field of natural language processing [42, 43, 20, 21, 44]. Recently, Some existing works apply memory networks to anomaly detection [19, 25]. Gong et al. [10] proposes MemAE, which uses a memory module to reduce the reconstruction ability of AE for anomalous samples. However, due to its simple memory module and the shortcomings of the combination method, it is difficult to be effective for tiny anomaly sample detection and localization tasks. Park et al. [19] proposes a combination network of AE and memory module for video anomaly detection. It constructs an update method for the memory items of the memory module and concatenates the outputs of AE and memory module as the input of the decoder. However, it is not suitable for the localization challenges of complex scenes because it focuses more on high-level features. Second, the output of AE as the input of the decoder results in the limitation of features stored in the memory module. Wang et al. [25] introduces an additional evaluation network on the MemAE framework as a discriminator to detect whether a sample is anomalous or not. However, it is difficult to achieve anomalous region localization with discriminators.

The most related study to the proposed method is DAAD. [11]. It uses the method in Figure 1 (2) to generate memory query of size $a\times b\times c$ ( $a\leq h$, $b\leq w$ ) and uses a discriminator to identify whether the samples are anomalous or not. DAAD uses larger query (generally $8\times 8\times c$) to make anomalies and normals that do not share the same block pattern. However, it destroys the integrity of the features learned by the convolutional network, resulting in the memory module learning reorganized logical information. Therefore, abnormal features can also be logically fitted to anomalies such that reconstruction errors are small. The proposed PMB module employs novel query of size $h\times w\times 1$ (Figure 1(3)) which enables the memory module to store normal feature learned from the convolutional network. Therefore, all queries can address the PMB module to get normal features. Unlike DAAD, it is not degrade the reconstruction capability of AE, but make the memory module unable to address abnormal features. In addition, DAAD introduces discriminator with adversarial loss to solve the problem of incorrect discrimination caused by cumulative errors, which results in additional resource consumption and the failure of anomaly localization. In this paper, we not only ensure that the normal features are well reconstructed, but also the anomaly localization results can be derived from the error map without additional resources.

This work proposes to solve the unsupervised anomaly detection problem based on the reconstruction scheme. The overview of the proposed framework is shown in Figure 3. The proposed framework contains three basic structures: 1) an AE, 2) a PMB module, and 3) a Histogram Error Estimation module.

Given the input sample $X$, the encoder first extracts the multi-scale latent representations $F=\{F_{1},F_{2},\ldots,F_{L}\}$. To make full utilization of the representation information provided by the memory modules and allow the memory module to focus on low-level features, we utilize the output of the PMBs at layer $1$ to layer $p$ of the encoder and the output of the subsequent $p+1$ to $L$ layers with the skip connection. Therefore, $p$ PMBs are designed to learn and store feature maps of different scales independently. The reads of the proposed PMB module are represented by $X_{M}=\{X_{M}^{1},X_{M}^{2},\ldots,X_{M}^{p}\}$. Then, the channel concatenation between $X_{M}$ and output of previous decoder $X_{D}$ as the input of the next decoder. The reconstructed image $X_{R}$ is obtained after the decoder. The original difference image $A$ is obtained by using Eq. 1,

Besides, the Histogram Error Estimation module is developed to eliminate reconstruction errors in normal regions to obtain a more accurate corrected error map $A^{\prime}$ for anomaly localization and detection. The anomaly scores and anomaly localization maps of the samples are obtained based on $A^{\prime}$.

As is shown in Figure 4, The PMB module is proposed to learn and store the potential feature representations of normal samples. It achieves the challenge of anomalous sample reconstruction error becoming large by exploring the proposed joint partition mechanism and a new query generation method.

To generate the query with critical feature information, a novel query generation method is developed, as shown in Figure 1 (3). It generates queries directly from each channel of the feature map. Thus, each query contains semantic information for a single channel, allowing the PMB module to store only normal semantic information. In addition, to enhance the learning ability and expressiveness of the memory module and to make the reconstruction error of normal regions small, a partition mechanism is proposed. It introduces multiple local units to store the normal features of different regions separately. Each memory unit is capable of storing detailed feature information for individual regions. Multiple memory units make the expressiveness of the memory module enhanced.

Specifically, for each feature map $F_{l}$ with the size of $h\times w\times c$ from $F$, it is partitioned by using the proposed query generation approach shown in Figure 4. Along the ‘$h$’ dimension of the feature map, $F_{l}$ is partitioned into $k$ regions with the size of $\frac{h}{k}\times w\times c$ , each of which is stored and queried by a separate memory unit. The feature map of individual channel in each partition is flattened to one $\frac{h}{k}\times w$ dimensional vector. Thus, we can get query sets $Q=\{\boldsymbol{q_{t}}|t\in[1,c]\}$ in each partition. Each PMB contains $k$ memory units $\boldsymbol{M}=\{\boldsymbol{M_{1}},\boldsymbol{M_{2}},\ldots,\boldsymbol{M_{k}}\}$, each memory unit contains $m$ memory items, and the size of the memory items is the same as the query $q_{t}$. Finally, as shown in Fig. 4 (4), suppose $\boldsymbol{q_{t}}$ is a query for $i$-th partition of $F_{l}$, $\boldsymbol{q_{t}}$ and the corresponding memory items in $\boldsymbol{M_{i}}$ are first normalized to improve the accuracy of the attention weight. The normalization operations are as follows:

where $norm(\cdot)$ represents the $l_{2}$-norm. $\boldsymbol{\widehat{q_{t}}}$ represents the normalized the query $\boldsymbol{q_{t}}$. $\boldsymbol{\widehat{M_{i}}}$ represents the memory items after normalization. Then the similarity between $\boldsymbol{\widehat{q_{t}}}$ and the memory item $\boldsymbol{\widehat{M_{i}}}$ is used to calculate the attention weight $\boldsymbol{\widehat{\omega_{t}}}$ as follows:

Here $<\cdot,\cdot>$ denotes the cosine similarity. To avoid too small weight of the query, this paper proposes to filter the weights by introducing a threshold, which can alleviate the reconstruction with most of memory items [10]. Consequently, when abnormal features are used as queries, the reads are difficult to fit and the abnormal regions are then difficult to be reconstructed.

where ${{\delta_{m}}}$ represents the threshold. By using the memory module, we can obtain a new feature vector $\boldsymbol{q_{t}^{\prime}}$ for each $\boldsymbol{q_{t}}$ by using the output of the memory unit $\boldsymbol{M_{i}}$.

Then a new feature map $F^{\prime}$ can be obtained by re-arranging all $\boldsymbol{q_{t}^{\prime}}$ and fed into the decoder to generate the reconstructed image $X_{R}$. It is worth noting that the stability of the reconstructed network can be improved by using $\boldsymbol{q_{t}^{\prime}}$ since the value of $\boldsymbol{q_{t}^{\prime}}$ does not change much.

For each image, the original error map $A$ is obtained by using Eq. 1, in which the value of each pixel represents the reconstruction error. In the proposed approach, the reconstruction error comes from the proposed PMB module and AE. The PMB module learns and stores the normal features, which enables to guarantee that the reconstruction errors of normal regions are obviously smaller than ones of abnormal regions. However, due to the presence of the memory module, AE is difficult to achieve zero-error reconstruction in the normal region. These accumulative reconstruction errors can result in false detection from the original error map. Therefore, we need to further mitigate the impact of accumulative reconstruction errors on detection and localization performance.

In this paper, we explore a simple, yet effective Histogram Error Estimation module to estimate the error $\mu$. As shown in Figure 3, the histogram of error map is first constructed. Due to the small reconstruction error of normal region, the pixels in normal region always lie on the left side of histogram. Therefore, we can quickly find normal pixels in the histogram. Then, a percentage $\delta_{H}$ of small reconstruction errors are chosen to estimate $\mu$ as shown in the dashed box of the histogram in Figure 3. The average value of errors of these selected pixels is utilized to estimate $\mu$. Then, the corrected difference image $A^{\prime}$ is obtained as follows:

The anomaly score and localization map of the samples are derived from $A^{\prime}$. $A^{\prime}$ is adopted as the localization map, while the anomaly score is calculated by the sum of the error values of the corrected difference image $A^{\prime}$.

To train the proposed model, the quality of the reconstructed image and the intuitive feeling of the human visual system are jointly explored. The reconstruction loss $\mathcal{L}_{MSE}$ and structural similarity index measure (SSIM) [55] loss ${\mathcal{L}}_{{\rm{SSIM}}}$ are introducesd , which are defined as:

Here $h$ and $w$ represent the height and width of the image. ${\rm{SSIM}}_{s}^{(i,j)}(X,{X_{\rm{R}}})$ represents structural similarity index between the image $X$ and image $X_{R}$ at the center $(i,j)$ position with kernel size of s.

This paper proposes to jointly explore the above loss to obtain the total training loss as follows:

where ${\lambda_{SSIM}}$ is the hyperparameter of the model that measures the importance of the SSIM loss.

(1) MVTec AD [4]. The MVTec AD dataset contains 5354 high-resolution color images in 15 different categories. Among them, There are 5 categories for textured images, such as “wood” or “leather”. The other 10 categories contain non-textured objects, such as “cable”. The challenge of this dataset comes from the fact that normal images and abnormal images come from the same category. The difference between the abnormal sample and the normal sample is subtle, such as scratches on the “leather”. Each category contains a dataset including only normal images for training and a dataset including normal and various abnormal images for testing.

(2) MNIST [56]. The dataset contains 10 different types of low-resolution grayscale images. A total of 70,000 images. 60,000 pictures for training, and 10,000 pictures for testing. In anomaly detection, a category is used as a normal sample and the remaining 9 categories are used as anomaly samples (called out-of-distribution detection [57]).

(3) Retinal-OCT [58] It is a medical dataset, which contains normal (healthy) retinal CT images and 3 categories of abnormal (damaged) retinal CT images.

(1) MVTec AD. For the MVTec AD dataset, this paper compares the state-of-the-art methods with distribution-based methods and AE methods. The distribution-based approach used pre-trained models with ImageNet dataset [64], such as CNN-Dict [45], MKDAD [32], SPADE [46], FAVAE [47], Cutpaste [48]. Among them, MKDAD uses the VGG-16 [65] model pre-trained on ImageNet, SPADE uses the ResNet-18 [66] model pre-trained on ImageNet, and CutPaste uses the EfficientNet (B4) [67] model pre-trained on ImageNet. AE methods include CAVGA-D_w [49], Puzzle-AE [50], ARFAD [24], memAE [10], AESc [51]. For example, Puzzle-AE introduces the complex self-supervised task of puzzle reduction, ARFAD introduces the self-supervised task of rotation and coloration, and MemAE introduces the traditional memory module.

From Table I, the proposed method achieves mean AUROC 91.8%. It is lower than Cutpaste due to the fact that Cutpaste uses a better pre-trained model and uses forged anomalous samples for training. However, the proposed method outperforms the remaining methods using pre-trained models by 4%-10%, validating the effectiveness of the AE combined with the PMB module. Among all AE methods, our method achieves state-of-the-art performance. It outperforms the CAVAG-D_w method that introduces real anomalous samples for training by 6%, validating that our method also performs very well without relying on anomalous samples. It is 13% and 8% higher than Puzzle-AE and ARFAD methods, indicating that our method can achieve excellent performance without designing complex image preprocessing operations. It outperforms the traditional memory model MemAE method by 11%, validating that the proposed PMB module stores normal features more efficiently.

(2) MNIST. For the MNIST dataset, this paper compares the state-of-the-art methods including DSVDD[28], CapsNetPP[52], OCGAN[53], LSA[54], MemAE[10], OCSVM[27], ARAE[18], , and MKDAD [32]. From Table II, it can be found that the proposed method achieves mean AUROC 98.1%[57] in the out-of-distribution detection setting. The advanced performance shows the superiority of the proposed method. It is slightly lower than the pre-trained method MKDAD method but far outperforms the remain distribution methods. Besides, compared with the traditional memory module MemAE, the proposed method outperforms MemAE, which validates the effectiveness of the novel query method and partition mechanism.

(3) Retinal-OCT. In the medical anomaly detection task, the proposed method achieves significant performance advantages. We compare with DSVDD [28], AnoGan [35], VAE-GAN [60], Cycle-GAN[61], GANomaly [36], P-Net [62], and MKDAD [32] on Retinal-OCT dataset. The results are shown in Table III. It can be found that the proposed method achieves the state-of-the-art performance with AUROC of 98.0%. It not only outperforms recent methods using GAN networks [60], [61], but even outperforms the MKDAD method. Our method does not rely on additional training data inductive bias specific to the pre-trained model. It is more suitable for medical image anomaly detection.

Compared with AE-based methods such as AE-SSIM [14] and AE-$L_{2}$ [14], the proposed method gains 5%-10% improvement. Because the proposed PMB module can store the learned normal sample features, effectively making the anomalous region reconstruction error larger. Compared with the two-stage method UTAD [63], the proposed method is simpler and more effective in training, and better localization results are obtained by exploring the corrected differential image $A^{\prime}$. Compared with the weakly supervised method CAVGA-$D_{w}$ [49] which uses additional abnormal region annotation information, the proposed method achieves better results. It shows that the proposed method outperforms the method with weakly supervised settings in the unsupervised setting, fully validating its superiority. To further demonstrate the localization performance of the proposed method, some examples of anomaly localization are illustrated in Figure 5. It can be observed that the proposed method can well solve the problem of abnormal region localization.

As shown in Figure 7, we show the reconstruction results of the PMB module for more abnormal and normal samples. In the ‘Normal sample’ on the left side of Figure 7, the reconstruction error is small. In the ’Abnormal Sample’ on the right, the reconstruction error of the anomalous region shown in the rectangular box is large, while the rest of the errors are small. It indicates that the PMB module does not weaken the reconstruction capability of AE, but makes it impossible for anomalous features to successfully address the abnormal feature output, while normal features can successfully address the normal feature output.

• $\bullet$

  AE-SSIM: Vanilla AE structure using SSIM loss function.

• $\bullet$

  AE-L₂: Vanilla AE structure using L₂ loss function.

• $\bullet$

  AE+skip: Vanilla AE structure using L₂ loss function with skip connections.

• $\bullet$

  MemAE: Vanilla AE structure using L₂ loss function with traditional memory module (Figure 1 (1)).

• $\bullet$

  DAAD: Vanilla AE is equipped with multi-scale block-wise memory module (Figure 1 (2)).

• $\bullet$

  DAAD+ : DAAD+ is DAAD complemented with the adversarially learned representation.

• $\bullet$

  OURS W/O H: Vanilla AE is equipped with PMB module.

• $\bullet$

  OURS: Vanilla AE is equipped with PMB module and parameter-free Histogram Error Estimation module.

The experimental results show that the proposed method improves 27%-20% compared to the vanilla AE method, and improves compared to both MemAE and DAAD methods. These indicate the effectiveness of the novel query generation method of the PMB module, which allows the memory module to store normal features. As in the ‘WOOD’ class, AE+skip performance can reach 97.7%, while MemAE, DAAD, DAAD+ all reduce the performance, and the proposed method further improves the anomaly detection performance. Comparing the results of MemAE, DAAD, and OURS W/O H, these show that the proposed query (Figure 1 (3)) is the most effective among the three query generation methods shown in Figure 1.