The diffusion process entails the incremental alteration of the original image x0 into xT by incrementally adding Gaussian noise, thereby achieving the goal of corrupting the image. This process is also known as the forward (positive) noise addition process, which the formula can represent: (1)xt=xt−1αt+ϵt−11−αt,ϵt−1∼N(0,1)

This process transforms the data sample x0 into a noise sample xt, where t is randomly drawn from {0,1,…,T}, ϵt−1∼N(0,1) represents Gaussian noise, and {αt}t=1T are predefined hyperparameters, known as the noise schedule. These αt values are crucial as they control the variance of the noise added at each step. A higher αt implies less noise addition, preserving more of the original image’s characteristics at step t, while a lower αt means more noise is introduced, rapidly transforming the image towards a pure noise state. Typically, αt is set in a decreasing sequence, often following a linear or cosine-based schedule, to ensure a smooth transition from the original image to noise. The choice of the αt schedule significantly impacts the model’s ability to learn the underlying data distribution; a well-designed schedule can facilitate better convergence during training and improve the quality of generated samples.

By iteratively deriving from the above formula, we can obtain the direct transformation formula from x0 to xt as follows: (2)xt=x0α¯t+ϵt1−α¯t,ϵt∼N(0,1)given that α¯t=∏i=0tαi, which is a hyperparameter set according to the noise schedule, and ϵt∼N(0,1) is also Gaussian noise. Both Eqs. (1) and (2) can describe the forward noise addition process, with the former used for gradually corrupting an image and the latter for corrupting an image in one step.

The diffusion process encompasses the addition of noise to the image, whereas the reverse process engages in denoising. Knowing the true distribution of each step in the reverse process, denoted as q(xt−1|xt), and beginning with random noise xt∼N(0,1), a gradual denoising procedure can yield an authentic sample. The parameter βt, related to αt through βt=1−αt, plays a significant role in the reverse process. A larger βt implies a more aggressive denoising step, which might introduce more variance in the generated samples but could also lead to faster convergence towards the original data distribution. In the reverse process equation: (3)xt−1=1αt(xt−1−αt1−α¯tϵθ(xt,t))+β~tz,z∼N(0,1)given β~t=1−α¯t−11−α¯tβt. Since the real noise ϵt in the forward process equations is not allowed to be used during the restoration process, the key is to obtain a model ϵθ(xt,t) that predicts the noise from xt and t, where θ represents the training parameters of the model, and z∼N(0,1) is another Gaussian noise used to represent the difference between the prediction and the actual noise. The balance between αt and βt in the reverse process determines the stability and quality of the generated output, with improper choices potentially leading to blurry or inconsistent results.

The reconstruction network is designed to predict the noise ε related to xt and t, based on a U-Net-like architecture integrating ResNet [21], PixelCNN++ [22], and Transformer [23]. In the noise prediction module, the encoder uses ResNet blocks with skip connections for hierarchical feature extraction, the middle layers adopt PixelCNN++ to model pixel-level spatial dependencies, and the decoder incorporates Transformer layers to capture long-range dependencies. For the module-to-module collaboration, U-Net’s skip connections transfer low-level details from the encoder to the decoder, enabling the decoder to combine high-level semantics with fine-grained information. ResNet, PixelCNN++, and Transformer components work together to enhance feature representation, refinement and synthesis. The choice of this combination is theoretically justified as U-Net suits detail recovery, ResNet boosts representational power, PixelCNN++ captures local structure, and Transformer overcomes the limitation of convolutional layers in long-range dependency modeling, achieving a balance between local and global understanding for better noise prediction.

In addition to using a U-Net-like structure as the backbone, each block in the decoder is connected to its corresponding encoder block through skip connections. Early diffusion works [24,25] have demonstrated the importance of skip connections in this architecture. On the one hand, in the encoder-decoder structure, they reduce the information loss during the downsampling process by directly passing the feature maps to the decoder, thus enriching its representation. On the other hand, skip connections accelerate the training process and improve the generation quality by providing additional gradient paths and preserving important features, enabling more efficient and effective diffusion modelling. However, the dense block-wise skip connections have become a bottleneck. A large number of block-wise skip features have to be concatenated to the features and merged along the channel dimension within the corresponding decoder blocks, which consumes excessive computational costs. Such inefficiency hinders the scalability of the model. To address this issue, we propose Sparse-Skip connections, in which the skip connections are applied only every few blocks instead of after each block. This approach improves the performance, and our experiments show that, compared with the densely connected skip design, it indeed yields better results. Fig. 2 illustrates the differences between the two U-Net structures.

Initially, the input image x0 is corrupted at a random time step t during the diffusion process, resulting in xt as per Eq. (2). As the time step increases, the input image x0 gradually loses its distinctive features—the anomalous pixels lose their sharp characteristics and approach an isotropic Gaussian distribution, where x0 represents a normal image or a composite abnormal image. Based on this, the reconstruction network can iteratively obtain the reconstructed anomaly-free x^0 through Eq. (3).

The segmentation model, based on U-Net, identifies and classifies each pixel in the image as normal or abnormal, thus achieving precise segmentation of the anomalous areas. It compares the reconstructed image x^0 with the original image x0 and learns to detect anomalies by learning the commonalities and differences between them.

