ReContrast: Domain-Specific Anomaly Detection via Contrastive Reconstruction

First, we revisit a simple yet powerful epistemic UAD method based on feature reconstruction, i.e., Reverse Distillation (RD4AD) [2]. The method is illustrated in Config.A of Figure 2. RD4AD consists of an encoder (teacher), a bottleneck, and a reconstruction decoder (student). Without loss of generality, the first three layers of a standard 4-layer network are used as the encoder. The encoder extracts informative feature maps with different spatial sizes. The bottleneck is similar to the fourth layer of the encoder but takes the feature maps of the first three layers as the input. The decoder is the reverse of the encoder, i.e., the down-sampling operation at the beginning of each layer is replaced by up-sampling. During training, the decoder learns to reconstruct the output of the encoder layers by maximizing the similarity between feature maps. During inference, the decoder is expected to reconstruct normal regions of feature maps but fails for anomalous regions as it has never seen such samples.

Let $f_{E}^{k},f_{D}^{k}\in\mathbb{R}^{C^{k}\times H^{k}\times W^{k}}$ denote the output feature maps from the $k^{th}$ layer of the encoder and decoder respectively, where $C^{k}$, $H^{k}$, and $W^{k}$ denote the number of channels, height, and width of the $k^{th}$ layer, respectively. Regional cosine distance [2; 6; 24] is used to measure the difference between $f_{E}^{k}$ and $f_{D}^{k}$ point by point, obtaining a 2-D distance map $\mathcal{M}^{k}\in\mathbb{R}^{H^{k}\times W^{k}}$:

$$\mathcal{M}^{k}\left({h,w}\right)=1-\frac{f_{E}^{k}\left({h,w}\right)^{T}\cdot f_{D}^{k}\left({h,w}\right)}{\left\|{f_{E}^{k}\left({h,w}\right)}\right\|~{}\left\|{f_{D}^{k}\left({h,w}\right)}\right\|},$$

(1)

where $\left\|\cdot\right\|$ is $\ell_{2}$ norm, and a set of $(h,w)$ locates a spatial point in the feature map. The calculation is depicted in Figure 3(a). A large value in $\mathcal{M}^{k}$ indicates anomaly at the corresponding position. Final anomaly score map $\mathcal{S}^{map}$ is obtained by up-sampling all $\mathcal{M}^{k}$ to image size and a sum operation. The maximum value of $\mathcal{S}^{map}$ is taken as the anomaly score for the image, denoted as $\mathcal{S}^{img}$. During training, the anomaly maps are minimized on normal images by the regional cosine distance loss, given as:

$$\mathcal{L}_{region}={\sum\limits_{k=1}^{3}{\frac{1}{~{}H^{k}W^{k}}{\sum\limits_{h=1}^{H^{k}}{\sum\limits_{w=1}^{W^{k}}{1-\frac{f_{E}^{k}\left({h,w}\right)^{T}\cdot f_{D}^{k}\left({h,w}\right)}{\left\|{f_{E}^{k}\left({h,w}\right)}\right\|~{}\left\|{f_{D}^{k}\left({h,w}\right)}\right\|}}}}.}}$$

(2)

The optimization of $\mathcal{L}_{region}$ is illustrated in the upper-left of Figure 3(b). Notably, though $\mathcal{L}_{region}$ was proposed clearly as the loss function of RD4AD [2], the official code¹¹1https://github.com/hq-deng/RD4AD implemented a different function that makes an extra dimension flatten operation (maybe by coding bug). However, we observe extreme instability during the training with $\mathcal{L}_{region}$, producing degenerated I-AUROC of 95.3% on MVTec AD compared with 98.5% reported in the paper. The I-AUROC of 95.3% is more consistent with other methods using $\mathcal{L}_{region}$ [6; 24]. In the next subsection, we will show that the seemingly unreasonable "bug" in the code improves the performance by explicitly introducing globality into optimization.

