Unsupervised Anomaly Detection via Masked Diffusion Posterior Sampling

1 Introduction

Anomaly detection (AD) is a fundamental computer vision task with widespread applications in medical diagnosis, intelligent manufacturing, autonomous driving and etc. However, due to the rarity and diversity of anomaly samples, recent studies mainly focus on unsupervised anomaly detection (UAD) [Diers and Pigorsch, 2023], in which the models only learn from normal samples but can detect anomaly data.

So far, there have existed various methods for UAD, among which reconstruction-based method is one of the earliest and most common neural network approaches [Ruff et al., 2021]. Given a test image, reconstruction-based method tries to reconstruct the corresponding normal image and compute the difference between the test image and the reconstruction to detect and localize the anomalies. Obviously, how to reconstruct a normal image is a key issue for reconstruction-based methods. Early reconstruction models include Autoencoder (AE), Generative Adversarial Network (GAN) and etc. However, AE can easily suffer from "identical shortcut" [You et al., 2022a] and blurry reconstruction [Baur et al., 2021], i.e., they reconstruct both normal and anomalous images and the results tend to be more blurry. Although GAN can alleviate the problems of AE, mode collapse and training instability make GAN challenging for UAD [Xia et al., 2022].

∗B. Lin is the corresponding author. Source code will be available at https://github.com/KevinBHLin/.

Recently, Diffusion Models (DMs) have attracted most researchers' attention with their powerful image generation ability [Ho et al., 2020]. Compared with previous generative models, DMs can effectively record image priors and generate various realistic images after simple training based on Gaussian denoising. Thus, they show promising applications for UAD to alleviate the problems of inferior reconstruction quality or insufficient coverage of the normal image distribution. However, two key problems arise when introducing DMs for UAD: First, although different DM-based methods are proposed for UAD, they are lack of strict mathematical theories or interpretability to ensure that the anomaly region of a test image can be reconstructed as the normal one. Second, most DMs unexpectedly suffer from low reconstruction quality, especially for the normal region of a test image, since Gaussian noise in DMs will destroy the original normal texture. This problem can easily result in misjudgment of anomalous pixels in the normal region. Thus, it requires further studies for DMs to maintain the texture in the normal region of a test image after reconstruction.

In this paper, we propose a novel Masked Diffusion Posterior Sampling (MDPS) method for UAD under Bayesian framework, which has high interpretability supported by relatively-strict mathematics. In our method, we firstly propose a Masked Noisy Observation Model which regards a test image as a masked noisy observation of a normal image to protect the normal region of a test image and enhance reconstruction quality. Then, we take a Denoising Diffusion Implicit Model (DDIM) trained by normal samples as image prior, and model the problem of normal image reconstruction as multiple diffusion posterior samplings for normal images based on the devised observation model and the DDIMbased prior. Third, using a metric designed from pixel-level and perceptual-level perspectives, we compute the difference maps between multiple normal posterior samples and the test image respectively, and average all difference maps to accurately obtain anomaly scores. Exhaustive experiments on MVTec and BTAD datasets prove that the proposed MDPS achieves excellent reconstruction quality and high accuracy of anomaly detection and localization compared with other state-of-the-art reconstruction-based methods.

2 RelatedWork

Recent reconstruction-based UAD methods can be roughly divided into three categories, including AE-based methods, GAN-based methods, and DM-based methods.

AE-based methods. Early reconstruction-based UAD methods commonly adopt AE due to their simple architectures and easy-to-implement training processes. Unfortunately, AEs can easily suffer from the problems of "identical shortcut" and blurry reconstruction. For the first problem, since Vision Transformer (ViT) can prevent "identical shortcut" [You et al., 2022a], recent studies try to design ViT-based AEs, e.g., [Mishra et al., 2021], [You et al., 2022b], [You et al., 2022a] and etc. For the second problem, various methods have been proposed: [Bergmann. et al., 2019] designs a loss function of structure similarity to improve reconstruction quality; [Liu et al., 2020] and [Dehaene et al., 2020] introduce Variational AE to achieve better reconstruction results; [Zavrtanik et al., 2021] and [Sun et al., 2023] propose a self-supervised model to reduce the influence of blurry reconstruction. Nevertheless, the above schemes cannot totally avoid blurry reconstruction, thus, leading to performance bottleneck for UAD.

