I TERATIVE ENERGY - BASED PROJECTION ON A NORMAL DATA MANIFOLD FOR ANOMALY LOCALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Autoencoder reconstructions are widely used for the task of unsupervised anomaly localization.
Indeed, an autoencoder trained on normal data is expected to only be able to reconstruct normal features of the data, allowing the segmentation of anomalous pixels in an image via a simple comparison between the image and its autoencoder reconstruction.
In practice however, local defects added to a normal image can deteriorate the whole reconstruction, making this segmentation challenging.
To tackle the issue we propose in this paper a new approach for projecting anomalous data on a autoencoder-learned normal data manifold, by using gradient descent on an energy derived from the autoencoder’s loss function.
This energy can be augmented with regularization terms that model priors on what constitutes the user-deﬁned optimal projection.
By iteratively updating the input of the autoencoder, we bypass the loss of high-frequency information caused by the autoencoder bottleneck.
This allows to produce images of higher quality than classic reconstructions.
Our method achieves state-of-the-art results on various anomaly localization datasets.
It also shows promising results at an inpainting task on the CelebA dataset.
I NTRODUCTION
Automating visual inspection on production lines with artiﬁcial intelligence has gained popularity and interest in recent years.
Indeed, the analysis of images to segment potential manufacturing defects seems well suited to computer vision algorithms.
However these solutions remain data hungry and require knowledge transfer from human to machine via image annotations.
Furthermore, the classiﬁcation in a limited number of user-predeﬁned categories such as nondefective, greasy, scratched and so on, will not generalize well if a previously unseen defect appears.
This is even more critical on production lines where a defective product is a rare occurrence.
For visual inspection, a better-suited task is unsupervised anomaly detection, in which the segmentation of the defect must be done only via prior knowledge of non-defective samples, constraining the issue to a two-class segmentation problem.
From a statistical point of view, an anomaly may be seen as a distribution outlier, or an observation that deviates so much from other observations as to arouse suspicion that it was generated by a different mechanism (Hawkins, 1980).
In this setting, generative models such as Variational AutoEncoders (VAE, Kingma & Welling (2013)), are especially interesting because they are capable to infer possible sampling mechanisms for a given dataset.
The original autoencoder (AE) jointly learns an encoder model, that compresses input samples into a low dimensional space, and a decoder, that decompresses the low dimensional samples into the original input space, by minimizing the distance between the input of the encoder and the output of the decoder.
The more recent variant, VAE, replaces the deterministic encoder and decoder by stochastic functions, enabling the modeling of the distribution of the dataset samples as well as the generation of new, unseen samples.
In both models, the output decompressed sample given an input is often called the reconstruction , and is used as some sort of projection of the input on the support of the normal data distribution, which we will call the normal manifold .
In most unsupervised anomaly detection methods based on VAE, models are trained on ﬂawless data and defect detection and localization is then performed using a distance metric between the input sample and its reconstruction (Bergmann et al., 2018; 2019; An & Cho, 2015;
Baur et al., 2018;
Matsubara et al., 2018).
One fundamental issue in this approach is that the models learn on the normal manifold, hence there is no guarantee of the generalization of their behavior outside this manifold.
This is problematic since it is precisely outside the dataset distribution that such methods intend to use the VAE for anomaly localization.
Even in the case of a model that always generates credible samples from the dataset distribution, there is no way to ensure that the reconstruction will be connected to the input sample in any useful way.
An example illustrating this limitation is given in ﬁgure 1, where a VAE trained on regular grid images provides a globally poor reconstruction despite a local perturbation, making the anomaly localization challenging.
In this paper, instead of using the VAE reconstruction, we propose to ﬁnd a better projection of an input sample on the normal manifold, by optimizing an energy function deﬁned by an autoencoder architecture.
Starting at the input sample, we iterate gradient descent steps on the input to converge to an optimum, simultaneously located on the data manifold and closest to the starting input.
This method allows us to add prior knowledge about the expected anomalies via regularization terms, which is not possible with the raw VAE reconstruction.
We show that such an optimum is better than previously proposed autoencoder reconstructions to localize anomalies on a variety of unsupervised anomaly localization datasets (Bergmann et al., 2019) and present its inpainting capabilities on the CelebA dataset (Liu et al., 2015).
We also propose a variant of the standard gradient descent that uses the pixel-wise reconstruction error to speed up the convergence of the energy.
[Figure]
B ACKGROUND 2. G ENERATIVE MODELS
In unsupervised anomaly detection, the only data available during training are samples x from a non-anomalous dataset X ⊂ R d .
In a generative setting, we suppose the existence of a probability function of density q , having its support on all R d , from which the dataset was sampled.
The generative objective is then to model an estimate of density q , from which we can obtain new samples close to the dataset.
Popular generative architectures are Generative Adversarial
Networks (GAN, Goodfellow et al.
, that concurrently train a generator G to generate samples from random, low-dimensional noise z ∼ p , z ∈ R l , l (cid:28) d , and a discriminator D to classify generated samples and dataset samples.
This model converges to the equilibrium of the expectation over both real and generated datasets of the binary cross entropy loss of the classiﬁer min G max D [ E x ∼ q [log( D ( x ))]
+ E z ∼ p
[log(1 − D ( G ( z )))] ] .
Disadvantages of GANs are that they are notoriously difﬁcult to train (Goodfellow, 2016), and they suffer from mode collapse, meaning that they have the tendency to only generate a subset of the original dataset.
This can be problematic for anomaly detection, in which we do not want some subset of the normal data to be considered as anomalous (Bergmann et al., 2019).
Recent works such as Thanh-Tung et al.
(2019) offer simple and attractive explanations for
GAN behavior and propose substantial upgrades, however Ravuri & Vinyals (2019) still support the point that GANs have more trouble than other generative models to cover the whole distribution support.
Another generative model is the VAE (Kingma & Welling (2013)), where, similar to a GAN generator, a decoder model tries to approximate the dataset distribution with a simple latent variables prior p ( z ) , with z ∈ R l , and conditional distributions output by the decoder p ( x | z ) .
This leads to the estimate p ( x )
=
(cid:82) p ( x | z ) p ( z ) dz , that we would like to optimize using maximum likelihood estimation on the dataset.
To render the learning tractable with a stochastic gradient descent (SGD) estimator with reasonable variance, we use importance sampling, introducing density functions q ( z | x ) output by an encoder network, and Jensen’s inequality to get the variational lower bound :
[Equation]
We will use L ( x ) as our loss function for training.
We deﬁne the VAE reconstruction, per analogy with an autoencoder reconstruction, as the deterministic sample f V AE ( x ) that we obtain by encoding x , decoding the mean of the encoded distribution q ( z | x ) , and taking again the mean of the decoded distribution p ( x | z ) .
VAEs are known to produce blurry reconstructions and generations, but Dai & Wipf (2019) show that a huge enhancement in image quality can be gained by learning the variance of the decoded distribution p ( x | z ) .
This comes at the cost of the distribution of latent variables produced by the encoder q ( z ) being farther away from the prior p ( z ) , so that samples generated by sampling z ∼ p ( z ) , x ∼ p ( x | z ) have poorer quality.
The authors show that using a second
VAE learned on samples from q ( z ) , and sampling from it with ancestral sampling u ∼ p ( u ) , z ∼ p ( z | u ) , x ∼ p ( x | z ) , allows to recover samples of GAN-like quality.
The original autoencoder can be roughly considered as a VAE whose encoded and decoded distributions have inﬁnitely small variances.
2. A NOMALY DETECTION AND LOCALIZATION
We will consider that an anomaly is a sample with low probability under our estimation of the dataset distribution.
The
VAE loss, being a lower bound on the density, is a good proxy to classify samples between the anomalous and nonanomalous categories.
To this effect, a threshold T can be deﬁned on the loss function, delimiting anomalous samples with L ( x ) ≥ T and
normal samples L ( x ) < T .
However, according to Matsubara et al.
(2018), the regularization term L KL ( x ) =
D KL ( q ( z | x ) (cid:107) p ( z )) has a negative inﬂuence in the computation of anomaly scores.
They propose instead an unregularized score L r ( x ) =
− E z ∼ q ( z | x ) log p ( x | z ) which is equivalent to the reconstruction term of a standard autoencoder and claim a better anomaly detection.
Going from anomaly detection to anomaly localization, this reconstruction term becomes crucial to most of existing solutions.
Indeed, the inability of the model to reconstruct a given part of an image is used as a way to segment the anomaly, using a pixel-wise threshold on the reconstruction error.
Actually, this segmentation is very often given by a pixel-wise (An & Cho, 2015;
Baur et al., 2018; Matsubara et al., 2018) or patch-wise comparison of the input image, and some generated image, as in Bergmann et al.
(2018; 2019), where the structural dissimilarity (DSSIM, Wang et al.
between the input and its VAE reconstruction is used.
Autoencoder-based methods thus provide a straightforward way of generating an image conditioned on the input image.
In the GAN original framework, though, images are generated from random noise z ∼ p ( z ) and are not conditioned by an input.
Schlegl et al.
(2017) propose with AnoGAN to get the closest generated image to the input using gradient descent on z for an energy deﬁned by:
[Equation]
The ﬁrst term ensures that the generation G ( z ) is close to the input x .
The second term is based on a distance between features of the input and the generated images, where f D ( x ) is the output of an intermediate layer of the discriminator.
This term ensures that the generated image stays in the vicinity of the original dataset distribution.
P ROPOSED METHOD
3. A DVERSARIAL PROJECTIONS
According to Zimmerer et al.
(2018), the loss gradient with respect to x gives the direction towards normal data samples, and its magnitude could indicate how abnormal a sample is.
In their work on anomaly identiﬁcation, they use the loss gradient as an anomaly score.
Here we propose to use the gradient of the loss to iteratively improve the observed x .
We propose to link this method to the methodology of computing adversarial samples in Szegedy et al.
After training a VAE on non-anomalous data, we can deﬁne a threshold T on the reconstruction loss L r as in (Matsubara et al., 2018), such that a small proportion of the most improbable samples are identiﬁed as anomalies.
We obtain a binary classiﬁer deﬁned by
[Equation]
Our method consists in computing adversarial samples of this classiﬁer (Szegedy et al., 2013), that is to say, starting from a sample x 0 with A ( x 0 )
= 1 , iterate gradient descent steps over the input x , constructing samples x 1 , . . .
x N , to minimize the energy E ( x ) , deﬁned as
[Equation]
An iteration is done by calculating x t +1 as
[Equation]
where α is a learning rate parameter, and λ is a parameter trading off the inclusion of x t in the normal manifold, given by L r ( x t ) , and the proximity between x t and the input x 0 , assured by the regularization term
|| x t − x 0 || 1 .
3. R EGULARIZATION TERM
We model the anomalous images that we encounter as normal images in which a region or several regions of pixels are altered but the rest of the pixels are left untouched.
To recover the best segmentation of the anomalous pixels from an anomalous image x a , we want to recover the closest image from the normal manifold x g .
The term closest has to be understood in the sense that the smallest number of pixels are modiﬁed between x a and x g .
In our model, we therefore would like to use the L 0 distance as a regularization distance of the energy.
Since the L 0 distance is not differentiable, we use the L 1 distance as an approximation.
3. O PTIMIZATION IN INPUT SPACE
While in our method the optimization is done in the input space, in the previously mentioned AnoGAN, the search for the optimal reconstruction is done by iterating over z samples with the energy deﬁned in equation 2.
Following the aforementioned analogy between a GAN generator G and a VAE decoder Dec , a similar approach in the context of a
VAE would be to use the energy
[Equation]
where the − log p ( z ) term has the same role as AnoGAN’s ||
f D ( x ) − f D ( G ( z ))
|| 1 term, to ensure that Dec ( z ) stays within the learned manifold.
We chose not to iterate over z in the latent space for two reasons.
First, because as noted in Dai & Wipf (2019) and Hoffman & Johnson (2016), the prior p ( z ) is not always a good proxy for the real image of the distribution in the latent space q ( z ) .
Second, because the VAE tends to ignore some details of the original image in its reconstruction, considering that these details are part of the independent pixel noise allowed by the modeling of p ( x | z ) as a diagonal Gaussian, which causes its infamous blurriness.
An optimization in latent space would have to recreate the high frequency structure of the image, whereas iterating over the input image space, and starting the descent on the input image x 0 , allows us to keep that structure and thus to obtain projections of higher quality.
3. O PTIMIZING GRADIENT DESCENT
We observed that using the Adam optimizer (Kingma & Ba, 2014) is beneﬁcial for the quality of the optimization.
Moreover, to speed up the convergence and further preserve the aforementioned high frequency structure of the input, [Table] cable 0. capsule 0.
hazelnut 0.
metalnut 0.
pill 0.
screw 0.
toothbrush 0.
transistor 0.
zipper 0.
we propose to compute our iterative samples using the pixel-wise reconstruction error of the VAE.
To explain the intuition behind this improvement, we will consider the inpainting task.
In this setting, as in anomaly localization, a local perturbation is added on top of a normal image.
However, in the classic inpainting task, the localization of the perturbation is known beforehand, and we can use the localization mask Ω to only change the value of the normal pixels in the gradient descent:
[Equation]
where (cid:12) is the Hadamard product.
For anomaly localization and blind inpainting, where this information is not available, we compute the pixel-wise reconstruction error which gives a rough estimate of the mask.
The term ∇ x E ( x t ) is therefore replaced with ∇ x E ( x t ) (cid:12) ( x t − f V AE ( x t )) 2 ) in equation 5:
[Equation]
where f V AE ( x ) is the standard reconstruction of the VAE.
Optimizing the energy this way, a pixel where the reconstruction error is high will update faster, whereas a pixel with good reconstruction will not change easily.
This prevents the image to update its pixels where the reconstruction is already good, even with a high learning rate.
As can be seen in appendix B, this method converges to the same performance as the method of equation 5, but with fewer iterations.
An illustration of our method can be found in ﬁgure 2.
3. S TOP CRITERION
A standard stop criterion based on the convergence of the energy can efﬁciently be used.
Using the adversarial setting introduced in section 3.1, we also propose to stop the gradient descent when a certain predeﬁned threshold on the
VAE loss is reached.
For example, such a threshold can be chosen to be a quantile of the empirical loss distribution computed on the training set.
E XPERIMENTS
In this section, we evaluate the proposed method for two different applications: anomaly segmentation and image inpainting.
Both applications are interesting use cases of our method, where we search to reconstruct partially corrupted images, correcting the anomalies while preserving the uncorrupted image regions.
4. U NSUPERVISED ANOMALY SEGMENTATION
In order to evaluate the proposed method for the task of anomaly segmentation, we perform experiments with the recently proposed MVTec dataset (Bergmann et al., 2019).
This collection of datasets consists of 15 different categories of objects and textures in the context of industrial inspection, each category containing a number of normal and anomalous samples.
We train our model on normal training samples and test it on both normal and anomalous test samples to evaluate the anomaly segmentation performance.
We perform experiments with three different baseline autoencoders: A “vanilla” variational autoencoder with decoder covariance matrix ﬁxed to identity (Kingma & Welling, 2013), a variational autoencoder with learned decoder variance (Dai & Wipf, 2019), a “vanilla” deterministic autoencoder trained with L 2 as reconstruction loss ( L 2 AE) and a deterministic autoencoder trained with DSSIM reconstruction loss (DSAE), as proposed by Bergmann et al.
For the sake of a fair comparison, all the autoencoder models are parameterized by convolutional neural networks with the same architecture, latent space dimensionality (set to 100), learning rate (set to 0.0001) and number of epochs (set to 300).
The architecture details (layers, paddings, strides) are the same as described in Bergmann et al.
(2018) and Bergmann et al.
Similarly to the authors in Bergmann et al.
(2019), for the textures datasets, we ﬁrst subsample the original dataset images to 512 × 512 and then crop random patches of size 128 × 128 which are used to train and test the different models.
For the object datasets, we directly subsample the original dataset images to 128 × 128 unlike in Bergmann et al.
(2019) who work on 256 × 256 images, then we perform rotation and translation data augmentations.
For all datasets we train on 10000 images.
Anomaly segmentation is then computed by reconstructing the anomalous image and comparing it with the original.
We perform the comparison between reconstructed and original with the DSSIM metric as it has been observed in
Bergmann et al.
(2018) that it provides better anomaly localization than L 2 or L 1 distances.
For the gradient descent, we set the step size α := 0 .
5 , L 1 regularization weight λ := 0 .
and the stop criterion is achieved when a sample reconstruction loss is inferior to the minimum reconstruction loss over the training set.
In table 1 we show the AUROC (Area Under the Receiver Operating Characteristics) for different autoencoder methods, with different thresholds applied to the DSSIM anomaly map computed between original and reconstructed images.
Note that an AUROC of 1 expresses the best possible segmentation in terms of normal and anomalous pixels.
For each autoencoder variant we compare the baseline reconstruction with the proposed gradient-based reconstruction (grad.).
As in Bergmann et al.
we observe that an overall best model is hard to identify, however we show that our method increases the AUC values for almost all autoencoder variants. Aggregating the results over all datasets and baselines, we report a mean improvement rate of 9.52%, with a median of 4.33%, a 25th percentile of 1.86%, and a 75th percentile of 15.86%. The histogram of the improvement rate for all datasets and baselines is provided in appendix F, as well as a short analysis.
In ﬁgure 3 we compare our anomaly segmentation with a baseline L 2 autoencoder Bergmann et al.
( L 2 AE) for a number of image categories.
For all results in ﬁgure 3, we set the same threshold to 0.2 to the anomaly detection map given by the DSSIM metric.
The visual results in ﬁgure 3 highlights an overall improvement of anomaly localization by our proposed iterative reconstruction ( L 2 AE-grad).
See appendix C for additional visual results of anomaly segmentation on remaining categories of MVTec dataset, and on remaining baseline models.
4. I NPAINTING
Image inpainting is a well known image reconstruction problem which consists of reconstructing a corrupted or missing part of an image, where the region to be reconstructed is usually given by a known mask.
Many different approaches for inpainting have been proposed in the literature, such as anisotropic diffusion (Bertalmio et al., 2000), patch matching (Criminisi et al., 2004), context autoencoders (Pathak et al., 2016) and conditional variational autoencoders (Ivanov et al., 2019).
N o r m a l
A no m a l ou s
L 2 A
E
L 2 A E - g r a d
If we consider that the region to be reconstructed is not known beforehand, the problem is sometimes called blind inpainting (Altinel et al., 2018), and the corrupted part can be seen as an anomaly to be corrected.
We performed experiments with image inpainting on the CelebA dataset (Liu et al., 2015), which consists of celebrity faces.
In ﬁgure 4 we compare the inpainting results obtained with a baseline VAE with learned variance ( γ -VAE) and Resnet architecture, as described by Dai & Wipf (2019), with the same VAE model, augmented by our proposed gradient-based iterative reconstruction.
Note that for the regular inpainting task, gradients are multiplied by the inpainting mask at each iteration (equation 7), while for the blind inpainting task, the mask is unknown.
See appendix D for a comparison with a recent method based on variational autoencoders, proposed by Ivanov et al.
R ELATED W ORK
Baur et al.
(2018) have used autoencoder reconstructions to localize anomalies in MRI scans, and have compared several variants using diverse per-pixel distances as well as perceptual metrics derived from a GAN-like architecture.
Bergmann et al.
(2018) use the structural similarity metric (Wang et al., 2004) to compare the original image and its reconstruction to achieve better anomaly localization, and also presents the SSIM autoencoder, which is trained directly with this metric.
Zimmerer et al.
(2018) use the derivative of the VAE loss function with respect to the input, called the score .
The amplitude of the score is supposed to indicate how abnormal a pixel is.
While we agree that the gradient of the loss is an indication of an anomaly, we think that we have to integrate this gradient over the path from the input to the normal [Figure] manifold to obtain meaningful information.
We compare our results to score-based results for anomaly localization in appendix A.
The work that is the most related to ours is AnoGAN (Schlegl et al., 2017).
We have mentioned above the differences between the two approaches, which, apart from the change in underlying architectures, boil down to the ability in our method to update directly the input image instead of searching for the optimal latent code.
This enables the method to converge faster and above all to keep higher-frequency structures of the input, which would have been deteriorated if it were passed through the AE bottleneck.
Bergmann et al.
(2019) compare standard AE reconstructions techniques to AnoGAN, and observes that AnoGAN’s performances on anomaly localizations tasks are poorer than AE’s due to the mode collapse tendency of GAN architectures.
Interestingly, updates on AnoGAN such as fast AnoGAN (Schlegl et al., 2019) or AnoVAEGAN (Baur et al., 2018) replaced the gradient descent search of the optimal z with a learned encoder model, yielding an approach very similar to the standard VAE reconstruction-based approaches, but with a reconstruction loss learned by a discriminator, which is still prone to mode collapse (Thanh-Tung et al., 2019).
C ONCLUSION
In this paper, we proposed a novel method for unsupervised anomaly localization, using gradient descent of an energy deﬁned by an autoencoder reconstruction loss.
Starting from a sample under test, we iteratively update this sample to reduce its autoencoder reconstruction error.
This method offers a way to incorporate human priors into what is the optimal projection of an out-of-distribution sample into the normal data manifold.
In particular, we use the pixel-wise reconstruction error to modulate the gradient descent, which gives impressive anomaly localization results in only a few iterations.
Using gradient descent in the input data space, starting from the input sample, enables us to overcome the autoencoder tendency to provide blurry reconstructions and keep normal high frequency structures.
This signiﬁcantly reduces the number of pixels that could be wrongly classiﬁed as defects when the autoencoder fails to reconstruct high frequencies.
We showed that this method, which can easily be added to any previously trained autoencoder architecture, gives state-of-the-art results on a variety of unsupervised anomaly localization datasets, as well as qualitative reconstructions on an inpainting task.
Future work can focus on replacing the L 1 -based regularization term with a Bayesian prior modeling common types of anomalies, and on further improving the speed of the gradient descent.
A C OMPARISON WITH Z IMMERER ET AL . (2019)
[Table]
hazelnut
[Table]
zipper
[Table]
Zimmerer et al.
(2019) proposed to perform anomaly localization using different scores derived from the gradient of the VAE loss.
In particular, it has been shown that the product of the VAE reconstruction error with the gradient of the KL divergence was very informative for medical images.
In table 2 we compare the pixel-wise anomaly detection
AUROC of these different scores with our method.
For all experiments, we use the same “vanilla” VAE as described in section 4.1.
It can be seen that other VAE-based methods using a single evaluation of the gradient are constantly outperformed by our method.
In ﬁgure 5 we compare the number of iterations needed to reach convergence with our two proposals for gradient descent: Standard update as in equation 5 and Tuned update using a gradient mask computed with the VAE reconstruction error, as in equation 8.
The model is a VAE with learned decoder variance (Dai & Wipf, 2019), trained on the
Grid dataset (Bergmann et al., 2019).
We compute the mean pixel-wise anomaly detection AUROC after each iteration on the test set.
We can see that the tuned method converges to the same performance as the standard method, with far fewer iterations. [Figure] L 2 A E - g r a d
D S A
E D S A E - g r a [Figure]
g r [Figure] Figure 9 illustrates our method principle. We start with a defective input x 0 whose reconstruction ˆ x 0 does not necessarily lie on the normal data manifold.
As the optimization process carries on, the optimized sample x 0 and its reconstruction look more similar and get closer to the manifold.
The regularization term of the energy function makes sure that the optimized sample stays close to the original sample.
[Figure]
• 8.3% of data points are under the 0 value delimiting an increase or decrease in AUC due to our method, and 91.7% data points are over this value.
Our method increases the AUC in a vast majority of cases.
• The median is at 4.33%, the 25th percentile at 1.86%, and the 75th percentile at 15.86%.