It is a convention to use global average pooling (GAP) to pool the 2-D feature map of the last network layer to a 1-D representation, in both supervised classification and self-supervised contrastive learning. In UAD, GAP will simply mess up all feature points together, losing the ability to distinguish normal and anomalous regions, as shown in the upper-right of Figure 3(b). Therefore, we try to directly bring globality into the loss function while maintaining the point-to-point correspondence between $f_{E}^{k}$ and $f_{D}^{k}$. Let $\mathcal{F}$ denote a flatten operation that casts a 2-D feature map $f\in\mathbb{R}^{C\times H\times W}$ to a vector ${v}\in\mathbb{R}^{CHW}$. Our proposed global cosine distance loss is given as:

$$\mathcal{L}_{global}={\sum\limits_{k=1}^{3}{1-\frac{{\mathcal{F}\left(f_{E}^{k}\right)}^{T}\cdot\mathcal{F}\left(f_{D}^{k}\right)}{\left\|{\mathcal{F}\left(f_{E}^{k}\right)}\right\|~{}\left\|{\mathcal{F}\left(f_{D}^{k}\right)}\right\|}}},$$

(3)

where the number of cosine dimension is $C\cdot H\cdot W$ instead of $C$ in $\mathcal{L}_{region}$, as depicted in Figure 3(a). Intuitively, $\mathcal{M}^{k}$ is also minimized along with the minimization of $\mathcal{L}_{global}$. It is easy to prove that minimizing $\mathcal{L}_{global}$ to 0 equals to minimizing $\mathcal{M}^{k}$ to 0. Though losses are not completely optimized to 0 in neural network training, the regional $\mathcal{S}^{map}$ still works as the anomaly map. The model trained with $\mathcal{L}_{global}$ is named as Config.B.

As shown in Figure 5(a), the performance of Config.B (blue) is much more favorable and stable than Config.A (purple) during training. Here, we briefly analyze the underlying mechanism by plotting the landscape of $\mathcal{S}^{map}$ (average of all samples) against model parameters near optimized minima in Figure 4. We observe two differences. First, the overall landscape of the model optimized by $\mathcal{L}_{global}$ is flatter than that by $\mathcal{L}_{region}$, which implies $\mathcal{S}^{map}$ is more stable during training with $\mathcal{L}_{global}$. The instability of $\mathcal{L}_{region}$ can be caused by the point-by-point distance, where the fitting of one region may cause under-fitting of the others. $\mathcal{L}_{global}$ can be regarded as the distance between manifolds of feature points, as depicted in the lower-left of Figure 3(b), which measures the consistency globally instead of excessively focusing on individual regions. Second, the landscape around the final minima of the model trained by $\mathcal{L}_{global}$ is sharper than that by $\mathcal{L}_{region}$. Previous explorations on the geometry of loss landscapes suggest that flat minima generalize better on unseen samples [25; 26], which is unwanted in UAD settings. The real mechanism can be more complicated and is still under further investigation.

It seems easy to adapt the network in the target domain by jointly optimizing the encoder and decoder, as illustrated in Figure 2 Config.C. However, previous attempts [9] suggest that training encoders would lead encoders to extract indiscriminative features that are easy to reconstruct, which is also described as pattern collapsing or converging to trivial solutions [2]. We observe a similar phenomenon that the training loss quickly decreases to nearly zero, as shown by the orange dotted line in Figure 5(a).

To inspect the behavior of the encoder, we probe the diversity of encoder feature maps by the standard deviation (std) of $\ell_{2}$ normalized feature points ${f_{E}^{k}(h,w)}/\left\|{f_{E}^{k}(h,w)}\right\|_{2}$. The larger the std, the more discriminative the feature, and the more distinctive between different regions. We plot the feature diversity of outputs of the $2^{nd}$ encoder layer during training in Figure 5(b). The std of Config.C drops quickly during training, while the std of Config.B keeps constant as a comparison. In addition, the std of Config.C does not completely collapse to zero, which explains the passable performance. The outputs of other middle layers follow the same trend. The observation indicates that the degeneration of performance is attributed to the degradation of encoder features. It is noted that the std starts from different values because of the different modes of batch normalization (BN) [28] during training.²²2Encoder BN is in eval mode in Config.A using pre-trained running mean and variation, and it is in train mode using batch mean and variation if optimized.

