Multiresolution Knowledge Distillation for Anomaly Detection

Introduction

Anomaly detection (AD) aims for recognizing test-time inputs looking abnormal or novel to the model according to the previously seen normal samples during training. It has been a vital demanding task in computer vision with various applications, like in industrial image-based product quality control or in health monitoring processes . These tasks also require the pixel-precise localization of the anomalous regions, called defects. This is pivotal for comprehending the dynamics of monitored procedures and triggering the apt antidotes, and providing proper data for the downstream models in industrial settings.

Traditionally, the AD problem has been approached in a one-class setting, where the anomalies represent a broadly different class from the normal samples. Recently, considering subtle anomalies has attracted attentions. This new setting further necessitates precise anomaly localization. However, performing excellently in both settings on various datasets is highly appreciated but is not yet fully achieved.

Due to the unsupervised nature of the AD problem and the restricted data access, availability of just the normal data in training, the majority of methods model the normal data abstraction by extracting semantically meaningful latent features. These methods perform well solely on either of the two mentioned cases. This problem, called the generality problem , highly declines trust in them on unseen future datasets. Moreover, anomaly localization is either impossible or poor in most of them and leads to intensive computations that hurt their real-time performance. Additionally, many earlier works suffer from unstable training, requiring unprincipled early stopping to achieve acceptable results.

Using the pre-trained networks, though not fully explored in the AD context, could potentially be an alternative track. This is especially very helpful when the sample size is small and the normal class shows large variations. Some earlier studies try to train their model based on the pre-trained features of normal data. These methods either miss anomaly localization , or tackle the problem in a region-based fashion , i.e. splitting images into smaller patches to determine the sub-regional abnormality. This is computationally expensive and often leads to inaccurate localization. To evade this issue, Bergmann et al. train an ensemble of student networks to mimic the last layer of a teacher network on the anomaly-free data. However, performing a region-based approach in this work, not only makes it heavily rely on the size of the cropped patches and hence susceptible to the changes in this size, but also intensifies the training cost severely. Furthermore, imitating only the last layer misses to fully exploit the knowledge of the teacher network . This makes them complicate their model and employ other complementary techniques, such as self-supervised learning, in parallel.

Lately, Zhang et al. have demonstrated that the activation values of the intermediate layers of neural networks are a firm perceptual representation of the input images. By this premise, we propose a novel knowledge distillation method that is designed to distill the comprehensive knowledge of an ImageNet pre-trained source network, solely on the normal training data, into a simpler cloner network. This happens by forcing the cloner’s intermediate embedding of normal training data at several critical layers to conform to those of the source. Consequently, the cloner learns the manifold of the normal data thoroughly, and yet earns no knowledge from the source about other possible input data. Hence, the cloner will behave differently from the source when fed with anomalous data. Furthermore, a simpler cloner architecture enables avoiding distraction by non-distinguishing features, and enhances the discrepancy in behavior of the two networks on anomalies.

In addition, we derive precise anomaly localization heat maps, without using region-based expensive training and testing, through exploiting the concept of gradient. We evaluate our method on a comprehensive set of datasets on various tasks of anomaly detection/localization where we exceed the SOTA in both localization and detection. Our training is highly stable and needs no dataset-dependent fine tuning. As we only train the cloner’s parameters, we require just one more forward pass of inputs through the source compared to a standard network training on the normal data. We also investigate our method through exhaustive ablation studies. Our main contributions are summarized as follows:

Enabling a more comprehensive transfer of the knowledge of the pre-trained expert network to the cloner one. Distilling the knowledge into a more compact network also helps concentrating solely on the features that are distinguishing normal vs. anomalous.

Our method has a computationally inexpensive and stable training process compared to the earlier work.

Our method allows a real-time and precise anomaly localization based on computing gradients of the discrepancy loss with respect to the input.

Conducting a huge number of diverse experiments, and outperforming previous SOTA models by a large margin on many datasets and yet staying competitive on the rest.

Related Work

