OmiAD: One-Step Adaptive Masked Diffusion Model for Multi-class Anomaly Detection via Adversarial Distillation

We are grateful for the time you spent reviewing our paper in detail. Your insightful comments have been extremely helpful, and we deeply appreciate your input.

Q1: Comparison with Recent Multi-class UAD Methods

A1：We have compared our results on the MVTEC and VISA datasets with MambaAD, ViTAD, and ReContrast, as summarized in the table below. All results are sourced from the original publications. As shown, OmiAD consistently achieves superior performance across datasets.

Q2: Speed of SimpleNet

A2: We revisited the official SimpleNet implementation and found that it upsamples both the anomaly score map and feature map to a shape of batch × 1536 × 288 × 288. The resulting tensor is then transferred from GPU to CPU and stored as a NumPy array (as seen in line 113 of the source code in common.py). This operation accounts for the majority of the inference time, resulting in significantly slower inference. Since our evaluation was based on the official public code, we observed these slower inference times accordingly.

Based on your valuable suggestion, we optimized the implementation by removing the unnecessary upsampling step and the conversion to NumPy arrays, and streamlined the overall code. This optimization resulted in inference speeds that are consistent with the expected performance of the embedding-based structure and are lower than most methods reported in Table 2 of our paper. The updated inference times on the four datasets are shown below. As observed, the inference time remains longer than that of OmiAD, and thus does not affect the conclusions presented in our paper.

Q3: Overall method

A3: Should the paper be accepted, we will include a comprehensive illustration of the overall method in the camera-ready version, covering all components of OmiAD, including both AMDM and ASD modules.

Q4：Single-class UAD results

A4: We present the single-class anomaly detection performance on the MVTec and MPDD datasets as follows. These results will be detailed in the appendix for further reference. The single-class results on VisA and Real-IAD are currently under evaluation and will be included in the camera-ready version.

Single class UAD results for MVTec dataset：

Single-class UAD results for MPDD dataset：

Q5: Application-oriented focus and alignment with ICML's interests

A5: Thank you for your comment. While this work is focused on multi-class anomaly detection, we did not make innovations in areas like anomaly score calculation or anomaly classification, which are task-specific. Instead, our innovation lies in enhancing generative capabilities and inference speed, which are more general improvements that are highly relevant in the current machine learning field. In the context of anomaly detection, the proposed method improves detection performance by enhancing the completion of anomalies with the powerful generator.

Besides, the core contribution of our work is a novel one-step adaptive masked diffusion model, which integrates a random masking strategy into the one-step diffusion process to enhance generation quality and efficiency. This design offers broader applicability to tasks demanding efficient and high-quality generation. Furthermore, the rapid and efficient generation characteristics of the proposed method are closely aligned with the challenges currently faced in industrial anomaly detection. Thus, we have applied and validated the capabilities of our method in this task. Our work aligns closely with ICML's focus on advancing machine learning techniques with practical applicability.

Finally, anomaly detection is also an important direction of machine learning, and it has strong correlation with unsupervised machine learning, statistical machine learning and so on. In recent years, many related papers have been published at ICML.

Due to character limitations, we would be happy to discuss and address any remaining questions during the discussion phase.

We truly value the time and effort you invested in carefully reading our paper. Your thoughtful and constructive feedback is highly appreciated.

Q1:Effectiveness of Adaptive Masking for Localized Anomalies

A1: Thank you for raising this important point. We agree that anomalies heavily dependent on local features may pose a challenge for global masking strategies. To address this, our adaptive masking mechanism varies the mask ratio with the diffusion step, which allows the model to preserve fine-grained local information at earlier stages and gradually increases difficulty during training. Furthermore, since our model operates on feature-level inputs extracted by EfficientNet, it benefits from strong localized representations.

Additionally, the Real-IAD dataset contains many small-scale anomalies. OmiAD demonstrates excellent performance on this dataset, further validating its effectiveness in handling fine-grained, localized anomalies.

Q2: Computational overhead of the adversarial distillation

A2: The computational overhead of adversarial distillation during training is comparable to that of training the teacher model once. This efficiency stems from the fact that the One-step Generator $g_θ$ directly inherits the architecture of the AMDM U-Net (as described in Algorithm 1), and the discriminator shares weights with the one-step generator’s encoder, resulting in minimal additional cost. Moreover, the distillation phase trains the One-step Generator for only 150 epochs—substantially fewer than the 1000 epochs required for AMDM.

At inference time, OmiAD reduces the number of sampling steps from 100 (in AMDM) to just one, leading to a significant speedup.

In summary, adversarial distillation introduces minimal training overhead while enabling substantial inference acceleration through drastic reduction in sampling steps.