Contrastive learning for SSL maximizes the similarity between the representations of two augmentations (views) of one image while trying to avoid collapsing solutions in which all inputs are encoded to a constant value. SimSiam [16] indicated that the stop-gradient operation plays an essential role in preventing collapsing for SSL methods with positive pairs only [16; 17]. Intriguingly, we found that contrastive SSL can be explained from the perspective of feature reconstruction. The predictor (a multi-layer perceptron, MLP) can be regarded as a reconstruction network that constructs the representation of one view $x$ to match the representation of another view $x\textquoteright$, as in Figure 1. Vice versa, feature reconstruction UAD can be transformed into contrastive learning as well, regarding the reconstruction network as a predictor.

To this end, we introduce the stop-gradient operation into feature reconstruction, transforming it into a contrastive-like paradigm. The network is optimized by contrasting the feature maps of the encoder and decoder, as shown in Config.D of Figure 2. The gradients do not propagate directly into the encoder, but back into the encoder through the decoder. It is implemented by modifying (3) as:

$$\mathcal{L}_{global}={\sum\limits_{k=1}^{3}{1-\frac{sg\left({\mathcal{F}\left(f_{E}^{k}\right)}\right)^{T}\cdot\mathcal{F}\left(f_{D}^{k}\right)}{sg\left(\left\|{\mathcal{F}\left(f_{E}^{k}\right)}\right\|\right)~{}\left\|{\mathcal{F}\left(f_{D}^{k}\right)}\right\|}}}$$

(4)

where sg denotes the stop-gradient operation. The training of this configuration can be intuitively explained as a "mutual reinforcement". The optimization of $f_{D}^{k}$ drives the encoder to be more specific and informative in the target domain by the gradients from the decoder, while the more domain-specific $f_{E}^{k}$ of encoder requires further optimization of $f_{D}^{k}$. There are previous works trying to explain the deeper mechanism of stop gradient in contrastive learning. In [16], the authors hypothesize this configuration as an implementation of an Expectation-Maximization (EM) algorithm that implicitly solves two underlying sub-problems with two sets of variables. In [29], the authors argue that there are flaws in the hypothesis of [16] and suggest that the decomposed gradient vector (center gradient) helps prevent collapse via the de-centering effect.

As shown in Figure 5(a) and (b), Config.D can constantly extract more discriminative features without degeneration, boosting it to perform better than Config.C at the beginning of training. Previous works [9; 3] suggested that the reconstruction decoder may learn an "identical shortcut" that well reconstructs both seen and unseen samples, which explains the performance drop of Config.D after reaching the peak in Figure 5(a). To handle this issue, contrastive pairs are introduced.

Contrastive learning methods make use of (positive) contrastive image pairs, based on the intuition that the semantic information of different augmented views of the same image should be the same. Without image augmentations ($x$=$x^{\prime}$), contrastive learning degrades into a self-reconstruction paradigm that the predictor predicts the input of itself, which may also collapse into an "identical shortcut" just like Config.D. In epistemic UAD methods, it is impossible to construct contrastive pairs by image augmentation, because 1) any augmentation (flip, cutout, color jitter, etc.) could be potential anomalies; and 2) the feature maps of two augmented views lost the regional correspondence after spatial augmentations (shift, rotate, scale, etc.). Therefore, we propose an augmentation-free method to construct two views of one image.

