A SAM-guided Two-stream Lightweight Model for Anomaly Detection

SAM (Kirillov et al., 2023) has attracted notable academic interest owing to its remarkable zero-shot capabilities and versatility across various vision applications. In our paper, we aim to harness the robust generalization abilities of SAM and render SAM more mobile-friendly.

First, we follow the idea of encouraging the student network to generate anomaly-specific features that diverge from those of the teacher network adopted by (Zhang et al., 2023d, c). This strategy effectively helps overcome the limitations posed by network architecture similarities and identical data flows. However, as (Zhang et al., 2023b) said, the test data can have large distribution shifts arising in many real-world applications and may correspond to multiple normal patterns, leading to imprecise representations of the teacher network. The plain student stream is trained to produce discriminative and general feature representations in both normal and anomalous regions as a superior alternative to the teacher network, demonstrated in Section 4.4. Similarly guided by SAM’s knowledge, the denoising student stream aligns its feature representations with those of the same images without any corruption, which effectively enhances the differentiation of the representations from two student streams when dealing with anomalous regions.

Second, the principal aim of our paper is to develop a model that can deliver satisfactory performance while significantly reducing the number of parameters and inference time. SAM has demonstrated its ability to function on resource-constrained devices, primarily thanks to its lightweight mask decoder. However, the default image encoder in the original SAM relies on ViT-H, which boasts over 600M parameters and is considered heavyweight. Considering the 3.2M decoder lightweight enough, We transition the default image encoder to ViT-Tiny (Wu et al., 2022) while kee** the original architecture of the mask decoder.

Consequently, SAM serves as the teacher network for the TLM. The output feature maps are extracted from the two-layer decoder of SAM, shown in Figure 4 (a). With no need for segmentation outcomes, we only use the Two-Way-Transformer blocks of the mask decoder, and introduce our FA module. The output features from the pre-trained SAM with normal images as input are denoted as $T_{P}^{1}$ and $T_{P}^{2}$, respectively. Similarly, those originating from pseudo-anomalous images are labeled $T_{D}^{1}$ and $T_{D}^{2}$. TLM adopts the image encoder of mobileSAM (Zhang et al., 2023a), i.e., ViT-Tiny (Wu et al., 2022), and their output features are denoted as $S_{P}^{1}$, $S_{P}^{2}$, $S_{D}^{1}$, and $S_{D}^{2}$, respectively.

To distill knowledge from SAM to the plain student stream, we minimize the cosine distance between features from $T_{P}^{k}$ and $S_{P}^{k}$, where $k$ = 1,2. Additionally, we minimize the cosine distance between features from $T_{D}^{k}$ and $S_{D}^{k}$, for $k$ = 1,2, to supervise the denoising student stream in reconstructing normal features. The cosine similarity can be computed through Equation (2) and the loss function for optimizing the network is formulated as Equation (3) and Equation (4).

where $k$ is the number of feature layers used in training. $i$ and $j$ stand for the spatial coordinate on the feature map.In particular, $i=1\ldots H_{k}$ and $j=1\ldots W_{k}$. $H$ and $W$ denote the height and width of $k$-th feature map. The cosine similarity between features from SAM and the plain stream denotes $X_{k}^{P}(i,j)$, while that between features from SAM and the denoising stream denotes $X_{k}^{D}(i,j)$.

Inspired by (Zhang et al., 2023a), the image embeddings generated by the student encoder can closely approximate that of the original teacher encoder. This observation leads to the conclusion that the utilization of distinct decoders for the TLM may not be imperative. Instead, a shared two-layer decoder as illustrated in Figure 4, is an efficient alternative. We observe that the shared decoder, while economical in terms of parameters and computation time, achieves a satisfactory result with a slight decrease compared to employing two separated decoders, shown in Figure 5.

The difference between our STLM and DeSTSeg (Zhang et al., 2023c) lies in 1) our method learns not only a denoising student but also a plain student in our TLM. The plain student stream can learn generalized knowledge related to the anomaly detection task, which enables it to effectively represent normal features and even previously unseen patterns in the training set, ensuring differences between features of two streams are well-captured. 2) The high-generalization-capable SAM, as the teacher network of our TLM, also helps the denoising student to generate high-quality reconstructed results for genuine anomalies.