Previous Methods: Autoencoder(AE)-based methods use the idea that by learning normal latent features, abnormal inputs are not reconstructed as precise as the normal ones. This results in higher reconstruction error for anomalies. To better learn these normal latent features, LSA trains an autoregressive model at its latent space and OC-GAN attempts to force abnormal inputs to be reconstructed as normal ones. These methods fail on industrial or complex datasets . SSIM-AE trains an AE with SSIM loss instead of MSE causing to perform just better on defect segmentation. Gradient-based VAE introduces an energy criterion, which is minimized at test-time by an iterative procedure. Both of the mentioned methods do not perform well on one-class settings, such as CIFAR-10 .

GAN-based approaches, like AnoGan , f-AnoGan , and GANomaly , attempt to find a specific latent space where the generator’s reconstructions, obtained from samplings of this space, are analogous to the normal data. f-AnoGan and GANomaly add an extra encoder to the generator to reduce inference time of AnoGan. Despite their acceptable performance in localization and detection on subtle anomalies, they fail on one-class settings.

Methods like uninformed-students , GT, and DSVDD keep only the useful information of normal data by building a compact latent feature space, in contrast to AE-based ones that try to miss the least amount of normal data information. To achieve this, they use self-supervised learning methods or one-class techniques. However, since we only have access to normal samples in an unsupervised setting, the optimization here is harder than in AE-based methods and usually converges to trivial solutions. To solve this issue, unprincipled early stopping is used that lowers the trust in these models on unseen future datasets. For example, GT fails on subtle anomaly datasets like MVTecAD while performs well on one-class settings.

Using Pre-trained Features: Some previous methods use pre-trained VGG’s last layer to solve the representation problem . However, sticks in bad local minima as it uses only the last layer. attempts to solve this by extracting lots of different patches from normal images. Then, it fits a Gaussian distribution on the VGG extracted embeddings of the patches. Although this might alleviate the problem, they fail to provide good localization or detection on diverse datasets because of using unimodal Guassian distribution and hand engineered size of patches.

Interpretability Methods: Determining the contribution of input elements to a deep function is investigated in interpretability methods. Gradient-based methods computes pixel’s importance using gradients as a proxy. While Gradients uses rough gradients, GuidedBackprop (GBP) filters out negative backpropagated gradients to only consider elements with positive contribution. As Gradients’ maps can be noisy, SmoothGrad adds small noises to the input and averages the maps obtained using Gradients for each noisy input. Several methods reveal some flaws in GBP by demonstrating that it reconstructs the image instead of explaining the outcome function.

Method

Given a training dataset $D_{train}=\{x_{1},...,x_{n}\}$ consisting only of normal images (i.e. no anomalies in them), we ultimately train a cloner network, $C$, that detects anomalous images in the test set, $D_{test}$, and localizes anomalies in those images with the help of a pre-trained network. As $C$ needs to predict the deviation of each sample from the manifold of normal data, it needs to know the manifold quite well. Therefore, it is trained to mimic the comprehensive behavior of an expert network, called the source network $S$. Earlier Work in knowledge distillation have conducted huge efforts to transfer one network’s knowledge to another smaller one for saving computational cost and memory usage. Many of them strive to teach just the output of $S$ to $C$. We, however, aim to transfer the intermediate knowledge of $S$ on the normal training data to $C$ as well.

In , it is shown that by using a single intermediate level hint from the source, thinner but deeper cloner even outperforms the source on classification tasks. In this work, we provide $C$ with multiple intermediate hints from $S$ by encouraging $C$ to learn $S$’s knowledge on normal samples through conforming its intermediate representations in a number of critical layers to $S$’s representations. It is known that layers of neural networks correspond to features at various abstraction levels. For instance, first layer filters act as simple edge detectors. They represent more semantic features when considering later layers. Therefore, mimicking different layers, educates $C$ in various abstraction levels, which leads to a more thorough final understanding of normal data. In contrast, using only the final layer shares a little portion of $S$’s knowledge with $C$. In addition, this causes the optimization to stuck in irrelevant local minima. On the contrary, using several intermediate hints turns the ill-posed problem into a more well-posed one. The effect of considering different layers is more investigated in Sec. 3.3.1.