We sincerely appreciate the time and effort you invested in carefully reviewing our paper. Your insightful and constructive comments are greatly valued and have helped us improve the clarity and rigor of our work.

Q1: Inference stage and Anomaly Score Computation

A1:In the inference stage, we process both normal and anomalous data, following established methods like UniAD and HVQ-Trans for generating anomaly score maps. The pseudocode below outlines this process:

1. Input: $img$: Input image, $g_θ$: Trained one-step generator, $EfficientNet$: Feature extractor, $t_{init}$: Initial timestep.
2. Output: S: Pixel-wise anomaly score map
3. Procedure:
   1. Feature Extraction: $x_0=EfficientNet(img)$
   2. Noising: $x_t=\sqrt{\bar\alpha_t} \cdot x_0+\sqrt{1 - \bar\alpha_t}\cdot\epsilon, \quad t=t_{\text{init}}, \epsilon\sim\mathcal{N}(0, I)$
   3. One-step Reconstruction: $\hat{{x}}_0=g_θ(x_t)$
   4. Anomaly Score Computation: $S=\|x_0-\hat{{x}}_0\|_2^2$

Q2: Ablation study for $t_{init}$ and the rationale for choosing $t_{init}$ = 960

A2: We conducted ablations with different $t_{init}$ values and found the model remains robust between 800 and 960. We recommend $t_{init} = 960$ for the best trade-off between semantic preservation and anomaly detection.

We attribute the effective reconstruction at $t_{init} = 960$ to the high similarity among images within the same category. Using a pre-trained EfficientNet for feature extraction improved the model's ability to accurately reconstruct images.

Q3: Formula 25 derivation

A3: Thank you very much for your careful derivation. After reviewing it, we found that there was a typographical error. When substituting equation (22) into equation (21), we mistakenly wrote $\bar\beta$ as $\beta$. We have corrected and updated the derivation process, and equation (25) is now as follows: $\mathbb{E}_{q(x\_t\mid x\_g, t)\,p\_\theta(x\_0)}\Big[\Big\langle \epsilon\_\phi(x\_t, t)-\epsilon\_\psi(x\_t, t),\bar{\beta}\_t \nabla\_{x\_t}\log p\_\theta(x\_t)\Big\rangle\Big]$

$=-\mathbb{E}_{x\_g \sim p\_\theta(x\_0)x\_t\sim p(x\_t \mid x\_g)} \Big[\Big\langle \epsilon\_\phi(x\_t, t)-\epsilon\_\psi(x\_t, t),\frac{x\_t-\bar{\alpha}\_t x\_g}{\bar{\beta}\_t} \Big\rangle\Big]$

=-\mathbb{E}_{x\_g \sim p\_\theta(x\_0), z,\boldsymbol\epsilon \sim \mathcal{N}(0,I)} \Big[\Big\langle \epsilon\_\phi(x\_t, t)-\epsilon\_\psi(x\_t, t),\epsilon \Big\rangle\Big]\

This typographical error does not affect the subsequent loss calculation.

Q4: Importance of EfficientNet

A4: ResNet and EfficientNet are commonly used pre-trained feature extractors for anomaly detection. Replacing EfficientNet with ResNet34 on MVTec led to a performance drop (Image/Pixel AUROC: 93.7/96.5). Similar trends were observed in UniAD, where EfficientNet consistently outperformed ResNet. We therefore recommend using EfficientNet for better performance.

Q5: Input to the One step generator

A5: The input to the one-step generator is the feature map extracted by EfficientNet.

Q6: The diffusion model architecture

A6: We use a U-Net architecture as the backbone for the diffusion model. To evaluate the generality of our approach, we replaced the U-Net in OmiAD with DiT and U-ViT[1]. The results demonstrate that OmiAD maintains strong performance across different architectures, indicating that our method is architecture-agnostic.

Q7: Method for Calculating FPR

A7: In our experiments, the threshold is set based on the number of anomaly samples. We rank the samples by anomaly score and set the threshold to the score ranked at the position matching the number of anomaly samples.

Q8: More ablation study for $p_{min}$ and $p_{max}$

A8: We conducted ablation experiments to address this point. For further details, please refer to our response to reviewer QjD4 Question 4.

Q9: Speed of SimpleNet

A9: We identified that the slow speed was due to the official code transferring data from GPU to CPU and storing it as a numpy array. For further details, please refer to our response to reviewer 96ub Question 2.

Q10: Comparison of speed with Transfusion

A10: We added a time comparison with TransFusion. Despite fewer steps, TransFusion predicts masks and anomalies at each step and operates at the pixel level rather than in the latent space, both of which limit its inference speed.