GAN-basedmethods. To overcome the drawbacks of AEs, [Schlegl et al., 2017] firstly proposed a GAN-based method named AnoGAN. Since then, various variants are proposed, e.g., [Akc¸ay et al., 2019], [Schlegl et al., 2019] and etc. Although GANs can empirically generate high definition results, two problems make GANs challenging for UAD [Xia et al., 2022]: First, GAN training is highly unstable to converge, sometimes bringing in meaningless reconstruction. Since the training process of AE is more stable, AE and GAN are usually combined to alleviate training instability, e.g., [Tang et al., 2020], [Contreras-Cruz et al., 2023]. Second, due to the vanish of discriminator gradient, it is hard to ensure sufficient coverage of the normal distribution during GAN training process. This phenomenon is called mode collapse, which can easily lead to reconstruction results with only a few modes. To alleviate this problem, self-supervised models are proposed based on simulated anomaly data to guide GAN training, e.g., [Song et al., 2022], [Liu et al., 2023] and etc.

DM-basedmethods. Recently, DMs have been popular for academic research due to their potentially-powerful generative ability. Typical DMs include DDPM [Ho et al., 2020], DDIM [Song et al., 2021a], SDE [Song et al., 2021b], etc. Compared with AEs and GANs, DMs have desirable properties, such as distribution coverage, a stationary training objective and easy scalability [Dhariwal and Nichol, 2021]. To generate desirable images based on DMs for certain requirements, conditional DMs are studied, e.g., [Dhariwal and Nichol, 2021] proposes a classifier-guided DM to generate images given a certain class label; [Chung et al., 2023] extends diffusion solvers to efficiently handle general noisy inverse problems via approximation of posterior sampling. The above works show promising applications for UAD to reconstruct high definition normal images.

So far, there have been a few DM-based AD methods. [Wolleb et al., 2022] and [Pinaya et al., 2022] firstly introduce DMs for weakly supervised AD and UAD respectively in medical diagnosis. [Wyatt et al., 2022] proposes a DM-based UAD method named AnoDDPM, which uses simplex noise to improve normal image reconstruction quality for brain MRI. [Gonzalez-Jimenez et al., 2023] uses score-based models to produce a gradient map highlighting anomaly areas. [Bercea et al., 2023] proposes a DM-based model named AutoDDPM which consists of three stages, i.e, mask, stitch, and re-sampling. [Lu et al., 2023] models the problem of normal image reconstruction as a DM-based denoising process. The above methods have enhanced the accuracy of DM-based AD from different aspects. However, these DM-based methods are lack of strictly-mathematical support, and unexpectedly suffer from low normal image reconstruction quality, especially for the normal region of a test image.

3 Background

In this section, we briefly review SDE-based diffusion models [Song et al., 2021b] and DDIM [Song et al., 2021a].

SDE-based diffusion models. For a diffusion process {x(t)}Tt=0 indexed by a continues time variable t ∈ [0, T], let x(0) ∼ p0(x) and x(T) ∼ pT (x), in which p0(x) and pT (x) denotes the data distribution of interest and a known spherical Gaussian distribution respectively. Then, the forward noising process x(0) → x(T) can be modeled as the following Itˆo stochastic differential equation (SDE):

where f(·, t) and g(·) are the drift and diffusion coefficient of x(t) respectively, ω is a standard Wiener process. The corresponding reverse process x(T ) → x(0) is given by