Besides, it is worth noting that our method exhibits superior performance and a more compact model compared to DeSTSeg.

We validate the effectiveness of our method using the MVTec AD (Bergmann et al., 2019a) dataset, a renowned benchmark in the field of anomaly detection and localization. MVTec AD comprises 5 texture and 10 object categories, each providing hundreds of normal images for training, along with a diverse set of both anomalous and normal images for evaluation. It also provides pixel-level ground truths for defective test images.

Additionally, we extend our experimentation on the VisA (Zou et al., 2022) and DAGM (Wieler and Hahn, 2007) dataset to showcase the generalization capabilities of STLM facing more complex datasets with. The VisA Dataset comprises 10,821 high-resolution color images (9,621 normal and 1,200 anomalous samples) encompassing 12 objects across 3 domains, establishing it as the most extensive anomaly detection dataset in the industrial domain to date. DAGM contains 10 textured objects with small abnormal regions that bear a strong visual resemblance to the background. For each class, we first move all anomalous samples from the original training set to the original test set and then move all normal samples from the test set to the training set. Furthermore, we randomly select 30 normal samples from the training set, designating them as ”good” samples in the test set, which are subsequently removed from the training set.

Additionally, we use the Describable Textures Dataset (DTD) (Cimpoi et al., 2014) as the source of anomaly images (denoted as A in Equation (1)).

In line with previous research, we employ AUC (Area Under the ROC Curve) to evaluate image-level and pixel-level anomaly detection. However, anomalous regions typically only occupy a tiny fraction of the entire image. Consequently, Pixel-AUROC may not accurately reflect the localization accuracy, as the false positive rate is primarily influenced by the vast number of anomalous-free pixels (Tao et al., 2022). To offer a more comprehensive measure of localization performance, we introduce the Per Region Overlap (PRO) (Bergmann et al., 2020), which assigns equal weight to anomaly regions of varying sizes. Pixel-level Average Precision (AP) (Saito and Rehmsmeier, 2015) is also introduced for comprehensive measuring. The PRO score treats anomaly regions of varying sizes equally, whereas AP is more suitable for highly imbalanced classes, particularly in the context of industrial anomaly localization where accuracy plays a critical role.

Our comparison SOTA methods include DRAEM (Zavrtanik et al., 2021a), CFLOW (Gudovskiy et al., 2022), CFA (Lee et al., 2022), PatchCore (Roth et al., 2022), RD4AD (Deng and Li, 2022), SimpleNet (Liu et al., 2023), DeSTSeg (Zhang et al., 2023c), RD++ (Tien et al., 2023), FastFlow (Yu et al., 2021), and FOD (Yao et al., 2023a).

Our initialized two-stream encoders of TLM are the ViT-Tiny (Wu et al., 2022) with MoblieSAM (Zhang et al., 2023a) weight, while SAM (Kirillov et al., 2023) serves as the teacher network. All images in the two datasets are resized to 1024 × 1024. We employ the Adam Optimizer (Kingma and Ba, 2014), setting the learning rate at 0.0005 for the TLM. For the FA module, we choose the Stochastic Gradient Descent Optimizer, using the same configuration as presented in (Zhang et al., 2023c). Each loss function in our paper carries equal weight. We conduct training for 200 epochs with a batch size of 2 and compute the average of the top 100 anomalous pixels as the image-level anomaly score. Notably, we use the augmentation method, i.e., pseudo anomaly generation strategy, that most state-of-the-art methods use for fair competition. Prototypical (Zhang et al., 2023d) presents more complex anomaly generation strategies to get better performance.

For the purpose of ensuring fair comparisons with other studies, we use the results reported in the original papers or the results re-evaluated by (Zhang et al., 2023c, d). In cases where results are not available, we indicate with a hyphen (‘-’).

Table 2 presents the results of anomaly detection and localization on the MVTec dataset. On average, our method achieves a competitive performance and secures the highest score for 9 out of 15 classes. To provide a comprehensive overview of anomaly localization capabilities, we include additional PRO metric results in Table 3. In 5 out of 15 classes, our method achieves a superior PRO score and it remains comparable to state-of-the-art in the remaining classes. The specific results in Figure 1 are shown in Table 2, our method achieves the most efficient method and the second-fastest inference speed, demonstrating the model efficiency and mobile-friendliness of our STLM.

