Focus the Discrepancy: Intra- and Inter-Correlation Learning for Image Anomaly Detection

Figure 2 overviews our proposed approach. The model consists of three branches: patch-wise discrepancy branch, intra-correlation branch, and inter-correlation branch. The input 2D image is first sent into a pre-trained backbone to extract multi-scale feature maps. At each feature level, we construct a subsequent transformer network to explicitly model the intra- and inter-image correlations for spotting the intra- and inter-discrepancies. Each network is implemented by stacking the I2Correlation blocks and feed-forward layers alternately. Suppose each network contains $K$ layers with length-$N$ input features $X\in\mathbb{R}^{N\times d}$, the output of $k$th layer $X_{k}\in\mathbb{R}^{N\times d_{m}}$ is calculated as follows:

where $X_{f}$ is the reference features used by the inter-correlation branch, $LN$ means LayerNorm, and $Z_{k}\in\mathbb{R}^{N\times d_{m}}$ is the $k$th layer’s hidden features. The final output features are calculated by linear projection: $\hat{X}=Z_{K}W_{o}$, where $W_{o}\in\mathbb{R}^{d_{m}\times d}$ is the output projection matrix.

Adopted from previous reconstruction-based AD methods, we employ feature reconstruction as our patch-wise recognition view for simplicity. In our approach, we construct a transformer network to reconstruct the input features. With the long-range dependency modeling ability, the features reconstructed by Transformer can have larger effective receptive fields [13], which are more conducive to detect hard global anomalies and logical anomalies. Moreover, with the further introduced inter-correlation branch (sec.3.4), the features can even perceive normal regions of images from the whole normal training set. Therefore, the reconstructed features generated by our model have better global perception and more discriminative semantic representation capability, which are more suitable for anomaly detection by patch-wise representation discrepancies.

Learning Objective. We can utilize classical reconstruction losses as the learning objective. We combine the $\ell 2$ distance and cosine distance to measure the feature distances between the reconstructed features $\hat{X}$ and the input features $X$. The loss function is defined as follows:

The intra-correlation learning branch aims to learn informative patch-to-patch correlations from the input patch sequence adaptively. As shown in the right part of Figure 2, we explicitly take advantage of the self-attention maps of transformers as intra-correlation matrices. Formally, the intra-correlation matrix of the $k$th layer is calculated by:

where $Q^{g}_{k},K^{g}_{k},V^{g}_{k}\in\mathbb{R}^{N\times d_{m}}$ represent the query, key, and value of the $k$th layer. $W_{k}^{Q_{g}},W_{k}^{K_{g}},W_{k}^{V_{g}}\in\mathbb{R}^{d_{m}\times d_{m}}$ represent the learnable projection matrices for $Q^{g}_{k},K^{g}_{k},V^{g}_{k}$, respectively. $S^{g}_{k}\in\mathbb{R}^{N\times N}$ denotes the learned intra-image patch-to-patch correlations. Since ${\rm Softmax}(\cdot)$ can convert the values in the similarity map into range $[0,1]$ along the horizontal axis, each row of $S^{g}_{k}$ can represent a discrete correlation distribution for each corresponding image patch. But different from the vanilla transformers, we further introduce target correlation matrices as the learning objective to explicitly optimize the intra-image correlations.

Target Correlation. Inspired by the contrastive learning method BarlowTwins [66], we can construct a target correlation matrix $T\in\mathbb{R}^{N\times N}$ as the learning target. The objective function of BarlowTwins measures the correlation matrix between the embeddings of two identical networks fed with distorted versions of a batch of samples, and tries to make this matrix close to the identity. In our work, the role of the target correlation is to introduce a prior correlation of patches as a pretext learning target, where each patch can be highly correlated to itself and also correlated to its neighborhood patches and the correlation decreases with the increase of distance. This allows us to optimize the intra- and inter-correlations in a self-supervised way. To this end, we use the radial basis function (RBF) to construct the target correlation matrix. We further adopt two learnable kernel variances $\sigma_{x}$ and $\sigma_{y}$ for horizontal and vertical axes to make the target correlation of each patch can adapt to the specific pattern of itself. The target correlation matrix of the $k$th layer is defined as:

where $||x_{ij}-x_{i^{\prime}j^{\prime}}||_{2}^{2}$ means the Square Euclidean distance between point $x_{ij}$ and $x_{i^{\prime}j^{\prime}}$, $i,i^{\prime}$ and $j,j^{\prime}$ means vertical and horizontal coordinates, respectively.

Next, we need to measure the distance between target- and intra-correlation distributions. This can usually be achieved by calculating the KL divergence. We can obtain a KL divergence value from each level of the network. Thus, we average all KL divergence values to combine the patch-to-patch correlations from multi-layer features into a more informative measure as follows:

where $KL(\cdot||\cdot)\in\mathbb{R}^{N}$ and its each element means the KL divergence between two discrete distributions corresponding to each row of $T^{g}_{k}$ and $S^{g}_{k}$.

Due to the rarity of anomalies and the dominance of normal patterns, the normal patches should build strong correlations with most patches in the whole image, while the weights of abnormal correlation distributions are harder to distribute to most patches and are more likely to concentrate on the adjacent image patches due to the neighborhood continuity. Since normal and abnormal patches have different correlation distributions, this is a distinguishable criterion for anomaly detection. Note that the intra-correlation branch explicitly exploits the spatial dependencies of each image patch, which are more informative than the patch-wise representations for anomaly detection.

Entropy constraint. Since the normal image patterns are usually diverse, the learned correlation distributions of the normal patches may also easily concentrate on the adjacent patches, which can cause the distinguishability between normal and abnormal to be downscaled. To address this, we further introduce an entropy constraint item for making normal patches establish strong associations with most normal patches in the whole image as much as possible. The entropy constraint item is defined as:

We will maximize the entropy constraint item. The loss function for the intra-correlation branch is defined as:

where $\lambda_{1}$ and $\lambda_{2}$ are used to trade off the loss items. The optimization of intra-correlation $\mathcal{S}^{g}$ is actually an alternating process with the guidance of target correlation $\mathcal{T}^{g}$ (see App. A.1), ultimately resulting that each normal patch can establish strong correlations with most normal patches.

Through the intra-correlation branch, we can establish patch-to-patch correlations within a single image. However, an image usually doesn’t contain all possible normal patterns, which may cause it difficult to distinguish some ambiguous abnormal patches (see rows 2, 3, 4 in Figure 1) only through the intra-correlations. To address this, we should effectively take advantage of the known normal patterns from the normal training set, which are more likely to contain more informative normal patterns. Specifically, we further propose an inter-correlation learning branch to explicitly model pairwise correlations with the whole normal training set. In this branch, the features of each patch establish a discrete inter-correlation distribution with the reference features extracted from all normal samples (see Figure 2). The inter-correlation matrix of the $k$th layer is similar to the corresponding intra-correlation, and is defined as:

where $X_{f}\in\mathbb{R}^{N_{e}\times d_{e}}$ represents the external reference features, $N_{e}$ is the length of the reference features and $d_{e}$ is the feature dimension. $W_{k}^{Q_{e}}\in\mathbb{R}^{d_{m}\times d_{m}}$ and $W_{k}^{K_{e}},W_{k}^{V_{e}}\in\mathbb{R}^{d_{e}\times d_{m}}$ are learnable matrices for $Q^{e}_{k},K^{e}_{k},V^{e}_{k}$. $S^{e}_{k}\in\mathbb{R}^{N\times N_{e}}$ denotes the learned inter-image correlations.

Loss. The loss function for the inter-correlation branch has the opposite optimization direction to $\mathcal{L}_{g}$. Because the external reference features contain more comprehensive normal patterns, normal patches can establish stronger correlations with the closest reference normal patterns, instead of establishing strong correlations with most reference patterns. By contrast, it shall be harder for anomalous patches to establish strong correlations with any of the reference patterns. So the inter-correlation distributions of anomalies are more dispersed, while the normal inter-correction distributions are more likely to be concentrated. To this end, we maximize the KL divergence and minimize the entropy item in the training process. The practical optimization strategy is also opposite to the intra-correlation branch (see App. A.1 for details). The loss function is defined as:

External Reference Features. External reference features are used for providing accumulated knowledge of normality for the inter-correlation learning branch. Thus, these features should represent all possible normal patterns of all normal samples from the whole normal training set. To this end, we can employ many methods to generate the reference features, such as sampling key features by coreset subsampling algorithm [35], generating prototype features by memory module [28], or learning codebook features through vector quantization [65] or sparse coding techniques [55]. However, because the RBF-kernel in $\mathcal{T}^{e}$ is position-sensitive, our reference features are better to preserve the positional information. The ablation results in App.B.2 show that the methods that can’t preserve the position information perform worse. From comprehensive ablation studies, we find that using patch-wise averaged features as the external reference features is a simple but effective way. We think that features extracted by deep neural networks are highly redundant [16, 52], and different normal patterns generally correspond to larger activation values at different channels in the feature vector [61]. So feature averaging will not eliminate some rare normal patterns [11], these patterns may be preserved at specific channels. And averaging can greatly reduce the feature redundancy, making the obtained reference features more representative. Formally, for position $(i,j)$, we first extract the set of patch features at $(i,j)$, $X_{ij}=\{x_{ij}^{k}\},k\in[1,N]$ from the $N$ normal training images. Then, the reference features at position $(i,j)$ is computed as $x_{ij}^{f}=\frac{1}{N}\sum_{k=1}^{N}x_{ij}^{k}$. The final external reference features are composed of averaged features at all locations and then flattened into 1D: $X_{f}={\rm Flatten}(\{x_{ij}^{f}\})$.