In what follows, we refer to the $i$-th critical layer in the networks as $CP_{i}$ ($CP_{0}$ stands for the raw input) and the source activation values of that critical layer as $a_{s}^{CP_{i}}$, and the cloner’s ones as $a_{c}^{CP_{i}}$. As discussed in knowledge distillation literature , the notion of knowledge can be seen as the value of activation functions. We define the notion of knowledge as both the value and direction of all $a^{CP_{i}}$s to intensify the full knowledge transfer from $S$ to $C$. Hence, we define two losses, $\mathcal{L}_{val}$ and $\mathcal{L}_{dir}$ to represent each aspect. The first, $\mathcal{L}_{val}$, aims to minimize the Euclidean distance between $C$’s and $S$’s activation values at each $CP_{i}$. Thus, $\mathcal{L}_{val}$ is formulated as

where $N_{i}$ indicates the number of neurons in layer $CP_{i}$ and $a_{.}^{CP_{i}}(j)$ is the value of $j$-th activation in layer $CP_{i}$. $N_{CP}$ represents total number of critical layers.

This means that the corresponding ReLU neuron would be activated by $a_{2}$, similar to $a^{*}$, while deactivated by $a_{1}$. To address this, using the cosine similarity metric, we define the $\mathcal{L}_{dir}$ as

where $vec(x)$ is a vectorization function transforming a matrix $x$ with arbitrary dimensions into a 1-D vector. This encourages the activation vector of $C$ be not only close to the $S$’s one in terms of Euclidean distance but also be in the same direction. Note that $\mathcal{L}_{dir}$ is $1$ for $a_{1}$, and is for $a_{2}$. The role of $\mathcal{L}_{dir}$ and $\mathcal{L}_{val}$ is more elaborated in Sec. 3.3.3. Using the two aforementioned losses, $\mathcal{L}_{total}$ is formulated as

where $\lambda$ is set to make the scale of both constituent terms the same. For this, we find the initial amount of error for each term on the untrained network and set $\lambda$ with respect to it. Training using $\mathcal{L}_{total}$, unlike many other methods , continues to fully converge, which is the only accessible criterion to measure when to stop training epochs.

Moreover, the architecture of $C$ is designed to be simpler than $S$ to enable knowledge “distillation”. This compression of the network facilitates the concentration on normal main features. While the source needs to be a very deep wide model to learn all necessary features to perform well on a large-scale domain dataset, like ImageNet , the goal of the cloner is simply acquiring the source’s knowledge of the normal data. Hence, superfluous filters are only detrimental by focusing on non-distinguishing features, present in both normal and anomalous data. Compressing the source prevents such distractions for the model. This can be of a greater vitality when dealing with normal data having a more restricted scope. The effect of the cloner’s architecture is explored in Sec. 3.3.2.

Anomaly Detection: To detect anomalous samples, each test input is fed to both $S$ and $C$. As $S$ has only taught the normal point of view to $C$, anomalies, inputs out of the normal manifold, are a potential surprise for $C$. On the other hand, $S$ is knowledgeable on anomalous inputs too. All this leads to a potential discrepancy in their behavior with anomalous inputs that is thresholded for anomaly detection using Eq. 4, which formulates this discrepancy.

Anomaly Localization: have shown that the derivative of loss function with respect to the input has meaningful information about the significance of each pixel. We employ gradients of $\mathcal{L}_{total}$ to find anomalous regions causing an increase in its value. To obtain our localization map for the input $x$, we first acquire the attribution map, $\Lambda$ by

To reduce the natural noises in these maps, we induce Gaussian blur and opening morphological filter on $\Lambda$. Hence, the localization map, $L_{map}$, is achieved by