All experiments are consistently conducted on an NVIDIA GeForce RTX 3090. Notably, we attempt to analyze the reason why our STLM has fewer parameters than DeSTSeg but longer inference times is the software and hardware optimization works for ResNet-like frameworks are more mature compared to Transformer (Li et al., 2023). The optimal inference speed of the image encoder requires further study, which we plan to optimize in future work.

To conduct more comprehensive evaluations of the anomaly detection capabilities, we further subject our networks to benchmarking using two additional widely recognized datasets VisA (Zou et al., 2022) and DAGM (Zavrtanik et al., 2021a). To provide a comprehensive overview of anomaly localization capabilities, we include 4 metrics: image-level AUROC, pixel-level AUROC, PRO and AP. The results, as presented in Table 2, indicate that our method achieves SOTA pixel-level PRO and AP metrics of 93.83/47.63 on the VisA dataset. Our method has also demonstrated its superiority on the DAGM dataset. On average, STLM performs exceptionally well, affirming its effectiveness and scalability when faced with increasingly complex datasets.

To further analyze the model performance, we visualize a qualitative evaluation of anomaly localization performance presented in Figure 7. The visual results demonstrate our model’s precise localization of anomalies. However, no one is perfect. In Figure 9, we analyze instances where our method does not perform as expected. As evidenced in the first and second columns, the susceptibility of our FA module holds the primary responsibility. But as mentioned in (Zhang et al., 2023c), we must acknowledge that several uncertain ground truths are accountable for the observed failures. Considering the results in the third and fourth columns where the ground truth highlights the designated location, our masks only cover the anomalous region on the same images.

We conducted substantial qualitative experiments on VisA datasets to visually demonstrate the superiority of our method in the accuracy of anomaly localization on other more difficult benchmarks. Also, the Feature Aggregation (FA) module is always effective. As shown in Figure 8, the visual results demonstrate our model’s robust generalization.

In Table 4, we assess the effectiveness of our four design components by conducting experiments where we eliminate the fixed teacher network and use a pre-trained mobileSAM as the teacher network to examine the contribution of SAM-knowledge, remove the plain student stream and use the fixed SAM as the plain stream in the TLM like early KD frameworks (but still use both normal and pseudo anomaly images to train denoising student stream), employee the image embedding exacted from the encoders as the FA module input and train with the cosine loss between paired outputs of two encoders, and substitute the FA module with an empirical feature fusion strategy (Wang et al., 2021). The best results are achieved when all four key design components are combined.

Notably, our work outperforms the “w/o-PlainS” manner, explaining that the plain student stream can learn generalized knowledge related to anomaly detection tasks. This enables our plain student stream to effectively represent previously unseen patterns, ensuring differences between features of two streams are well-captured.

We conduct ablation studies to investigate the impact of two distinct training strategies mentioned in Section 3.3, as detailed in Table 5. The results reveal that one-stage training delivered superior performance compared to separate training of the TLM and FA module proposed in (Zhang et al., 2023c), particularly in terms of PRO scores. This suggests that training the TLM with additional supervised signals of the pseudo anomaly mask $M$ enhances the network’s ability to locate tiny defects or accurately delineate defect boundaries.

As mentioned in Section 3.3, the input of our feature aggregation module is two element-wise products of the feature maps from TLM, as defined by Equation (2). To validate the rationality of this setting, we evaluate three alternative feature combinations as input.

The first way computes the cosine distance between the feature maps of our two-stream lightweight model, making use of more prior information from STLM, which is trained through optimization of the cosine loss function. The second way involves concatenating features of the normal image and the anomalous residual representation defined by (Zhang et al., 2023d). The anomalous feature residuals denote the element-wise Euclidean distance between tensors from two streams. The third way directly concatenates the feature maps $S_{P}^{k}$ and $S_{D}^{k}$ of TLM as the input of our FA module.