Due to character limits, we welcome further discussion and are happy to address any remaining questions during the discussion phase.

[1]Bao F, et al. All are worth words: A vit backbone for diffusion models.

We greatly appreciate your thorough review of our paper. Your valuable feedback and constructive suggestions have provided us with a clearer direction for improvement.

Q1: Inference stage and Anomaly Score Computation

A1: In the inference stage, we process both normal and anomalous data. The methodology for generating anomaly score maps follows established approaches, such as those outlined in UniAD[1] and HVQ-Trans[2]. The pseudocode below outlines the process for the inference stage, including the generation of anomaly score maps:

1. Input: $img$: Original input image, $g_θ$: Trained one-step generator, $EfficientNet$: Feature extractor, $t_{init}$: Initial timestep.
2. Output: S: Pixel-wise anomaly score map
3. Procedure:
   1. Feature Extraction: $x_0=EfficientNet(img)$
   2. Noising: $x_t=\sqrt{\bar\alpha_t} \cdot x_0+\sqrt{1 - \bar\alpha_t}\cdot\epsilon,\quad t=t_{\text{init}}, \epsilon\sim\mathcal{N}(0, I)$
   3. One step Reconstruction: $\hat{{x}}_0=g_θ(x_t)$
   4. Anomaly Score Computation: $S=\|x_0-\hat{{x}}_0\|_2^2$

Q2: Justification for estimating $x_0$ using masked features $x_m^t$

A2: Anomaly detection requires leveraging global information to reconstruct the anomalous part into the normal modality. We use Equation (10) to predict $x_0$, and apply Equation (11) as a constraint. This formulation allows $\epsilon_\theta$ to compensate for the anomalous part and generate the normal part. We conducted experiments where we replaced $x_m^t$ with $x^t$ in Equation (10). The resulting performance, with Image/Pixel AUROC scores of 90.1/92.9, is inferior to that of AMDM, which achieved 98.4/97.5, thereby further validating our conclusions. In addition, some works, such as DiffMAE[3] and MaskDiT[4], utilize unmasked regions to predict $\hat{{x}}_0$

Q3: Fixed Mask Strategy

A3: Unlike A-Mask, where the mask probability varies with each time step, the mask probability in F-Mask remains constant. We set four different mask probabilities [0.1,0.2,0.3,0.4] and conducted experiments, selecting the best-performing combination as the F-Mask strategy for comparison.

Q4: More ablation study for $p_{min}$, $p_{max}$ and $t_{init}$

A4: We fixed $p_{min}$=0.1, with varying $p_{max}$, AMDM performance:

We fixed $p_{max}$=0.4, with varying $p_{min}$, AMDM performance:

The results show that AMDM maintains stable and strong performance across different $p_{\text{min}}$ and $p_{\text{max}}$ settings. The best performance is observed when $p_{\text{min}}=0.1$ and $p_{\text{max}}=0.4$.

Ablation study for $t_{init}$

We conducted ablation experiments on OmiAD at different $t_{init}$. The experimental results show that the model maintains performance robustness within a wide range of initial time steps (800 to 960). However, when $t_{init}$ = 1000, performance drops significantly due to the noisy input resembling pure noise, which hinders image recovery and reduces anomaly detection ability. We recommend using tinit = 960, as it provides optimal performance at both image and pixel levels.

Q5: Handling Large Anomalies

A5:Thanks to its multi-step generation process and global modeling capability, AMDM is capable of progressively reconstructing missing regions. We visualized AMDM's inference process on samples with missing transistors: before step 860, the model focuses on recovering the missing component, while after step 860, it shifts attention to refining fine-grained image details.

Since OmiAD is distilled from AMDM, it inherits this reconstruction ability and achieves normal restoration from anomalous inputs in a single step.

For pure white or pure black images, where no valid semantic information is available, the model produces meaningless outputs.

Q6: Essential References Not Discussed

A6:Thank you for the suggestion. We understand the importance of citing classical anomaly detection methods and will include references to SPADE, Padim, Cflow-AD, and other diffusion-based AD methods in the revised version of the paper.

We also release our single-class anomaly detection results as a public benchmark. For detailed performance, please refer to Reviewer 96ub, Question 3.

Due to character limitations, we would be happy to discuss and address any remaining questions during the discussion phase.

[1]You, Z, et al. A unified model for multi-class anomaly detection.

[2]Lu, R, et al. Hierarchical vector quantized transformer for multi-class unsupervised anomaly detection.

[3]Wei C, et al. Diffusion models as masked autoencoders.

[4]Zheng H, et al. Fast training of diffusion models with masked transformers.