where g denotes a Gaussian filter with standard deviation of $\sigma$. $\ominus$ and $\oplus$ represent morphological erosion and dilation by a structuring element $B$, respectively. Together, called opening, these operations remove small sporadic noises and yield clean maps. The structuring element, $B$, is a simple binary map usually in shape of an ellipse or disk. Instead of using simple gradients as in Eq. 5, some other gradient-based interpretability methods can be employed to further illuminate the role of each pixel on loss value. We discuss different methods more in Sec. 3.3.4. Our proposed framework is shown schematically in Figure 2. Note that we need only two forward passes for detection and one backward pass through $C$ for localization.

2 Settings

VGG features have shown great performance in classification and transfer learning . This highlights the practicality of its filters in different domains. By transferring the knowledge of an ImageNet VGG-16 to a simple cloner, we exploit the discrepancy of features between $C$ and $S$ to find anomalies. In our VGG-16 source network, we choose the four final layers of each convolutional block, i.e. max-pooling layers, to be the critical points ($CP_{i}$s). Selecting critical points is explored more in Sec. 3.3.1.

For the cloner network, for all experiments and datasets, we use the architecture described in Figure 2, which is smaller than the source. As a result, it can benefit from the advantages of compression discussed in Sec. 3. The role of cloner architecture is discussed more in Sec. 3.3.2. Note that, similar to , we avoid using bias terms in our cloner’s network. As proven by , networks with bias in any layer can easily learn constant functions, independent of the input. In our work, though it can be negligible on datasets with diverse normal data, it can be detrimental when normal images are roughly the same. To be more specific, for some layers $l$ and $l+1$ that are between any $i$-th and $(i-1)$-th $CP$, the cloner can generate a specific constant activation vector, $a_{C}^{CP_{i}}$, regardless of the input, only by setting the $l$-th layer’s weight to zero and adjusting the $l+1$-th layers’s bias. As the normal training images are much alike, the source’s intermediate activations are also highly similar for them. Therefore, those constant $a_{c}^{CP_{i}}$s can be arbitrarily close to the source’s correlated intermediate activations for any training input, which is the goal of training phase while harming the test procedure since they are constant outputs indeed. To avoid this, we use a bias-less network for $C$.

In all experiments, we use Adam optimizer with learning rate $=0.001$ and batch size $=64$ for optimization.

3 Ablation Studies

In this experiment, we examine the effect of involving the last, the last two, and the last four max-pooling layers as $CP_{i}$s on MVTecAD and MNIST. We report average AUROC of all classes in Figure 3.3.1. Obviously, a consistent growing trend exist that shows the effectiveness of considering more layers. Notice that some MVTecAD classes (e.g “screw”) have near random AUCROC in “just the last layer“ setting. This suggests that using just the last layer makes the problem ill-posed and hard to optimize.

3.2 Distillation Effect (Compact C𝐶C)

As originally motivated in the knowledge distillation field, smaller $C$ plays an important role in our approach by eliminating non-distinguishing filters causing various distractions. It is especially more important when performing on normal data where the scope is dramatically limited. Here, we probe the effect of the cloner’s architecture. As in Figure 4, anomaly detection, on MVTecAD, using a compact $C$ network outperforms a $C$ network with equal size to $S$. This is especially noticeable on classes in which anomalies are partial (like in “toothbrush” or “screw”). Overall, the smaller network performs better with a margin of $\sim 3\%$.

In this part, we discuss each loss component’s effect to show the insufficiency of solely considering the Euclidean distance or directional loss in practice. The high impact of using $\mathcal{L}_{total}$ can be seen in Fig. 5. We report the mean AUROC over all the classes in the datasets. For more ablation studies, refer to Supplementary Materials for a a class-detailed report. Discarding the directional loss term drastically harms the overall performance on cases where anomalies are essentially different from normal cases and are more diverse, like in CIFAR-10. Using $\mathcal{L}_{dir}$ alone, however, shows top results. On the other hand, when considering cases with subtle anomalies MSE loss performs noticeably better and $\mathcal{L}_{dir}$ fails in comparison. However, in both cases, our proposed $\mathcal{L}_{total}$, which is a combination of the two losses, can achieve the highest performance. Theses results highlight the positive impact of considering a direction-wise notion of knowledge in addition to an MSE approach.