We further combine the intra-correlation and inter-correlation branches to form the I2Correlation block. The output features of the intra- and inter-correlation branches are defined as: $Z^{g}_{k}=S^{g}_{k}V^{g}_{k}$ and $Z^{e}_{k}=S^{e}_{k}V^{e}_{k}$, respectively. Then, we use the residual feature $Z^{ge}_{k}=Z^{g}_{k}-Z^{e}_{k}$ as the output of the I2Correlation block. The feature $Z^{e}_{k}$ is generated from the external reference features, which can contain rich normal patterns. Thus, by subtracting $Z^{e}_{k}$ from the feature $Z^{g}_{k}$, it is conducive to spotlight the abnormal patterns in the $Z^{g}_{k}$. This is beneficial for anomaly detection.

The total loss function consists of the there branch loss functions, and is combined as follows:

We utilize reconstruction errors as the baseline anomaly criterion and incorporate the normalized correlation distribution distances into the reconstruction criterion. The final anomaly score of the $i$th patch in the input patch sequence is shown as follows:

where $\odot$ is the element-wise multiplication.

Datasets. We extensively evaluate our approach on two widely used industrial AD datasets: the MVTecAD [3] and BTAD [26], and one recent challenging dataset: the MVTecAD-3D [5]. MVTecAD is established as a standard benchmark for evaluating unsupervised image anomaly detection methods. This dataset contains 5354 high-resolution images from 15 real-world categories. 5 classes consist of textures and the other 10 classes contain objects. A total of 73 different anomaly types are presented. BTAD is another popular benchmark for unsupervised image anomaly detection, which contains 2540 RGB images of three industrial products. All classes in this dataset belong to textures. MVTecAD-3D is recently proposed, which contains 4147 2D RGB images paired with high-resolution 3D point cloud scans from 10 real-world categories. Even though this dataset is mainly designed for 3D anomaly detection, most anomalies can also be detected only through RGB images without 3D point clouds. Since we focus on image anomaly detection, we only use RGB images of the MVTecAD-3D dataset. We refer to this subset as MVTec3D-RGB.

Evaluation Metrics. The standard metric in anomaly detection, AUROC, is used to evaluate the performance of AD methods. Image-level AUROC is used for anomaly detection and a pixel-based AUROC for evaluating anomaly localization [3, 4, 11, 64].

Implementation Details. We use EfficientNet-b6 [45] to extract two levels of feature maps with $\{8\times,16\times\}$ downsampling ratios, the pre-trained networks are from the timm library [53]. Then, we construct a subsequent transformer network (see Figure 2) at each feature level to reconstruct patch features and learn patch-to-patch correlations. The parameters of the feature extractor are frozen in the training process, only the parameters of the subsequent transformer networks are learnable. All the subsequent transformer networks in our model contain 3 layers. We set the hidden dimension $d_{m}$ as $\{256,512\}$ and the number of heads as $8$. The hyperparameters $\lambda_{1}$ and $\lambda_{2}$ are set as $0.5$ and $0.5$ to trade off two parts of the $\mathcal{L}_{g}$ and $\mathcal{L}_{e}$ loss functions (see App. B.2 for hyperparameter sensitivity analysis). We use the Adam [19] optimizer with an initial learning rate of $1{\rm e}^{-4}$. The total training epochs are set as 100 and the batch size is 1 by default. All the training and test images are resized and cropped to $256\times 256$ resolution from the original resolution.

SOTA Methods. We extensively compare our method with those published SOTA methods in the past three years. The comparison results on MVTecAD are shown in Table 1. Then, we select five reproducible methods to report the detailed results on MVTecAD, these methods are representative and SOTA AD methods in the mainstream categories of image anomaly detection as we discuss in Related Work, including: DRAEM [64], PatchSVDD[60], PatchCore [35], MKD[41], CFLOW[15]. For a fair comparison, we reproduce all these methods with the same backbone as in our model. Thus, despite using the unmodified code from the official repositories, we are not able to exactly reproduce the original results, but our numbers are very close. The detailed image-level AUROC results are shown in Table 2, and detailed pixel-level AUROC results are in the App. Table 6. Most of the methods in Table 1 don’t report results on the BTAD and MVTec3D-RGB datasets, so we reproduce the five representative methods on the two datasets for comparison. The image-level and pixel-level AUROC results on BTAD and MVTec3D-RGB are shown in Table 3. Additional detailed results are in the App. Table 7, 8.