The second and the third ways preserve the information from TLM more effectively. The results are presented in Table 6, indicating our method maintains a balance between prior knowledge and representation of information.

We investigate the importance of the fixed SAM teacher, i.e., only utilizing TLM during training and inference phases, as illustrated in Figure 12. The results are shown in Table 7. We find that the knowledge distillation framework, comparing the two-stream representations, is indispensable for detecting anomalous regions.

The probability of the Pseudo Anomaly Generation process being activated governs the extent to which the training datasets deviate from the original normal datasets. To be specific, a higher probability results in a larger proportion of abnormal images and a smaller proportion of normal images. The probability of 0.5 denotes an optimal balance between pseudo-anomalous images and normal images, as shown in Figure 10, achieving the best performance. Since the “good” samples in the test dataset, the probability of 1.0 lets normal images be “unseen” images during training, leading to a high false negative. We design experiments to verify this explanation in Section 4.4.6

In Table 10, we provide further quantitative results of the false negative rate (FNR) (Altman and Bland, 1994) to verify the analysis in the ablation study that the probability of 1.0 lets normal images be “unseen” images during training, leading to a high false negative. The false negative rate, i.e., miss rate, is defined to measure the proportion of instances that are incorrectly identified as negative when it is actually positive. The FNR is calculated as:

where false negative (FN) refers to the number that is actually positive but incorrectly predicted as negative and true positive (TP) represents the total number that is actually positive.

In Equation (3) and Equation (4), we define the layers of the feature extracted from the Mask Decoder as the $k$-th layer, which is also presented in Figure 2: Element-wise Production. We conducted a series of experiments to investigate the impact of $k$, and the corresponding results are presented in Figure 11. We denoted the two layers as L1 and L2. We include image-level AUROC, pixel-level AUROC and PRO metrics to provide a comprehensive cooperation of the anomaly detection and localization capabilities of each layer.

It is evident that features extracted from L1 alone already yield excellent image-level AUROC performance, while pixel-level AUROC and PRO metrics benefit from incorporating information from both L1 and L2. It indicates that it is necessary to exchange information across different layers. Therefore, we selected the combination of L1 + L2 as our default setting.

We compare the results of employing pre-trained parameters with MobileSAM (Zhang et al., 2023a) and utilizing randomly initialized parameters for ViT-Tiny (Wu et al., 2022) in our TLM, shown in Table 8. We observe that a pre-trained image encoder allows our method to achieve optimal results in fewer epochs. At first glance, the two methods seem comparable for AD tasks. However, upon closer inspection, the ”w/o initial weight” method outperforms in texture classes, while our STLM is better suited for object classes. This leads us to analyze in future research whether SAM’s prior knowledge, centered around object segmentation, plays a role in its superior performance in object-related tasks, despite its less effective representation in texture classes.

SAM generates high-quality segmentation results for previously unseen images, making it suitable for logit distillation. As (Zhao et al., 2022) said, distillation methods often focus on distilling deep features from intermediate layers, while the significance of logit distillation is greatly overlooked. We compare the feature-based method with the logit-based method, which minimizes the focal loss (Lin et al., 2017) and L1 loss (Girshick, 2015) between segmentation masks of teacher SAM and the student streams (13). The experimental results, shown in Table 10, indicate that the performance of feature distillation is superior on average, compared with directly utilizing segmentation results from a SAM-like architecture.

To evaluate the pseudo anomaly generation strategies, we present experimental results with the renowned strategies NSA (Schlüter et al., 2022) and Cutpaste (Li et al., 2021b) in Table 11. In most SOTA methods, i.e., DRAEM (Zavrtanik et al., 2021a), Prototypical (Zhang et al., 2023d) and DeSTSeg (Zhang et al., 2023c), the authors attempt to leverage synthetic anomalies generated by NSA, which are also utilized in our paper. Furthermore, based on NSA and Cutpaste,  (Kim et al., 2024; Yao et al., 2023b; Zhang et al., 2023d; Tien et al., 2023) employ a variety of complex or artificial operations that more closely resemble the real anomalies. Notably, SimpleNet (Liu et al., 2023) simulates anomalies by introducing Gaussian noise to normal features, a process that requires careful tuning of appropriate hyperparameters.