3.4 Localization using Interpretability Methods

In addition to simple gradients explained in Eq. 6, in this section, other interpretability methods are also used for anomaly localization in our framework. In Table 1, the results on MVTecAD images are shown with and without applying Guassian filter. As expected, SmoothGrad highlights the anomalous parts better than others as it discards wrongly highlighted pixels by Gradients, through calculating an average over gradients of noisy inputs. GBP, however, performs weaker than others since it tends more to reconstruct the image instead of staying faithful to the function . Anyway, after applying the noise-removing filters, the methods perform almost the same. Hence, we use simple Gradients in the rest of our experiments instead of SmoothGrad that requires severe additional computations.

Experiments

In this section, extensive experiments have been done to demonstrate the effectiveness of our method. Code to reproduce the results is provided at https://github.com/Niousha12/Knowledge_Distillation_AD. Unlike other methods that report their maximum achieved results, we report an average on our trained models, sampled every 10 epochs after convergence, to show our training stability. Variances are also reported. Finally, we emphasize that $S$ is pre-trained on ImageNet and has not seen any data of the tested datasets. Hence, the comparison is fair.

Datasets: We test our method on 7 datasets as follows: MNIST : 60k training and 10k test $28\times 28$ gray-scale handwritten digit images. Fashion-MNIST : similar to MNIST (with 10k more training images) made up of 10 fashion product categories. CIFAR-10 50k training and 10k test $32\times 32$ color images in 10 equally-sized natural entity classes. MVTecAD : an industrial dataset with over 5k high-resolution images in 15 categories of objects and textures. Each category has both normal images and anomalous images having various kinds of defects (only for testing). All images have been down scaled to the size $128\times 128$. Retinal OCT Images (optical coherence tomography) : consisting of 84,495 X-Ray images and 4 categories. HeadCT : a medical dataset containing 100 $128\times 128$ normal head CT images and 100 with hemorrhage. Each image comes from a different person. BrainMRI for brain tumor detection : consisting of 98 $256\times 256$ normal MRI images and 155 with tumors.

Evaluation Protocol: Medical datasets: 10 random normal images + all anomalous ones for test, the rest normal ones for training. MVTecAD & Retinal-OCT: datasets train and test sets are used. Others: one class as normal and others as anomaly, at testing: the whole test set is used.

2 Results

First, we evaluate our method on the conventional AD task on MNIST, Fashion-MNIST, and CIFAR-10 as described in Sec. 4.1. This targets detecting anomalies disparate from the normal samples in essence and not only slightly. As CFIAR-10 images are natural images, they have been resized and normalized according to ImageNet’s properties. No normalization and resizing is done for other datasets.

For evaluation, similar to previous works, the area under the receiver operating characteristic curve (AUROC) is used. This allows comparison using different thresholds on the anomaly score. We compare our method with an exhaustive set of state-of-the-art approaches, including generative, self-supervised and autoencoder-based methods, in Table 2. We outperform all other methods on F-MNIST and CIFAR-10, while staying comparatively well on MNIST, though avoiding complicated training procedures. Note that some methods, like U-Std, apply dataset-dependent fine-tuning. We, however, avoid such fine-tunings.

2.2 MVTecAD

Detection: In this part, we report the results of our method performance on AD using MVTecAD. As shown in Table 3, our method outperforms all others with a large margin of $\sim 10\%$. This is remarkable since other methods fail to perform well in both one-class setting and defect detection simultaneously. In contrast, we achieve SOTA in both cases.

Localization: We not only accomplish SOTA in AD but outperform previous SOTA methods in anomaly localization. As stated in 3.3.4, we use simple gradients to obtain maps. We use Gaussian filter with $\sigma=4$ and a $3\times 3$ ellipse structuring element kernel. We compare our method against others, including AE-based and generative methods in Table 4. We use AUROC, based on each pixel anomaly score, to measure how well anomalies are localized. Vividly, we outperform all previous methods. Fig. 6 shows our localization maps on different defects’ types in MVTecAD.