We trained the reconstruction model and segmentation model together. The reconstruction model learns the entire distribution of a normal sample (y=0) by minimizing the following loss function: (7)ℒnoise=(1−y)(‖ϵts−ϵθ(xts)‖2+‖ϵtb−ϵθ(xts)‖2)2

The segmentation model uses the commonalities and differences between x0 and x^0s to predict pixel-level anomaly scores that are as close as possible to the ground truth. Segmentation loss is defined as: (8)ℒmask=Smoothℒ1(M,M^)+γℒfocal(M,M^)where M is the truth mask of the input image, and M^ is the output of the segmentation model. Smoothℒ1 and ℒfocal were applied simultaneously to reduce oversensitivity to outliers and accurately segment difficult exception examples. γ∈R+ is the hyperparameter that controls the importance of ℒfocal. Therefore, the total loss of Diff-Fastener use for combined training is: (9)ℒtotal=ℒnoise+ℒmask

To achieve higher inference speed and reconstruction quality, we still use one-step norm-guided estimation in the inference stage. After the segmentation model predicts the pixel-level anomaly fraction M^, we take the average of the first K anomaly pixels in M^ as the image-level anomaly fraction. This inference method is hundreds of times faster while maintaining sampling quality.

Fig. 8 shows the visualization of detection results for various types of fastener anomalies. It can be observed that Diff-Fastener does not require prior learning of specific anomaly categories. By simply training with samples of healthy fastener images, it can detect various fastener anomalies and successfully repair them, while maintaining high image fidelity. Among them, the model performs best in detecting two types of anomalies: broken and missing fasteners. It indicates that scenarios where objects transition from absent to present are clearly more suitable for anomaly detection based on generative models.

By using cross-validation, we have obtained a more reliable assessment of the performance of both the supervised and unsupervised models. The average performance metrics from cross-validation provide a better understanding of the models’ generalization ability.

(1) Supervised models

This paper trains and tests the current advanced supervised detection model on the rail fastener dataset, and compares the results based on the recall rate and accuracy index at the label-level (i.e., at the whole image level). The comparison results are presented in Table 1, and the visualization of the comparison results can be seen in Fig. 9. Diff-Fastener learns fastener features by simply adding noise to and denoising healthy data, it can achieve approximately 99.9% accuracy rates and 99.9% recall rates in the label-level, that is, it can judge whether there is an anomaly in each fastener image almost perfectly. This result is of great significance, as it demonstrates the feasibility of the unsupervised detection model in the field of abnormal fastener detection.

(2) Unsupervised models

The MVTec [29] and VisA [30] datasets are essential benchmarks for anomaly detection. MVTec, with its 5000+ high-resolution images across 15 categories and pixel-precise anomaly annotations, is a gold standard for evaluating industrial inspection algorithms. VisA, offering 9621 normal and 1200 abnormal images in 12 categories, provides a new testbed for diverse industrial scenarios.

We compared the anomaly detection performance of Diff-Fastener with three feature embedding-based methods and three generative model-based methods; the results are presented in Table 2. The visualization of the comparison results can be seen in Fig. 10. The results show that for average Image AUROC, our model outperforms the feature embedding-based SOTA by 1.6% and the generative model-based SOTA by 0.3%. For average Pixel AUROC, our model outperforms the feature embedding-based SOTA by 2.9% and the generative model-based SOTA by 0.4%. The performance of Diff-Fastener is significantly better than previous models.

Table 3 shows the results of the ablation study. From the first and second rows of the table, compared to traditional iterative denoising, the inference speed of One-Step denoising is greatly improved, approximately 320 times faster than before. In the third row, the use of the Norm-Guided paradigm enhances the model’s ability to preserve image details while effectively removing anomaly regions of varying sizes, achieving more accurate anomaly detection, with a slight decrease in speed. From the fourth to sixth rows, it is evident that the introduction of CBAM (Convolutional Block Attention Module) [31] and dilated convolutions [32] improves the model’s semantic awareness while maintaining relatively efficient speed. In addition, sparse skip has enhanced the network performance and detection efficiency. Ablation experiments have demonstrated the effectiveness of the above methods.

Fig. 11 shows the effects of the Norm-Guided paradigm and DACM. Fig. 11a–f represents several different anomalous samples. A indicates a diffusion model structure using only one-step denoising. The introduction of one-step denoising leads to a loss of detail in the reconstruction results, limiting adaptability to different types of anomalies. From the image, it can be seen that the reconstructed semantic edges of the fasteners are unclear, with some even experiencing spatial distortion, and the overall tone of the image appears unnatural. B represents a one-step denoising diffusion model after the introduction of the Norm-Guided paradigm, which allows for more refined reconstruction rather than simple rough recovery. Additionally, the model can effectively handle anomalies in varying-sized areas with two scales of noise, enhancing the model’s adaptability. However, there is still an issue of excessive image smoothing. After introducing DACM, the model can emphasize important features while suppressing less significant ones during the reconstruction process. It also captures a broader context in the feature extraction stage, further enhancing the model’s reconstruction capability. As shown in Fig. 11, our reconstructed images contain richer detail information and are more aligned with real-world situations.