where ˆω is a standard Wiener process running backwards, and dt is an infinitesimal negative time step. ∇x log pt(x) is the score function of each marginal distribution which can be approximately estimated by training a time-dependent scorebased model sθ(x, t), i.e., sθ(x, t) ≃ ∇x log pt(x).

DDIMs. DDIMs can be regarded as an acceleration version of DDPMs [Ho et al., 2020] with the same training procedure. However, different fromDDPMs, the forward noising process of DDIM is defined as a non-Markov process:

where {αt ∈ (0, 1]|t ∈ [0, T]} is a decreasing sequence decided by a predetermined schedule, and ϵ is white noise drawn from a standard normal distribution. The accelerated reverse process of DDIM is described as:

where σt = √(1-αs)/(1-αt)√1-αt/αs, s < t, and ϵθ(xt, t) is a U-net denoiser which estimates the white noise ϵ from xt, and can be trained by minimizing the following objective:

[Song et al., 2021b] has pointed out that DDIM is a discrete form of SDE when f(x, t) = -xtβ(t)/2, g(t) = √β(t), and αt = ∏t (1 - β(t)) in Eqn.(1). Thus,

4 Masked Diffusion Posterior Sampling

In this section, we firstly establish a masked noisy observation model for normal image reconstruction. Then, based on the observation model, we propose a posterior sampling method for normal images using DDIMs. Thirdly, we design imagewise and pixel-wise anomaly scores and propose a mask generation scheme based on the designed anomaly scoring for the proposed observation model to enhance the reconstruction quality of posterior samples.

4.1 Masked Noisy ObservationModel

Let y and x0 denote an anomaly image and the corresponding normal image respectively. Let m denote a mask image with the same size of y in which pixel values are set to 0 in the normal regions of y and set to 1 otherwise. Then, we establish a masked noisy observation model to describe the relationship between y and x0 as the following:

where ⊙ denotes Hadamard product, n denotes zero-mean isotropic Gaussian noise and n ∼ N(0, σ2I). Obviously, the above modeling assumes that y is a noisy observation of x0 in the anomaly region whereas y shares the same pixel values as that of x0 in the normal region, i.e.,

As y would contain unknown types of anomaly, the assumption of Gaussian noise for anomaly is reasonable. Furthermore, the mask 1-m can guarantee the reconstruction quality of normal image x0, since it maintains pixel values in the normal region of y for reconstruction of x0. Then, based on Eqn.(7), we can obtain the distribution of y given x0 for the anomaly region indicated by m as the following:

4.2 Posterior Sampling for Normal Images

To obtain multiple normal samples of x0 given y, we propose to conduct posterior sampling for p(x0|y). According to Bayes' Theorem, p(x0|y) is related to the normal image prior p(x0) and the observation distribution p(y|x0) given in Eqn.(9). In our model, we introduce a DDIM trained by normal samples to model the normal image prior p(x0), since this model has been widely adopted for its powerful image generation ability in recent years. Then, following [Chung et al., 2023], we regard the posterior sampling process as a diffusion process based on the introduced DDIM for p(x0), and obtain multiple normal samples drawn from p(x0|y). The details are shown as the following.

Specifically, let ϵθ(xt, t) and ϵφ(xt, y, t) denote the denoisers of two DDIMs for sampling of p(x0) and p(x0|y) respectively. Since ϵθ(xt, t) is known and trained by normal images, we discuss the design of ϵφ(xt, y, t). According to Eqn.(3) and Eqn.(8a), we can accurately estimate ϵ as:

where ϵmφ (xt, y, t) denotes the values of ϵφ(xt, y, t) in the normal region indicated by 1 -m.

However, for the anomaly region indicated by m, according to Bayes' Theorem, Eqn.(6) and Eqn.(8b),

where ϵmφ (xt, y, t) denotes the result of ϵφ(xt, y, t) for the anomaly region indicated by m. Obviously, ϵmφ (xt, y, t) can be replaced by the above equation base on ϵθ(xt, t) and ∇xt log p(y|xt). Since ϵθ(xt, t) is a known denoiser trained by normal samples, we focus on the modeling of p(y|xt).

