SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ruiying Lu $^{1\dagger}$ , YuJie Wu $^{2\dagger}$ , Long Tian $^{2*}$ , Dongsheng Wang $^{3}$ , Bo Chen $^{3}$ , Xiyang Liu $^{2}$ , Ruimin Hu $^{1}$
School of Cyber Engineering<sup>1</sup>, Software Engineering Institute<sup>2</sup>
National Key Laboratory of Radar Signal Processing<sup>3</sup>
Xidian University
{luruiying,tianlong}@xidian.edu.cn
Unsupervised image Anomaly Detection (UAD) aims to learn robust and discriminative representations of normal samples. While separate solutions per class endow expensive computation and limited generalizability, this paper focuses on building a unified framework for multiple classes. Under such a challenging setting, popular reconstruction-based networks with continuous latent representation assumption always suffer from the "identical shortcut" issue, where both normal and abnormal samples can be well recovered and difficult to distinguish. To address this pivotal issue, we propose a hierarchical vector quantized prototype-oriented Transformer under a probabilistic framework. First, instead of learning the continuous representations, we preserve the typical normal patterns as discrete iconic prototypes, and confirm the importance of Vector Quantization in preventing the model from falling into the shortcut. The vector quantized iconic prototypes are integrated into the Transformer for reconstruction, such that the abnormal data point is flipped to a normal data point. Second, we investigate an exquisite hierarchical framework to relieve the codebook collapse issue and replenish frail normal patterns. Third, a prototype-oriented optimal transport method is proposed to better regulate the prototypes and hierarchically evaluate the abnormal score. By evaluating on MVTec-AD and VisA datasets, our model surpasses the state-of-the-art alternatives and possesses good interpretability. The code is available at https://github.com/RuiyingLu/HVQ-Trans.
Anomaly detection is an essential task with increasingly wide applications in various areas, such as video surveillance [1], industrial inspection [2], and medical image analysis [3]. Due to the scarcity of anomalous samples, the unsupervised anomaly detection [4-7] methods gain wide attention by modeling the distribution of normal data only, and then identify the samples deviates from the normal profile as anomalies. Common approaches follow the one-for-one scheme [8, 9] by training separate models for different classes of objects, which is time-memory-consuming for real application and uncongenial to the object class with large intra-class diversity. Recently, a newly emerging one-for-all scheme [10] tries to use a unified model to detect anomalies from all the different object classes without any fine-tuning. Modeling high-dimensional data is notoriously challenging, and the problem becomes even more difficult to capture the multi-class distribution in a unified model precisely.
†Equal Contribution
* Corresponding Author
Under the unsupervised setting, a powerful approach to modeling data distribution follows the deep autoencoding frameworks, reckoning that a well-trained model with normal data will always reconstruct normal patterns regardless of the defects present in the input data. Thus, it is generally assumed that the reconstruction error will be larger for the anomalous input, making them distinguishable from the normal samples. However, this assumption may not always hold that sometimes the abnormal inputs can also be well reconstructed, which is named as "identical shortcut" issue [9-11]. Intuitively, compared to working extremely hard to learn the joint distribution, returning a direct copy of the input disregarding its content appears as a far easier solution. This phenomenon has been observed in existing researches [4, 9]. Furthermore, under the unified case, the "identical shortcut" issue becomes even more severe as the distribution of multi-class data is more complex [10]. This motivates us to enhance the discriminability of model encountering normal and anomalous samples.
Learning representations with continuous features have been the focus of many previous works [12- 14]. However, these methods lack a reliable mechanism to encourage the model to induce large reconstruction error on the anomalies, restricting the performances by the under-designed representation of the latent space. In recent researches, a branch of approaches [8, 9, 15] investigate the memory-augmented networks for mitigating the "identical shortcut" issue of AEs. Those approaches augment the deep autoencoder with a memory module to record the normal patterns in the normal training data, manifesting in different forms such as the memory set in the latent space [9], the fixed Transformer value matrix in the attention layer [10], or neighborhood-aware patch-level memory bank [8]. This kind of method hopes to obtain low reconstruction error for normal samples and highlight the reconstruction error if the input is not similar to normal data, that is, an anomaly. The most relevant items in the memory are retrieved and weighted averaging all the related memory content are aggregated into the decoder for reconstruction. However, the discrete memory items are recombined and weighted averaged, falling into an unknown continuous latent space which might be distorted. Intuitively, some anomalous regions can not be reconstructed by the discrete latent memory but could be decoded from the unknown latent space. To intrinsically mitigate the problems, we aim to learn a representative and discriminative discrete latent space for anomaly detection.
To preserve the typical normal patterns in the discrete latent space, we hope to successfully model critical features that usually span many dimensions in the normal data space, as opposed to focusing or spending capacity on noise and imperceptible details. Incorporating ideas from vector quantization (VQ), we model the discrete latent space as codebooks for each category, consisting of iconic prototypes learned from normal training data. During reconstruction, we replace the original encoding features with the nearest iconic prototypes, and then decoded with a VQ-based transformer decoder to intensify the use of iconic prototypes. As a result, the abnormal data point is flipped to a normal data point, highlighted by large reconstruction errors, as shown in Fig. 1. However, the model may suffer from the codebook collapse issue [16, 17]: At some point during training, a part of latent codes in the codebook may no longer work and the modeling capacity is limited by the discrete representations, resulting in collapsed reconstruction [18].
Thus, we further investigate the hierarchical nature of images and propose a hierarchical VQ framework by merging fine-grained and abstract features to prevent codebook collapse, which could also reduce the decoding search time and retain high inference speeds. In addition, most abnormal scoring methods are constrained to the observation space and can be fallible to complex data distributions. Therefore, we have introduced a hierarchical prototype-oriented optimal transport (OT) based optimization and anomaly detection scoring method to enhance the robustness and discriminability of our model for normal and anomaly samples.
In conclusion, we carefully tailor a variational autoencoding framework for unsupervised anomaly detection, called hierarchical vector quantized Transformer (HVQ-Trans). Our work contributes in the following ways: (1) We realize the learning of discrete normal representations by extracting
Figure 1: By replacing the continuous latent features with the normal iconic prototypes of corresponding category, the normal regions are reconstructed as normal patterns (shown in yellow boxes), while the anomalies are also reconstructed as normal (shown in red boxes).
(a) Overall Framework
(b) Iconic Prototype-oriented Learning
(c) Switching mechanism
Figure 2: (a) The overall framework of our HVQ-Trans. (b) Each VQ-based Layer replaces continuous features with iconic prototypes, equipped with the POT module to promote better learning and scoring. (c) The codebook and expert network are switched for individual image. (d) The detailed structure of each VQ-based Transformer decoder, where the prototypes are integrated via cross-attention.
(d) VQ-based Transformer Decoder
prototypes and propose a VQ-based transformer to address the "identical shortcut" issue by inducing large feature discrepancy for anomalies. (2) We develop a hierarchical VQ-based approach with switching mechanism to overcome the "prototype collapse" problem and effectively use multilevel feature representations to maximize the nominal information available. (3) A hierarchical prototype-oriented learning and anomaly scoring method is developed to guide prototype learning and dexterously measure the feature level anomaly score to robustly and accurately identify anomalies. (4) Extensive experiments demonstrate our method achieves state-of-the-art performances for anomaly detection and localization, and possesses enhanced interpretability through prototype visualization.
2.1 Overview
Our proposed model is a Transformer-based reconstruction method that assumes normal and anomaly cannot be reconstructed with comparable performances. In contrast to other Transformer-based autoencoding [19] with continuous embeddings, we focus on compressing input images into discrete representations and achieve discriminative reconstruction for anomaly detection and localization. We denote the set of normal images available at training time as  $\mathcal{X}_N$  ( $\forall \boldsymbol{x} \in \mathcal{X}_N : \boldsymbol{y}_x = 0$ ), with  $\boldsymbol{y}_x \in \{0,1\}$  denoting if an image  $\boldsymbol{x}$  is normal (0) or anomalous (1). Accordingly, we define the test sample as  $\forall \boldsymbol{x} \in \mathcal{X}_T : \boldsymbol{y}_x = \{0,1\}$ , including both the normal test images and abnormal test images. The model pipeline, shown in Figure 2, can be summarized as follows: i) The input image is fed into the pre-trained EfficientNet [20] to extract visual tokens by splitting 3-D feature maps; ii) The extracted tokens are passed through the cascaded vanilla transformer encoder for non-local multi-level feature aggregation; iii) Aggregated features at certain layers are hierarchically fed into the corresponding  $VQ$ -based layer to select the most relevant iconic prototypes; iv) The visual tokens are then successively fused with vector quantized embeddings via cascaded  $VQ$ -based transformer decoder for reconstruction; v) Decoded tokens are passed through the switching experts to reconstruct features, which possess flexibility in high diversity multi-class image scenarios; vi) Finally, anomaly detection and localization are achieved through a calibrated anomaly score map refined via prototype-oriented module by measuring the OT-based hierarchical feature discrepancy.
2.2 Improving Feature Reconstruction with Hierarchical Vector Quantization
Motivation: Normal memory augmentation was initially introduced by Gong et al. [9] and has obtained wide interests in unsupervised anomaly localization and detection. To record the "normal" appearance, image features are augmented by weighted averaging the similar patterns in the memory matrix. This augmentation is, however, rebuilding a continuous latent space again which might be distorted to contain abnormal patterns. Relied on the vector quantization, discrete variational autoencoder[21] learns the discrete latent space, but suffers from the issue of codebook collapse.
Therefore, to simultaneously learn the discrete representation and avoid codebook collapse, we proposed the hierarchical VQ-based framework.
Hierarchical Vector Quantized (HVQ) Transformer: Our proposed HVQ-Trans can be viewed as a communication system serving as an information bottleneck to better capture the normal patterns during training, which could be further generalized to test unknown images with arbitrary anomalies. We denote the input image as  $\pmb{x} \in \mathbb{R}^{H \times W \times 3}$ , then the  $N$  visual tokens  $\pmb{h}^0 = f_{\phi^0}(\pmb{x}) \in \mathbb{R}^{N \times C}$  are extracted by the pre-trained EfficientNet  $\phi^0$  [20] to be fed into the Transformer encoder. As shown in the graphical model of Fig. 3, HVQ-Trans comprises a cascaded vanilla Transformer encoder (vanTrans-enc) parameterized by  $\phi^l$  that encodes multi-layer patch embeddings as  $\pmb{h}^l = f_{\phi^l}(\pmb{h}^{l-1}) \in \mathbb{R}^{N \times C}$ . To further enlarge the normality and suppress the anomaly, we subsequently develop hierarchical VQ-based layers (VQLayer) to layer-wisely quantize the refined visual tokens  $\{\pmb{h}^l\}_{l=1}^L$  to the prototypes  $e_k^l$  in the learnable codebooks  $E^l \in \mathbb{R}^{K \times C}$ , as:
$$
\begin{array}{l} \boldsymbol {\theta} = \text {Q u a n t i z e} \left(\Upsilon^ {L} (\boldsymbol {h} ^ {L})\right) = \boldsymbol {e} _ {i} ^ {L}, \quad i = \min  _ {j} \| \Upsilon^ {L} (\boldsymbol {h} ^ {L}) - \boldsymbol {e} _ {j} ^ {L} \| _ {2} ^ {2}, \\ \boldsymbol {z} ^ {l} = \text {Q u a n t i z e} \left(\Upsilon^ {l} \left(\left[ \boldsymbol {h} ^ {l - 1}, \boldsymbol {\theta} \right]\right)\right) = \boldsymbol {e} _ {k} ^ {l}, \quad k = \min  _ {r} \| \Upsilon^ {l} \left(\left[ \boldsymbol {h} ^ {l - 1}, \boldsymbol {z} ^ {L} \right]\right) - \boldsymbol {e} _ {r} ^ {l} \| _ {2} ^ {2}, \quad l \in [ 1, \dots , L ], \tag {1} \\ \end{array}
$$
where  $\Upsilon^l (\cdot)$  refers to the layer-wise embedding function, and  $[\cdot ]$  denotes the concatenation operation. Intuitively, we hierarchically replace visual tokens  $h^{l}$  with their most similar prototypes in the codebook  $\pmb{E}^{l}$  as quantized vector  $\pmb{z}^{l}$  (note that this process can be lossy). Moreover, we find that merging fine-grained concrete information with abstraction-level semantics is critical for robust anomaly detection. Hence, we fuse the multi-level visual tokens with the global quantized vector  $\theta$  to learn hierarchical prototypes, maximizing the preserved nominal information.
To make full use of the quantized multi-level visual tokens  $z^l$  derived from Eq. 1, we further developed a cascaded VQ-based Transformer decoder (VQTrans-dec). As shown in Fig. 2 (d), given the quantized visual tokens  $z^l$  of the  $l^{th}$  layer and its
Figure 3: Graphical model of our HVQ-Trans.
corresponding output  $d^{l+1}$  from the last VQTrans-dec layer, the operation in each VQTrans-dec layer can be expressed as:
$$
\begin{array}{l} \boldsymbol {q} ^ {l} = \operatorname {M S A} (\text {q u e r y} = \mathbf {W} _ {q} \boldsymbol {d} ^ {l + 1}, \text {k e y} = \mathbf {W} _ {k} \boldsymbol {d} ^ {l + 1}, \text {v a l u e} = \mathbf {W} _ {v} \boldsymbol {d} ^ {l + 1}) + \boldsymbol {d} ^ {l + 1}, \\ \tilde {\boldsymbol {d}} ^ {l} = \operatorname {M C A} (\text {q u e r y} = \mathbf {W} _ {q} ^ {\prime} \boldsymbol {q} ^ {l}, \text {k e y} = \mathbf {W} _ {k} ^ {\prime} \boldsymbol {z} ^ {l}, \text {v a l u e} = \mathbf {W} _ {v} ^ {\prime} \boldsymbol {z} ^ {l}) + \boldsymbol {q} ^ {l}, \quad \boldsymbol {d} ^ {l} = \operatorname {F F N} (\tilde {\boldsymbol {d}} ^ {l}) + \tilde {\boldsymbol {d}} ^ {l}, \tag {2} \\ \end{array}
$$
where  $\mathrm{MSA}(\cdot)$  and  $\mathrm{MCA}(\cdot)$  share the same architecture as the standard multi-head self attention and multi-head cross attention in vanilla Transformer [22]. An essential aspect of our method is that, in the  $\mathrm{MCA}(\cdot)$  operation, the refined visual tokens from the previous layer  $q^{l}$  crossly attend to the prototypes of normal images  $z^{l}$ . Hence, the values at abnormal regions of  $q^{l}$  will be suppressed, and the abnormal signals could be rarely transmitted to the output terminal for reconstruction. Namely, the fewer reconstructed anomalies there are, the larger the reconstruction difference will be, which in turn leads to better performance in localizing and detecting anomalies.
During training, the typical normal patterns are recorded in the discrete variables, i.e., iconic prototypes. When encountering the anomalous during testing, the abnormal patterns will also be quantized as the normal prototypes, leading to larger feature migration and information loss, highlighted by higher reconstruction error. It is worth noting that while information loss triggered by VQ is exist for normal images, it is significantly more pronounced for anomaly images. This discrepancy in information loss serves as a key factor in effective anomaly detection. By investigating this difference, we can enhance the accuracy of our model in distinguishing abnormal regions.
Switching Mechanism: We adopt the switching mechanism to make the proposed model more suitable for multi-class anomaly detection. On the one hand, we develop the switching codebooks by assembling the independent iconic prototypes for each class. On the other hand, we develop the switching experts for flexible reconstruction of multi-classes, inspired by the Mixture of Experts (MoE) models [23, 24] and its sparsely-activated version [25]. Here, we choose to reconstruct at the
feature level rather than the pixel level, due to the invariance to subtle noise, rotation, and translation at the pixel level.
Specifically, the switching mechanism contains a multi-category classifier,  $M$  codebooks, and  $M$  reconstruction experts. The multi-category classifier takes the image feature as input and outputs the classification probability over the  $M$  category. In order to fit the data diversity property in the one-for-all setting, we switch the specific codebook (including a group of prototypes) from  $M$  codebooks according to the classification probability. Furthermore, we switch the individual reconstruction network (dubbed as expert) for feature reconstruction. The visual tokens from the last VQTrans-dec layer are depicted as  $\pmb{d}^{0} \in \mathbb{R}^{N \times C}$ , which is expected to reconstruct the input patch features  $\pmb{h}^{0} \in \mathbb{R}^{N \times C}$  as:
$$
\tilde {\boldsymbol {h}} ^ {0} = \Psi_ {m} (\boldsymbol {d} ^ {0}), \quad m = \underset {j} {\operatorname {a r g m a x}} p _ {j} (\boldsymbol {x}), \quad p _ {j} (\boldsymbol {x}) = \frac {\exp \left(\Omega (\boldsymbol {x}) _ {j}\right)}{\sum_ {j = 1} ^ {M} \exp \left(\Omega (\boldsymbol {x}) _ {j}\right)}, \tag {3}
$$
where  $\Omega(\cdot)$  is a classifier for producing logits, which are then normalized via a Softmax function over the total  $M$  experts.  $p_j(\cdot)$  is the probability of selecting the codebook and reconstruction expert, as shown in Fig. 2 (c). The codebook and expert  $\Psi_m$  with the highest probability are employed for reconstruction. The switching mechanism under the one-for-all setting will classify each input image into a single category and choose the corresponding codebook and expert for reconstruction. For the normal images, it is highly likely to be classified into the correct category and thus switch the proper reconstruction expert and codebook. For the abnormal images, there remains big uncertainty that which reconstruction expert and codebook will be switched, because the anomalies are unseen during training. Thus, the reconstruction uncertainty of the abnormal image is increased. Noting that the difference between normal and abnormal is the key factor deciding the anomaly detection performance. To this end, the uncertainty during anomalous sample switching could facilitate multi-class anomaly detection.
2.3 Prototype-oriented Learning and Scoring for Anomaly Detection
Motivation: During training with normal data, the HVQ-Trans enhances the point-wise correlations of the selected iconic prototypes from codebooks and the continuous visual features of normal images. However, it may despise the global relations between the above two sets. Motivated by previous efforts on OT theory and applications [26-29], we propose a hierarchical prototype-oriented optimal transport (POT) for anomaly detection, which is a transport solver defined within the scope of the basic distance between two unknown sampling sets to improve the tightness between the codebooks and the normal features. Meanwhile, at the testing stage, it is worth noting that the anomaly score adopted in the existing methods [10, 17] mainly concern the most significant difference of the score map between the input features and the reconstructed ones, as measured by the Euclidean distance. However, the importance of hierarchical differences at multiple feature levels is neglected. To this end, we also proposed a hierarchical POT-based anomaly scoring method to reinforce the identification of the score map and further boost the anomaly detection performance.
Learning Iconic Prototypes with POT: Each POT module, included within each VQLayer as shown in Fig. 2 (b), is responsible for enhancing consistency between codebooks and normal image features per layer. This enables the prototypes in codebooks to be more representative of normal patterns and less so of anomalous patterns. At the  $l$ -th layer, the codebook of each category contains a group of prototypes  $e^l = [e_1^l,\dots,e_K^l ]\in \mathbb{R}^{K\times C}$ . We omit the index  $l$  in the following for simplicity without causing ambiguity. To assemble the normal patterns conveyed by images, we represent  $N$  patches per image as an empirical distribution  $\mathbb{P}_h = \sum_{i = 1}^N\frac{1}{N}\delta_{h_i}$ , where  $h\in \mathbb{R}^{N\times C}$  is the features sampled from the latent variables of the HVQ-Trans. The prototypes serve to represent normal patterns across different classes. As a result, when attempting to identify suitable prototypes to reconstruct a specific normal image, each prototype is given equal importance. Thus, the distribution over normal prototypes can also be expressed as an empirical distribution  $\mathbb{P}_e = \sum_{j = 1}^K\frac{1}{K}\delta_{e_j}$ . In this way, the transport matrix  $M^{*}\in \mathbb{R}^{N\times K}$  from  $\mathbb{P}_h$  to  $\mathbb{P}_e$  can be estimated by  $M^{*} = \min_{M}\sum_{i = 1}^{N}\sum_{j = 1}^{K}\boldsymbol{M}_{i,j}\boldsymbol{C}_{i,j}$ , where the transport matrix  $M$  should satisfy  $\Pi ([\frac{1}{K} ],[\frac{1}{N} ]) = \{M|M\mathbf{1}_K = [\frac{1}{K}],M^T\mathbf{1}_N = [\frac{1}{N} ]\} .[\frac{1}{K} ]$  and  $[\frac{1}{N} ]$  are two uniform distributed prior defined in  $\mathbb{P}_h$  and  $\mathbb{P}_e$ , respectively. The cost matrix  $C\in \mathbb{R}^{N\times K}$  is defined as  $C_{i,j} = \sqrt{(h_i - e_j)^2}$ . In order to learn the prototypes of normal codebook at certain layer,
we define the average POT loss inspired by Sinkhorn algorithm [30] as:
$$
\mathbb {L} _ {P O T} = \min  _ {\boldsymbol {E}} \mathbb {E} _ {\boldsymbol {h} \sim F _ {\phi} (\boldsymbol {x})} \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {K} M _ {i, j} ^ {*} C _ {i, j} + \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {K} M _ {i, j} ^ {*} \ln M _ {i, j} ^ {*}. \tag {4}
$$
Calibrating Anomaly Score with POT: The anomaly score computed via Transformer-based methods always suffers from the sub-optimal distance measurement, which is usually calculated as the point-wise L2 norm of the reconstruction differences as  $s_{org} = \| f_{org} - f_{rec} \|_2^2$ . In this paper, on the one hand, we alleviate the mismatch by restricting the distance of prototypes and visual features during training. On the other hand, we further calibrate the anomaly score with multi-level POT at test time. Accordingly, in our proposed method, we note that the anomaly degree could also be reflected by the dissimilarity between visual features and normal iconic prototypes in the codebooks. As for the  $l$ -th layer, the dissimilarity is evaluated by  $s_{POT}^l = M^* C$ . The transport matrix  $M^*$  acts as a probability to re-weight the cost with different prototypes  $C$ , which measures the importance of different distances between image features and the normal prototypes. Therefore, we calibrate the multi-level anomaly score as  $s_{cab} = s_{org} + \lambda \sum_{l=1}^{L} s_{POT}^l$  for better anomaly detection.
2.4 Overall Optimization