In previous explorations, the encoder is adapted to the target image domain and generates feature representations in a domain-specific view. We additionally introduce a frozen encoder network without any adaptation throughout the training, which perceives images in the view of pre-train image domain. As shown in Config.E of Figure 2, the decoder and bottleneck take the features of the domain-specific encoder to reconstruct the frozen encoder; vice versa, they simultaneously take the features of the frozen encoder to reconstruct the domain-specific encoder, building a complete two-view contrastive paradigm. This configuration can also be interpreted as cross-reconstruction, as the decoder reconstructs the information of the frozen encoder given the input of the adapted encoder; and reconstructs the information of the adapted encoder given the input of the frozen encoder. With six pairs of feature maps, the $\mathcal{S}^{map}$ is obtained by adding up the six up-sampled cosine distance maps. As shown in Figure 5(b), the encoder feature diversity of Config.E not only holds but slightly increases during training.

Due to network information capacity, intrinsic reconstruction errors exist in available normal regions. The errors of normal details and edges (hard-normal) are generally higher than other plain regions (easy-normal), leading to confusion between the intrinsic error of hard-normal regions and the epistemic error of unseen anomalies, as demonstrated in the lower-left of Figure 5(c). It was proposed to hard-mine the hard-normal regions by simply discarding the "easy" pairs of $f_{E}^{k}\left({h,w}\right)$, $f_{D}^{k}\left({h,w}\right)$ with small cosine distance in $\mathcal{L}_{region}$ [24]. However, arbitrarily discarding feature points in $f_{D}^{k}$ breaks the global feature point manifolds in $\mathcal{L}_{global}$. Therefore, we propose a hard mining strategy, namely $\mathcal{L}_{global-hm}$, focusing on the optimization of hard-normal regions to widen the gap between the epistemic errors caused by unseen anomalies and the errors of all normal regions. The gradients of easy $f_{D}^{k}(h,w)$ are discarded instead of eliminating $f_{D}^{k}(h,w)$ itself, by modifying $f_{D}^{k}$ as:

$$f_{D}^{k}(h,w)=\left\{\begin{matrix}{sg\left({f_{D}^{k}\left({h,w}\right)}\right),~{}if~{}\mathcal{M}^{k}\left({h,w}\right)<\mu\left({\mathcal{M}^{k}\left({h,w}\right)}\right)~{}\text{+}~{}\alpha~{}*\sigma\left({\mathcal{M}^{k}\left({h,w}\right)}\right)}\\ {f_{D}^{k}\left({h,w}\right),~{}else}\\ \end{matrix},\right.$$

(5)

where $\mu\left({\mathcal{M}^{k}\left({h,w}\right)}\right)$ is the average regional distance on the dataset (approximated on mini-batch), $\sigma\left({\mathcal{M}^{k}\left({h,w}\right)}\right)$ is the standard deviation of the regional distance, and $\alpha$ is a hyper-parameter to control the discarding rate. The optimization of $\mathcal{L}_{global-hm}$ is depicted in the lower-right of Figure 3(b). As shown in Figure 5(c), the model trained by $\mathcal{L}_{global-hm}$ activates fewer normal regions than the model trained by $\mathcal{L}_{global}$. The usage of hard mining strategy can also mitigate the over-fitting of easy training samples.

Datasets. MVTec AD is the most widely used industrial defect detection dataset, containing 15 categories of sub-datasets (5 textures and 10 objects). VisA is a challenging industrial defect detection dataset, containing 12 categories of sub-datasets. OCT2017 is an optical coherence tomography dataset [21]. APTOS is a color fundus image dataset, available as the training set of the 2019 APTOS blindness detection challenge [22]. ISIC2018 is a skin disease dataset, available as task 3 of ISIC2018 challenge [23]. Detailed information is presented in Appendix B.

Implementation. WideResNet50 [30] pre-trained on ImageNet [1] is utilized as the encoder by default. AdamW optimizer [31] is utilized with $\beta$=(0.9,0.999) and weight decay=1e-5. The learning rates of new (decoder and bottleneck) and pre-trained (encoder) parameters are 2e-3 and 1e-5, respectively. The network is trained for 3,000 iterations on VisA, 2,000 on MVTec AD and ISIC2018, and 1,000 on APTOS and OCT2017. The $\alpha$ in equation (5) linearly rises from -3 to 1 in the first one-tenth iterations and keeps 1 for the rest training. More experimental details and environments are presented in Appendix C.

Metrics. Image-level anomaly detection performance is measured by the Area Under the Receiver Operator Curve (I-AUROC). F1-score (F1) and accuracy (ACC) are adopted for medical datasets following [32]. The operating threshold of F1 and ACC is determined by the optimal value of F1. For anomaly segmentation, pixel-level AUROC (P-AUROC) and Area Under the Per-Region-Overlap (AUPRO) [33] are used as the evaluation metrics. AUPRO is more meaningful than P-AUROC because of the unbalanced amount of normal and anomalous pixels [20].

Anomaly detection results on MVTec AD are shown in Table 1. Our approach achieves a superior average I-AUROC of 99.5%. In the aspect of error (100%$-$I-AUROC), we achieve only 0.5%, reducing the previous SOTA error of PatchCore (0.9%) by relatively 44%. Anomaly segmentation results are shown in Table 2 (average of 15 categories). Our ReContrast achieves SOTA performance measured by both P-AUROC and AUPRO. In AUPRO, we produce 95.2%, exceeding the previous SOTA RD4AD by 1.3%. We also present SOTA performances on texture (Text. Avg.) and object (Obj. Avg.) categories, respectively.

Anomaly detection and segmentation results on VisA are shown in Table 3 (average of 12 categories). Our approach achieves a superior average I-AUROC of 97.5%, exceeding previous SOTA RD4AD by 1.5%. In AUPRO, we produce 92.6%, outperforming previous SOTA PatchCore by 1.4%.

Conventional UAD settings train separate models for different object categories. In UniAD [9], the authors proposed to train a unified model for multiple classes. We evaluate our ReContrast under the unified setting on MVTec AD and VisA. Our model is trained for 5,000 iterations on all images containing all 15 MVTec AD categories (or 12 categories in VisA) and evaluated on each category following [9]. The results are presented in Table 4. On MVTec AD, our method produces an I-AUROC of 98.2%, exceeding the previous SOTA 96.5% of UniAD by a large 1.7%. On VisA, both our ReContrast (95.1%) and baseline RD4AD (92.7%) outperform the previous multi-class SOTA UniAD (91.5%), while the proposed method exceeds RD4AD by 2.4%.

The experimental results on three medical image datasets are presented in Table 5. Apart from industrial UAD methods, we also compare our ReContrast with approaches designed for medical images (f-AnoGan [7], GANomaly [8], AE-flow [32]), and contemporaneous methods proposed for better domain adaption ability (CFA [14], SimpleNet [15]). We reproduce the methods that were not evaluated on the corresponding dataset if they are open-sourced. Our method achieves the best results in all metrics on APTOS and ISIC2018, and comparable best results on the relatively easy OCT2017, demonstrating that ReContrast can directly adapt to various medical image modalities.

To verify the effect of the proposed elements, we conduct thorough ablation studies on MVTec AD and APTOS. Because some configurations are unstable during training on some categories, we report the last iteration’s performances and those of the best checkpoint during training (best by I-AUROC, evaluate every 250 iterations). The results are reported in Table 6. The use of $\mathcal{L}_{global}$ (Config.B) boosts the performance of Config.A trained by $\mathcal{L}_{region}$. Directly optimizing the encoder and decoder together (Config.C) causes a great degeneration because of pattern collapse, while either stop-gradient (Config.D) or contrastive pairs (Config.C$+cp$) recovers some of it. By introducing both (Config.E), the performance is further improved to exceed the frozen encoder baseline. Finally, training with $\mathcal{L}_{global-hm}$ yields SOTA performances (Ours). Meanwhile, $\mathcal{L}_{global-hm}$ can also improve the baseline reconstruction method alone (Config.B$+hm$). More ablation studies, experimental results with error bars, limitations, and qualitative visualization are presented in Appendix.