As has been proven in [Chung et al., 2023],

where p(y|x0) has been given in Eqn.(9), and p(x0|xt) is determined by the reverse process of DDIMs for sampling of p(x0). Based on Eqn.(3) and Eqn.(4),

Apparently, in the reverse process of DDIMs, ϵ of Eqn.(3) turns to be an unknown fixed constant, and can be estimated by ϵθ(xt, t). The estimation errors between ϵ and ϵθ(xt, t) are reduced by minimizing the objective function of Eqn.(5) during the training process of DDIMs for sampling of p(x0). Thus, we can assume that ϵ ∼ N(ϵθ(xt, t), λ-1 t I) in which λt indicates the estimation precision, and obtain

Then, for the anomaly region indicated bym, p(y|xt) can be calculated based on Eqn.(9), Eqn.(12) and Eqn.(14) as:

For simplicity, we let ρt = ρ, ρ is an adjustable guidance scale. ϵmφ (xt, y, t) can be obtained based on Eqn.(15) as:

Combining Eqn.(16) and Eqn.(10), we can obtain the denoiser ϵφ(xt, y, t) of DDIMs for sampling of p(x0|y) as:

4.3 Anomaly Scoring andMask Generation

To detect and locate anomaly, we have to predict image-wise and pixel-wise anomaly scores respectively. For prediction of pixel-wise anomaly score, we have to calculate the difference between the test image and the corresponding normal reconstruction result. Traditionally, pixel-level metrics such as L1norm, and MSE are adopted to measure the differences of two

Algorithm 1 Masked Diffusion Posterior Sampling

Input: test image y, DDIM denoiser ϵθ(xt, t), mask image m, guidance scale ρ, noise level T , sampling times N Output: normal image x0

images. Recently, [Zhang et al., 2018] proposed a perceptuallevel metric named LPIPS based on feature maps extracted by a pretrained convolutional neural network. Since LPIPS can better match human perceptual similarity judgments, it has been widely adopted for reconstruction-based AD method, e.g., [Defard et al., 2021], [Roth et al., 2022] and etc. Following recent works, the difference between y and x0 at the k-th spatial position is calculated as:

in which L1-norm and LPIPS are incorporated to measure the differences of two images from pixel-level and perceptuallevel perspectives, F (k) i denotes the i-th stage output feature of ResNet-101 proposed by [Zagoruyko and Komodakis, 2016] pretrained on ImageNet at the k-th spatial position, and η is a hyperparameter to balance the weight between the pixel-level and perceptual-level metrics. J denotes the set of different stages of ResNet. In our paper, J is set to {1, 2, 3}. Since our model will generate multiple normal samples given a single test image, the pixel-wise anomaly score map for y is defined based on Eqn.(18) as:

where x(j) 0 denotes the j-th reconstructed normal sample, and Ns denotes the number of normal samples.

Then, based on Eqn.(19), we further design the image-wise anomaly score for y as the average of the largest S pixelwise anomaly scores in D to mitigate false positives caused by image noise. In our paper, S is set to 500.

Using the above image-wise and pixel-wise anomaly scores, we design a mask image generation scheme for our observation model of Eqn.(7). To obtain an accurate mask m, we firstly set m to 1, i.e., all pixels in the test image y are regarded as potential anomaly pixels, and run a group of posterior sampling for normal images. Then, we obtain the score map D for y based on Eqn.(19), and estimate m as:

where λ is a hyperparameter, and Tth is a threshold for anomaly score determined by λ. Using the above mask generation scheme, we re-run another group of posterior sampling for normal images and calculate anomaly scores for each image and each pixel respectively. The proposed MDPS is summarized in Algorithm 1.

5 Experiments

In this section, we compare our MDPS with other UAD methods, and conduct ablation studies to validate the designs.