Anomaly Detection. On MVTecAD, we set the SOTA performance on the mean detection AUROC, which is slightly higher than the best competitor, PatchCore [35]. Note that the results in Table 2 show that our method can achieve much better results than PatchCore when using the same backbone. What’s more, in addition to the pill, screw and toothbrush classes, our method achieves more than 99% AUROC in all other classes, while other methods only achieve more than 99% AUROC in most nine classes. This shows that our method is more stable and effective in real-world applications. In the classes (e.g. metal nut, screw, and transistor) with more global and logical anomalies, our method can achieve significantly better results than those methods depending on patch-wise discrepancy (e.g. DRAEM [64] and MKD [41]), and can also achieve better results compared with other methods. This verifies that the three discrepancies are complementary factors and our model can simultaneously spot these discrepancies to recognize harder global and logical anomalies. On BTAD, our FOD can achieve 96.0% mean detection AUROC, which can outperform the best competitor CFLOW [15] by a margin of 1.2%. On MVTec3D-RGB, we can outperform all previous SOTA methods by a margin of 3.3%. Note that this dataset is much more challenging than the MVTecAD dataset when comparing the best results (99.2% for MVTecAD vs. 88.4% AUROC for MVTec3D-RGB). This demonstrates the robustness of our method.

Anomaly Localization. Our method can achieve comparable results with the best competitors on all three datasets. Our method is slightly lower than CFLOW [15] and PatchCore [35] on the three datasets, but we observe that our model that considers the intra- and inter-correlations outperforms the vanilla reconstruction AD models (e.g. DRAEM [64] and MKD [41]) on all three datasets, which verifies the effectiveness of correlation modeling. Compared with patch-wise reconstruction methods, our method is more effective and robust to detect hard global and logical anomalies (see Figure 3).

Qualitative Results. Samples in Figure 1 show qualitative impressions of accurate anomaly localization from our method. It can be found that our approach can achieve the best anomaly score maps by combining the three discrepancies. Additional visualizations are in the App. Figure 6.

Ablation Study. To explain how our model works effectively, we further investigate the effect of the three key designs in our model: recognition views, entropy constraint and reference features. The quantitative results are shown in Table 4, more results can be found in App. Table 9. The entropy constraint is quite effective and necessary in the intra-correlation branch. Specifically, it brings a remarkable 7.7% averaged absolute AUROC promotion, which demonstrates that the entropy constraint strategy is really conducive to increase the distinguishability between abnormal and normal. Only utilizing pure reconstruction criterion or pure KL divergence can’t get the most superior detection results, the combined criterion can outperform each single criterion consistently by a margin of 1.0% and 9.8%. Thus the reconstruction errors and the intra- and inter-correlations can collaborate to improve detection performance. For the external reference features, we compare the simple mean features with more elaborate coreset features [35] (see App. Table 15 for more comparison methods). The results show that the simple mean features are more effective than the coreset features (0.923 vs. 0.836), which means that mean features are effective enough to represent all possible normal patterns and preserve the positional information. Finally, our proposed FOD surpasses the pure reconstruction Transformer by 4.0% absolute improvement. These verify that our proposed explicit correlation learning approach is effective.

Effect for Different Anomalies. To illustrate the effectiveness of our model intuitively, we explore the anomaly localization quality under different types of anomalies (see Figure 3). E.g., simple local anomalies (Column 1) are obviously different from normal visuals; hard global anomalies (Columns 2,3) have less obvious visual appearances, which usually need to be effectively compared with normal regions to detect; logical anomalies (Columns 4,5) may be locally normal and can be detected correctly through the overall semantic understanding. We can find that our FOD is more distinguishable in general. For simple local anomalies, patch-wise reconstruction methods can also achieve good results. While for hard global anomalies, these methods cannot detect or only detect partial anomalies. However, for logical anomalies, these methods cannot detect them at all. In contrast, our method is conducive to detect diverse anomalies because of its complementary recognition views. This verifies that our method can make more precise detection and reduce the false-negative rate compared with the pure patch-wise reconstruction methods.

Correlations Visualization. To explain what our model has learned intuitively, we visualize some learned correlations in Figure 4 with the MVTecAD dataset. For the intra-correlation branch, the normal patches learn to build strong correlations with most patches in the whole image, while the correlation distributions of anomalies usually concentrate in the adjacent image patches. It can be found that the normal correlations are much spread and the abnormal correlations are more concentrated in the adjacent regions. For the inter-correlation branch, most normal correlations only concentrate in one point and the abnormal correlations are more spread, which exactly means that each normal patch can build strong correlations with one special normal pattern and abnormal patches are harder to establish correlations with normal patterns.