Source-linked AI summary
A Unified Model for Multi-class Anomaly Detection
Zhiyuan You, Lei Cui, Yujun Shen, Kai Yang, Xin Lu, Yu Zheng, Xinyi Le
TL;DR
Existing anomaly detection methods often require separate models for different object classes, while unified multi-class detection is challenged by reconstruction shortcuts that can recover anomalies. UniAD addresses this with a unified reconstruction framework using layer-wise query decoding, neighbor masked attention, and feature jittering, achieving state-of-the-art results on MVTec-AD and CIFAR-10.
Problem
Existing methods train separate models for different object classes, while unified detection must model multiple class distributions without class labels or fine-tuning.
Method
UniAD combines a layer-wise query decoder, neighbor masked attention, and feature jittering to reduce reconstruction shortcuts and information leakage in unified anomaly detection.
Results
UniAD achieves state-of-the-art performance on MVTec-AD and CIFAR-10, significantly outperforming existing alternatives under the unified task setting.
Takeaways & Limitations
A unified model can handle multiple object classes without distinguishing them, supporting anomaly detection across categories within one framework.
Takeaways & Limitations
The study does not use category labels that might help the model fit multi-class data, and incorporating such labels remains future work.
Abstract
from arXiv · showhide
Despite the rapid advance of unsupervised anomaly detection, existing methods require to train separate models for different objects. In this work, we present UniAD that accomplishes anomaly detection for multiple classes with a unified framework. Under such a challenging setting, popular reconstruction networks may fall into an "identical shortcut", where both normal and anomalous samples can be well recovered, and hence fail to spot outliers. To tackle this obstacle, we make three improvements. First, we revisit the formulations of fully-connected layer, convolutional layer, as well as attention layer, and confirm the important role of query embedding (i.e., within attention layer) in preventing the network from learning the shortcut. We therefore come up with a layer-wise query decoder to help model the multi-class distribution. Second, we employ a neighbor masked attention module to further avoid the information leak from the input feature to the reconstructed output feature. Third, we propose a feature jittering strategy that urges the model to recover the correct message even with noisy inputs. We evaluate our algorithm on MVTec-AD and CIFAR-10 datasets, where we surpass the state-of-the-art alternatives by a sufficiently large margin. For example, when learning a unified model for 15 categories in MVTec-AD, we surpass the second competitor on the tasks of both anomaly detection (from 88.1% to 96.5%) and anomaly localization (from 89.5% to 96.8%). Code is available at https://github.com/zhiyuanyou/UniAD.
1 Introduction
The paper frames unified anomaly detection as detecting anomalies across multiple object classes with one model, without class labels or fine-tuning. It identifies reconstruction shortcuts as a central obstacle and introduces three architectural and training changes to address them.
- Existing anomaly detectors commonly train separate models for each object class, which can become memory-consuming as the number of classes grows.
- The unified task trains on normal samples from multiple categories and must detect anomalies across all categories without class labels or fine-tuning.
- A unified model must capture multiple class distributions while treating samples from other categories as anomalies when assessing a particular category.
- Popular reconstruction networks can learn an identical shortcut that reconstructs anomalous inputs well, undermining anomaly detection.
- UniAD uses a layer-wise query decoder, neighbor masked attention, and feature jittering to prevent shortcuts and information leakage during feature reconstruction.
2 Related work
Related work covers classical, pseudo-anomaly, reconstruction-based, and transformer-based anomaly detection. The paper emphasizes that reconstruction architectures and transformer shortcut behavior remain open concerns for unified detection.
- Classical anomaly detection methods extend one-class classification, while pseudo-anomaly methods convert detection into supervised classification, denoising, or segmentation.
- Reconstruction-based methods assume models trained on normal samples reconstruct normal regions successfully but fail on anomalous regions.
- Reconstruction models can also learn tricks that restore anomalies well, motivating strategies such as instructional information, memory, iteration, masking, and pseudo-anomalies.
- Transformer-based anomaly detectors use attention for image recovery, but prior methods directly apply vanilla transformer architectures without resolving the shortcut problem.
3 Method
UniAD revisits feature reconstruction for unified multi-class anomaly detection, where conventional networks can learn an identical shortcut that reconstructs anomalies too well. Its layer-wise query decoder, neighbor masked attention, and feature jittering are designed to prevent this shortcut and model normal data across classes.
- Revisiting feature reconstruction: MLP, CNN, and transformer reconstruction errors decrease while anomaly performance drops, revealing an identical shortcut that copies inputs to outputs.UniAD instead reconstructs anomalies as normal samples, making them distinguishable from normal features.
- Revisiting feature reconstruction: 18.1% and 13.4% are the transformer performance drops in anomaly detection and localization, respectively, when query embedding is removed.The query embedding is therefore important for modeling the normal distribution and preventing anomaly reconstruction.
- Revisiting feature reconstruction: The unified MLP case has smaller reconstruction error but much worse performance than the separate case, indicating a more severe identical shortcut.The separate-case curves average all categories, whereas the unified model must handle multiple categories simultaneously.
- UniAD framework: UniAD combines a Neighbor Masked Encoder and Layer-wise Query Decoder, with learnable queries fused at every decoder layer to model complex multi-class distributions.The decoder outputs reconstructed features after successive query-based feature fusion.
- UniAD framework: Neighbor masked attention prevents each feature point from relating to itself or neighboring regions, reducing information leakage that enables copying.The encoder and decoder replace full attention with this masking strategy.
- UniAD framework: Feature jittering perturbs feature tokens so the model learns normal-sample knowledge through denoising and recovers the correct message from noisy inputs.The disturbance is sampled from a Gaussian distribution and applied with a fixed jittering probability.
4 Experiment
Experiments evaluate UniAD’s unified anomaly detection and localization across MVTec-AD and CIFAR-10, including comparisons, qualitative analysis, and ablations. UniAD remains effective without fine-tuning across categories and benefits from its query, masking, and jittering designs.
- Anomaly detection on MVTec-AD: 96.5% anomaly-detection AUROC shows UniAD nearly matches its separate-case performance of 96.6% on unified MVTec-AD.The unified model detects anomalies across all categories without fine-tuning.
- Anomaly detection on MVTec-AD: 8.4% is UniAD’s margin over DRAEM for unified MVTec-AD anomaly detection, while DRAEM drops near 10% from the separate case.CutPaste drops 18.6% under the unified setting.
- Anomaly localization on MVTec-AD: 96.8% anomaly-localization AUROC is achieved by UniAD in the unified MVTec-AD setting, improving slightly over its separate-case result of 96.6%.UniAD surpasses PaDiM by 7.3%, whereas US and FCDD decline by 12.1% and 28.7%, respectively.
- Anomaly localization on MVTec-AD: UniAD reconstructs global and local structural anomalies and scattered or scratched texture anomalies, then localizes them through reconstruction differences.Figure 6 compares a normal reference, anomaly, reconstruction, ground truth, and predicted anomaly map.
- Anomaly detection on CIFAR-10: 8.3% is UniAD’s improvement over FCDD+OE on unified CIFAR-10 anomaly detection without auxiliary Outlier Exposure.US and FCDD almost lose their detection ability when five classes serve jointly as normal samples.
- Ablation studies: 18.1% and 13.4% are the gains from adding one query embedding to vanilla transformer for anomaly detection and localization, respectively.Layer-wise queries add 7.4% image-level and 3.7% pixel-level AUROC improvements; neighbor masking and feature jittering provide further gains.
5 Conclusion
UniAD unifies anomaly detection across multiple object classes while addressing the identical shortcut through query decoding, neighbor masked attention, and feature jittering. It achieves state-of-the-art performance on MVTec-AD and CIFAR-10, while category-label integration remains open.
- UniAD unifies anomaly detection for multiple classes with a framework that avoids fine-tuning across categories.
- The method uses layer-wise query decoding, neighbor masked attention, and feature jittering to counter the identical shortcut and input information leakage.
- UniAD achieves state-of-the-art performance on MVTec-AD and CIFAR-10, significantly outperforming existing alternatives.
- Category labels are not used, and incorporating them into the unified model remains an open direction.
- The unified setting differs from semantic anomaly detection, which treats category membership as the anomaly criterion rather than modeling normal samples within each category.
B Network architecture and training configurations
The architecture section describes CNN, Transformer, and MLP reconstruction baselines alongside UniAD’s neighbor masked encoder and layer-wise query decoder. Features are resized, concatenated, channel-reduced, reconstructed, and projected back to the original channels.
- The baselines share UniAD’s assistant modules and training configurations, enabling architectural comparisons under matched settings.
- CNN: CNN modifies ResNet-34 by preserving feature-map size and using channel dimensions C, C/2, C/2, and C across four stages.
- Transformer: Transformer uses four-layer encoder and decoder stacks with self-attention, cross-attention, feed-forward networks, residual connections, and positional embeddings.
- MLP: MLP replaces attention with linear projections and removes position embeddings while retaining layer normalization and ReLU activation.
- UniAD: UniAD includes a neighbor masked encoder and layer-wise query decoder, with feed-forward details provided separately.
- Feature extraction: Selected features are resized and concatenated into a feature map, then tokenized into H × W tokens except for CNN.
- Channel reduction: Channel reduction maps C_org to C before reconstruction, then a projection restores C_org to produce the reconstructed feature map.
- Visualization: A visualization-only decoder renders reconstructed backbone features into images using a reversed ResNet-34 architecture with transposed-convolution upsampling.
B.4 Training configurations on CIFAR-10
For CIFAR-10, the configuration uses multi-stage features to capture semantic differences between anomalous objects and normal samples.
- CIFAR-10 inputs use image size 224 × 224 and feature size 14 × 14.
- Features from stages 1 through 5 are resized and concatenated into a 720-channel feature map because deep features provide semantic information.
C Ablation studies
The ablation studies evaluate UniAD components on MVTec-AD under the unified setting using AUROC, with comparisons covering query embeddings, neighbor masked attention, and feature jittering.
- Ablation experiments are conducted on MVTec-AD and evaluated with AUROC under the unified case.
- Table A2 studies layer-wise query embedding, neighbor masked attention, and feature jittering, with default settings marked in blue.
C.1 Full ablation studies of our three designs
The ablation studies show that Neighbor Masked Attention and Feature Jittering each improve vanilla-transformer performance, while their combination is strongest with layer-wise query embedding.
- 96.2% for detection and 96.6% for localization are achieved when NMA and FJ are added together to a vanilla transformer with one query embedding.Adding either component individually also obviously improves the results.
- NMA and FJ are therefore effective individually in the one-query setting, while their combined effect is strongest under layer-wise query embedding.
- 96.5% for detection and 96.8% for localization are achieved when layer-wise query embedding, NMA, and FJ are combined.Under layer-wise query embedding, FJ performs best when cooperating with NMA.
C.2 Layer-wise query decoder
The Layer-wise Query Decoder fuses learnable queries, encoder embeddings, and previous-layer outputs through two attention modules, with design choice (a) selected as the final architecture. UniAD remains effective across several regression losses.
- Layer-wise query decoder: Each decoder layer after the first fuses a learnable query embedding, encoder embedding, and previous-layer output through two attention modules.Six assignments of these inputs to attention roles are evaluated.
- Layer-wise query decoder: Designs (a), (b), and (f) outperform the vanilla transformer, with design (a) achieving the best result without NMA or FJ.The previous layer’s outputs serve as attention key and value in these effective designs.
- Layer-wise query decoder: Adding NMA and FJ stably improves the three effective decoder designs, and design (a) remains the best-performing final choice.
- Loss function: Anomaly localization uses L2 differences for normalized MSE and cosine distance for cosine-distance loss.
- Loss function: The three loss functions achieve similar results, supporting UniAD’s use with different regression objectives.MSE is chosen because it is the most commonly adopted regression or reconstruction loss.
C.4 Backbone
Freezing the backbone is important for UniAD’s performance, and the model’s advantage does not come from larger capacity. EfficientNet and moderate-sized backbones are particularly suitable.
- Trainable or frozen: -30.6% for anomaly detection and -31.3% for localization occur when the backbone is trained rather than frozen.The authors speculate that training extracts indiscriminative features that are easy to reconstruct but unhelpful for anomaly detection.
- Backbone architecture: EfficientNet performs obviously better than ResNet, especially for anomaly detection.
- Backbone architecture: Moderate parameter size is more suitable for anomaly detection, including ResNet-34, EfficientNet-b2, and EfficientNet-b4.
- Complexity: UniAD’s advantage does not come from a larger model capacity when inference FLOPs and learnable parameters are compared at 224 × 224 image size.
D.2 Visualization results
Visualization results show that standard reconstruction networks can recover anomalous regions through an identical shortcut, whereas UniAD reconstructs anomalies as normal samples and localizes varied defect types across all 15 categories.
- Reconstruction visualization: MLP, CNN, and transformer networks tend to recover anomalous regions through an identical shortcut, while UniAD reconstructs anomalies as normal samples.
- Anomaly localization: Qualitative localization results cover all 15 MVTec-AD categories with a unified model.
- Anomaly localization: Figure A3 compares a normal reference, anomaly, reconstruction, ground truth, and predicted anomaly map.