Datasets. We conduct all experiments on the MVTec and BTAD Datasets. The MVTec dataset is an industrial AD benchmark [Bergmann et al., 2021], which contains 15 categories (5 textural categories and 10 object categories) with about 200 normal samples and 100 anomaly samples for each class. It provides various types of anomaly with pixel-level segmentation ground truths such as scratches, cracks, color, and missing components, posing a great challenge to AD. In our experiment, following [Roth et al., 2022], we resize all images to 256×256 and center crop the images to 224×224 for MVTec. The BTAD dataset contains approximately 2500 real-world industrial images of three products [Mishra et al., 2021], which is more challenging for AD. Since some anomalies are located in the edge regions of the image in BTAD, we only resize all images to 256 × 256 without center cropping.

Evaluation Metrics. To evaluate the results of AD methods, we employ the metrics of Area Under the Receiver Operating characteristic Curve (AUROC). Specifically, we adopt Image-AUROC and Pixel-AUROC to evaluate the accuracy of anomaly detection and anomaly localization respectively.

Implementation details. We adopt the U-net architecture proposed by [Dhariwal and Nichol, 2021] to implement the denoiser ϵθ(xt, t) of DDIM for sampling of p(x0). For each category of normal samples in MVTec/BTAD, we train a UNet denoiser ϵθ(xt, t) separately within 2000 epochs using an Adam optimizer (learing rate: 1e-4, weight decay: 5e-2) based on a single GeForce RTX 3090 GPU. In the training process, the batchsize is set to 8, and the timestep of DDIM is set to be 1000. After training, we utilize the trained denoiser for the proposed MDPS, and let T = 200, N = 10, ρ = 100. In Section 5.2, we will further discuss the selection of hyperparameters for MDPS in details.

5.1 Comparison with State-of-the-art

MVTec. We compare MDPS with several representative reconstruction-based methods on the MVTec dataset, including DRÆM [Zavrtanik et al., 2021], UniAD [You et al., 2022a], AnoGAN [Schlegl et al., 2017], AnoSeg [Song et al., 2022], AnoDDPM [Wyatt et al., 2022], AutoDDPM [Bercea et al., 2023], and RAN [Lu et al., 2023]. The results are shown in Table 1 and Figure 2.

From Table 1 and Figure 2, we can find that MDPS achieves the best performance in reconstruction quality as well as anomaly detection and localization especially when Ns = 16. Specifically, the Image-AUROC of MDPS with Ns = 16 outperforms the second best comparison method DRÆM by 0.8%, and outperforms the second best DMbased method RAN by 5.7%. The Pixel-AUROC of MDPS with Ns = 16 outperforms the second best DM-based method RAN by 0.6%, and shares the same values as the Pixel-AUROC of DRÆM. Although DRÆM shows the same anomaly localization performance as MDPS, the reconstruction results of DRÆM suffer from excessive blurring. Besides, DRÆM is a self-supervised method and performs badly for the real anomalies which differ significantly from the pseudo ones generated for training [Lu et al., 2023]. However, MDPS does not require any pseudo training samples, thus, has higher generalization ability than DRÆM.

BTAD. To further demonstrate the generalizability and superiority of our method, we compare MDPS with several recent state-of-the-art AD methods on the BTAD dataset, including four reconstruction-based methods (DRÆM ,VTADL [Mishra et al., 2021], AnoDDPM and AutoDDPM), two representation-based methods (PatchCore [Roth et al., 2022] and PaDiM [Defard et al., 2021]), and a normalizing flow based method (PyramidFlow [Lei et al., 2023]). The results are shown in Table 2. From Table 2, we can find that MDPS still achieves the competitive performance for anomaly detection and localization under the metric of average AUROC. Note that DRÆM shows a significant decrease performance on the BTAD dataset compared with the results of MVTec due to its limited generalization ability [Lu et al., 2023].