2.3 Medical Datasets

To further evaluate our method in various domains, we use 3 medical datasets and compare ours method on them against others. First, we use Retinal-OCT dataset, a recent dataset for detecting abnormalities in retinal optical coherence tomography (OCT) images. According to Table 5, our method outplays all SOTA methods by a huge margin. This shows that the knowledge of the pre-trained netowrk, $S$, has been highly valuable to the cloner, $C$, even in an entirely different domain of medical retinal OCT inputs. Furthermore, the unawareness of $C$ about the outside of the normal data manifold, in contrast to $S$, intensifies the discrepancy between them. This expresses the generality of our method to even future unseen datasets, something missed in many methods.

Moreover, we validate our performance on brain tumor detection using brain MRI images. In this dataset, images with tumors are assumed as anomalous while healthy ones are considered as normal. In Table 6, our method achieves SOTA results alongside LSA. While slightly ($\sim 0.5\%$) less than LSA, our method shows a significantly less variance, magnifying its stability, compared to others. It is also noteworthy that LSA fails substantially on other tasks such as on CIFAR10 and MVTecAD anomaly detection with AUROCs $\sim 23\%$ and $\sim 25\%$ below our method’s, respectively.

Lastly, using HeadCT (hemorrhage) dataset, we discuss an important aspect of our model. Performing on head computed tomography (CT) images for AD, we ouperform OCGAN and GT by a huge margin, and perform $\sim 3\%$ below LSA. Here, since the training data is dramatically limited, our method can possibly face difficulties transferring the $S$’s knowledge to $C$. However, this can be addressed by using simple data augmentations. We use 20 degree rotation in addition to scaling in range $[0.9,1.05]$ to augment the images. These augmentations are generic non-tuned ones aiming solely to increase the amount of data with no dependency to the dataset. In Table 6, it is showed that using augmentation, the proposed method achieves similar results to LSA’s, while outshining it on other tasks significantly.

Conclusion

We show that “distilling” the intermediate knowledge of an ImageNet pre-trained expert network on anomaly-free data into a more compact cloner network, and then using their different behavior with different samples, sets a new direction for finding distinctive criterion to detect and localize anomalies. Without using intensive region-based training and testing, we leverage interpretability methods in our novel framework for obtaining localization maps. We achieve superior results in various tasks and on many datasets even with domains far from ImageNet’s domain.

References

Appendix A Intermediate Knowledge

The performance of our framework using different layers as critical points for distillation was discussed in Sec. 3.3.1. Here, we provide the class-detailed performance on MVTecAD and MNIST in Table 7 and Table 8. As discussed in the paper, the performance is enhanced when more intermediate hints are considered. Note that the “only the last layer” setting performs roughly the same as a random detector (AUC=50%) on some MVTecAD classes.

Appendix B Distillation Effect (Compact Cloner)

In this section, we provide the details of results in Sec. 3.3.2 of the paper. As mentioned in the paper, using a more compact cloner network outperforms when a network with equal size to the source is employed for cloner. Here in Tables. 9 and 10 we present a class-detailed comparison for MVTecAD and CIFAR-10 datasets.

In this part, we present a classed-detailed report for the effect of each loss component as discussed in Sec. 3.3.3 in the paper. We report the AUROC for all the classes in MVTecAD and CIFAR-10 datasets in Table 9 and Table 10 respectively. As investigated in the paper, $\mathcal{L}_{total}$, which is a combination of the directional and MSE loss, achieves the best performance. Theses results highlight the positive impact of considering a direction-wise notion of activations’ knowledge in addition to an MSE approach.

Appendix D Localization using Interpretability Methods

Here, we report detailed results of Sec. 3.3.4 in the paper. In Table. 11, the AUROC for all MVTecAD classes are shown with and without applying the Guassian filter. As discussed in the paper, SmoothGrad highlights the anomalous parts better than others and GBP performs weaker than others. Anyway, after applying the noise-removing filters, the methods perform almost the same.