Besides, we compare MDPS with several state-of-the-art AD methods on time consumption in Table 3, including PatchCore, DRÆM and AutoDDPM. Thanks to the acceleration of DDIM, MDPS shows comparable computational cost compared with AutoDDPM and PatchCore, but cannot compete with the AE-based method DRÆM.

5.2 Ablation Study

In this section, we conduct a series of ablation experiments on MVTec to discuss the selection of hyperparameters and validate the designs of MDPS and anomaly scoring.

Selection of Hyperparameters. We empirically set T = 200 and N = 10 for MDPS. Then, we mainly focus on selection of the guidance scale ρ for Eqn.(17), since ρ is related to anomalies and is a key hyperparameter to influence the reconstruction quality. Several normal image reconstruction examples of MVTec and BTAD are displayed to show the influence of ρ in Figure 3. In Figure 3, partial normal texture details are destructed when ρ < 100, which would lead to misjudgment of anomalous pixels in the normal region; however, when ρ > 100, anomalous texture details appear in the reconstructed image, which would lead to misjudgment of normal pixels in the anomalous region. From Figure 3, we can find that ρ can control the sensitivity for anomalies in the normal image reconstruction, and MDPS has better reconstruction quality when ρ is set to 100.

Effectiveness of MDPS. To validate the designs of MDPS, we conduct a group of ablation experiments and display the results in Figure 4. In Figure 4, "vanilla DDIM" represents only using the sampling process of DDIM, i.e, m = 1 and ρ = 0; "w/o mask" represents m = 1; 'w/o posterior' represents ρ = 0, i.e., the problem of normal image reconstruction is modeled as prior sampling instead of posterior sampling for normal images. From Figure 4, we can find that MDPS achieves higher values of Image-AUROC and Pixel-AUROC than MDPS with m = 1 or ρ = 0 no matter for textural or object categories, which validates the modeling of posterior sampling and the design of mask image m in MDPS.

To further validate the effectiveness of MDPS, we display qualitative results of MDPS in Figure 5. From Figure 5, we can find that MDPS can reconstruct high-quality normal images given various test images with different anomalies.

Effectiveness of anomaly scoring. The effectiveness of anomaly scoring is related to the difference metric for two images and the number of posterior samples Ns. To validate the designed metric of Eqn.(18), we compare the results of Eqn.(18) with those of only pixel-level or perceptual-level metrics in Table 4. The results show that the anomly scores of "Pixel+Perceptual" can improve AUROC efficiently.

Then, we discuss the influence of Ns on anomaly scoring. The results with different values of Ns are shown in Table 5. It can be observed that larger value of Ns leads to further improvements of AUROC. In Figure 6, we display multiple normal samplers of a test image reconstructed by MDPS. These normal images exhibit some subtle differences. But the average of difference maps for these normal samplers can rectify misjudgment caused by a single one.

6 Conclusion

This paper proposes MDPS, a novel and highly interpretable UAD method. MDPS generates multiple normal images based on diffusion posterior sampling under Bayesian framework. Using a combination of pixel-level and perceptuallevel metrics, MDPS averages all difference maps between multiple reconstructed normal images and the test image to obtain the anomaly scores accurately. Exhaustive experiments show MDPS achieves high reconstruction quality and state-of-the-art performance for anomaly detection and localization compared with recent UAD methods, including 98.8% Imgae-AUROC and 97.3% Pixel-AUROC on MVTec, 99.5% Imgae-AUROC and 97.6% Pixel-AUROC on BTAD. However, MDPS suffers from high computational cost caused by diffusion posterior samplings. In future work, we would reduce the computational cost through knowledge distillation.

Acknowledgments

The work is supported in part by Natural Science Foundation of Guangdong Province (2022A1515010493), in part by Shenzhen Science and Technology Program (No. JCYJ20210324140407021), in part by Natural Science Foundation of China (No. 62372082) and Natural Science Foundation of Sichuan Province (No.2023NSFSC0485).