Unilaterally Aggregated Contrastive Learning with Hierarchical Augmentation for Anomaly Detection

Unsupervised Contrastive Learning. Unsupervised contrastive learning aims to learn representations from unlabeled data. The premise is that similar samples as positive pairs are supposed to have similar representations. The practical way to create positive pairs is to apply random augmentation to the same sample independently, e.g. two crops of the same image or multi-modal views of the same scene. Formally, let $x$ be an anchor, $D_{x}^{+}$ and $D_{x}^{-}$ be the sets of positive and negative samples w.r.t. $x$, respectively. We consider the following common form of the contrastive loss:

$$\begin{split}&{\cal L}_{\rm cons}(x,D_{x}^{+},D_{x}^{-}):=\\ &-\frac{1}{|D_{x}^{+}|}\sum_{x^{\prime}\in D_{x}^{+}}{{\rm log}\frac{e^{z(x^{\prime})^{T}z(x)/\tau}}{\sum_{x^{\prime}\in D_{x}^{+}\cup D_{x}^{-}}{e^{z(x^{\prime})^{T}z(x)/\tau}}}},\end{split}$$

(1)

where $|D_{x}^{+}|$ denotes the cardinality of $D_{x}^{+}$, $z(\cdot)$ extracts the $\ell_{2}$-normalized representation of $x$ and $\tau>0$ is a temperature hyper-parameter. We, in this work, specifically consider the simple contrastive learning framework i.e. SimCLR based on instance discrimination. For a batch of unlabeled images ${\cal B}:=\{x\}_{i=1}^{N}$, we first apply a composition of pre-defined identity-preserving augmentations $\cal T$ to construct two views ${\tilde{x}_{i}^{1}}:=t_{1}(x_{i})$ and ${\tilde{x}_{i}^{2}}:=t_{2}(x_{i})$ of the same instance $x_{i}$, where $t_{1},t_{2}\sim{\cal T}$. The contrastive loss of SimCLR is defined as follows:

$$\begin{split}{\cal L}_{\rm SimCLR}({\cal B};{\cal T})=\frac{1}{2N}\sum_{i=1}^{N}{\cal L}_{\rm cons}(\tilde{x}^{1}_{i},\{\tilde{x}^{2}_{i}\},{\cal\tilde{B}}-\{\tilde{x}^{1}_{i},\tilde{x}^{2}_{i}\})\\ +{\cal L}_{\rm cons}(\tilde{x}^{2}_{i},\{\tilde{x}^{1}_{i}\},{\cal\tilde{B}}-\{\tilde{x}^{2}_{i},\tilde{x}^{1}_{i}\}),\end{split}$$

(2)

where ${\cal\tilde{B}}:={\cal\tilde{B}}^{1}\cup{\cal\tilde{B}}^{2}$, ${\cal\tilde{B}}^{1}:=\{\tilde{x}^{1}\}_{i=1}^{N}$ and ${\cal\tilde{B}}^{2}:=\{\tilde{x}^{2}\}_{i=1}^{N}$. In this case, $D_{\tilde{x}_{i}^{1}}^{+}:=\{\tilde{x}_{i}^{2}\}$, $D_{\tilde{x}_{i}^{2}}^{+}:=\{\tilde{x}_{i}^{1}\}$ and $D_{\tilde{x}_{i}^{2}}^{-}=D_{\tilde{x}_{i}^{1}}^{-}:={\cal\tilde{B}}-\{\tilde{x}_{i}^{2},\tilde{x}_{i}^{1}\}$.

Supervised Contrastive Learning. SupCLR [29] is a supervised extension of SimCLR by considering class labels. Different from the unsupervised contrastive loss in Eq. 1 where each sample has only one positive sample, i.e. the augmented view of itself, there are multiple positive samples sharing the same class label, resulting in multiple clusters in the representation space corresponding to their labels. Formally, given a batch of labeled training samples ${\cal C}:=\{(x_{i},y_{i})\}_{i=1}^{N}$ with class label $y_{i}\in{\cal Y}$, ${\cal\tilde{C}}^{1}:=\{(\tilde{x}^{1}_{i},y_{i})|{\tilde{x}}^{1}_{i}\in{\cal\tilde{B}}^{1}\}_{i=1}^{N}$ and ${\cal\tilde{C}}^{2}:=\{(\tilde{x}^{2}_{i},y_{i})|{\tilde{x}}^{2}_{i}\in{\cal\tilde{B}}^{2}\}_{i=1}^{N}$ are the two sets of augmented views. The supervised contrastive loss is given as follows:

$${\cal L}_{\rm SupCLR}({\cal C};{\cal T})=\frac{1}{2N}\sum_{i=1}^{2N}{\cal L}_{\rm cons}({\tilde{x}_{i}},D_{x_{i}}^{+},D_{x_{i}}^{-}),$$

(3)

where $D_{\tilde{x}_{i}}^{+}:=\{x|(x,y_{i})\in{\cal\tilde{C}}^{1}\cup{\cal\tilde{C}}^{2}\}-\{{\tilde{x}_{i}}\}$ and $D_{\tilde{x}_{i}}^{-}:={\cal\tilde{B}}-\{\tilde{x}_{i}\}-D_{\tilde{x}_{i}}^{+}$.

Recall that a good representation distribution for AD entails a concentrated grouping of inliers and an appropriate dispersion of outliers. Given that only inliers are available for training, a natural question arises: how to obtain outliers? Due to the inaccessibility of real-world outliers, some attempts are investigated to create virtual outliers, aiming at a trade-off in various manners, such as through transformations [51, 3, 26, 33] or by sourcing them from additional datasets, known as OE [25]. These methods, relying on virtual outliers, display the superiority over their counterparts based on inliers only; however, they all fall short in fully addressing both the requirements of a good representation distribution for AD.

Following the success of introducing virtual outliers, in this work, we directly treat the goal of encouraging inlier concentration and outlier dispersion as the optimization objective via a pure contrastive learning framework. We particularly design a novel contrastive loss, namely UniCLR, to unilaterally aggregate inliers and disperse outliers. Different from the existing contrastive learning methods for AD [49, 48, 61, 51] that equally treat each instance from inliers and virtual outliers as one class and perform universal instance discrimination, we take all inliers as one class while each outlier itself a distinct class.

Formally, given a training inlier set ${\cal D}_{\rm in}$, we first apply distributionally-shifted augmentation ${\cal S}$, e.g. rotation, to inliers to create a set of virtual outliers ${\cal D}_{\rm vout}\equiv\{s(x)|x\in{\cal D}_{\rm in}\ \land s\in{\cal S}\}$. Note that ${\cal D}_{\rm in}$ and ${\cal D}_{\rm vout}$ are disjoint. For each image $x_{i}\in{\cal D}_{\rm in}/{\cal D}_{\rm vout}$, we further apply identity-preserving augmentations ${\cal T}$ to create two views of $x_{i}$ and finally obtain ${\cal\tilde{D}}_{\rm in}:={\cal\tilde{D}}_{\rm in}^{1}\cup{\cal\tilde{D}}_{\rm in}^{2}$ and ${\cal\tilde{D}}_{\rm vout}:={\cal\tilde{D}}_{\rm vout}^{1}\cup{\cal\tilde{D}}_{\rm vout}^{2}$, based on which we prepare a batch ${\cal\tilde{B}}:={\cal\tilde{D}}_{\rm in}\cup{\cal\tilde{D}}_{\rm vout}$ for training. The contrastive objective is given as:

$${\cal L}_{\rm UniCLR}({\cal D}_{\rm in}\cup{\cal D}_{\rm vout};{\cal T})=\frac{1}{|{\cal\tilde{B}}|}\sum_{i=1}^{|{\cal D}_{\rm in}|+|{\cal D}_{\rm vout}|}{{\cal L}^{i}_{\rm UniCLR}},$$

(4)

$$\begin{split}&{\cal L}_{\rm UniCLR}^{i}=\left\{\begin{aligned} &{\cal L}_{\rm cons}({\tilde{x}_{i}}^{1},{\cal\tilde{D}}_{\rm in}-\{{\tilde{x}_{i}}^{1}\},{\cal\tilde{D}}_{\rm vout})+\\ &\quad{\cal L}_{\rm cons}({\tilde{x}_{i}}^{2},{\cal\tilde{D}}_{\rm in}-\{{\tilde{x}_{i}}^{2}\},{\cal\tilde{D}}_{\rm vout}),x_{i}\in{\cal D}_{\rm in},\\ &{\cal L}_{\rm cons}({\tilde{x}_{i}}^{1},\{{\tilde{x}_{i}}^{2}\},{\cal\tilde{B}}-\{{\tilde{x}_{i}}^{2},{\tilde{x}_{i}}^{1}\})+\\ &\quad{\cal L}_{\rm cons}({\tilde{x}_{i}}^{2},\{{\tilde{x}_{i}}^{1}\},{\cal\tilde{B}}-\{{\tilde{x}_{i}}^{1},{\tilde{x}_{i}}^{2}\}),x_{i}\in{\cal D}_{\rm vout}.\end{aligned}\right.\end{split}$$

(5)

Our formulation is structurally similar to SimCLR (Eq. 2) and SupCLR (Eq. 3), with some modifications for AD in consideration of both inlier aggregation and outlier dispersion. Though our method is originally designed for inliers without class labels, it can be easily extended to the labeled multi-class setting where inliers sharing the same label are positive views of each other while samples from either other classes or augmented by $\cal S$ are negative.

Soft Aggregation. Data augmentation plays a central role in contrastive representation learning [10]. Though the commonly adopted augmentation pipeline in contrastive learning has witnessed the advanced progress in diverse downstream tasks, we observe that excessive distortion applied on the images inevitably shifts the original semantics, inducing outlier-like samples [59, 1]. Aggregating these semantic-drifting samples hinders the inlier concentration. A straightforward solution is to restrict the augmentation strength and apply weak augmentations to inliers; however, it cannot guarantee learning class-separated representations, i.e. reliably aggregating different instances with similar semantics [22, 60]. To take advantage of the diverse samples by strong augmentations while diminishing the side effects of outliers, we propose to aggregate augmented views of inliers with a soft mechanism based on the magnitude of deviation from the inlier distribution, and follow the notions defined in Eq. 1 to formulate it as follows:

$$\begin{split}&{\cal L}_{\rm soft\_cons}(x,D_{x}^{+},D_{x}^{-}):=\\ &-\frac{1}{\sum_{x^{\prime}\in D_{x}^{+}}{w_{x}}{w_{x^{\prime}}}}\sum_{x^{\prime}\in D_{x}^{+}}{{\rm log}\frac{{w_{x}}{w_{x^{\prime}}e^{z(x^{\prime})^{T}z(x)/\tau}}}{{\sum_{x^{\prime}\in D_{x}^{+}\cup D_{x}^{-}}{e^{z(x^{\prime})^{T}z(x)/\tau}}}}},\end{split}$$

(6)

where $w_{x}$ is the soft weight indicating the importance of sample $x$ in aggregation. According to Eq. LABEL:eq:soft, a positive pair of $x$ and $x^{\prime}$ receives more attention only if the corresponding $w_{x}$ and $w_{x^{\prime}}$ are both sufficiently large. Specifically, we measure $w_{x}$($w_{x^{\prime}}$) for $x(x^{\prime})$ by calculating the normalized average similarities with other inliers, i.e. $D_{x}^{+}$, as follows:

$$\omega_{x_{i}}=\frac{\sum_{x_{j}\in D_{x}^{+}\backslash\{x_{i}\}}e^{z(x_{i})^{T}z(x_{j})/\tau_{\omega}}}{\sum_{x_{k}\in D_{x}^{+}}{\sum_{x_{j}\in D_{x}^{+}\backslash\{x_{k}\}}}e^{z(x_{k})^{T}z(x_{j})/\tau_{\omega}}},$$

(7)

where $\tau_{\omega}$ controls the sharpness of the weight distribution. Intuitively, if one is away from all other inliers, there is a high probability that it is an outlier and vice versa. We apply soft aggregation (SA) only on inliers, i.e., replacing ${\cal L}_{\rm cons}$ with ${\cal L}_{\rm soft\_cons}$ only for $x_{i}\in{\cal D}_{\rm in}$ in Eq. 5.

Though the proposed UniCLR is reasonably effective for aggregating inliers and separating them from outliers, one can further improve the performance by prompting a more concentrated distribution for inliers. Inspired by the success of deep supervision in classification, we propose to aggregate inliers with hierarchical augmentation (HA) at different depths of the network based on the strengths of data augmentation. In alignment with the feature extraction process that shallow layers learn low-level features while deep layers emphasize more on high-level task-related semantic features, motivated by curriculum learning (CL) [2], we set stronger augmentation strengths for deeper layers and vice versa, aiming at capturing the distinct representations of inliers from low-level properties and high-level semantics at shallow and deep layers, respectively. To this end, we apply a series of augmentations at different network stages and gradually increase the augmentation strength as the network goes deep. Each stage is responsible for unilaterally aggregating inliers and dispersing virtual outliers generated with the corresponding augmentation strengths.

Formally, we have four sets of augmentation $T_{i}$, corresponding to four stages $res_{i}$ in ResNet [24]. Each $T_{i}$ is composed of the same types of augmentations but with different augmentation strengths. Extra projection heads $g_{i}$ are additionally attached at the end of $res_{i}$ to down-sample and project the feature maps with the same shape as in the last stage. With $T_{1}\sim T_{4}$ applied to inliers ${\cal D}_{\rm in}$ and outliers ${\cal D}_{\rm vout}$, we extract their features $z_{i}(x)$ with $res_{i}$ and $g_{i}$:

$$z_{i}(x)=g_{i}(res_{i}(T_{i}(x))),i=1,2,3,4.$$

(8)

Based on the features extracted by projector $g_{i}$, we separately perform unilateral aggregation for inliers and dispersion for outliers at each stage. The overall training loss can be formulated as:

$${\cal L}_{\mathrm{all}}=\frac{1}{4}\sum_{i=1}^{4}{\lambda_{i}{\cal L}_{\rm UniCLR}({\cal D}_{\rm in}\cup{\cal D}_{\rm vout});T_{i}),}$$

(9)

where $\lambda_{i}$ balances the loss at different network stages.

Through enforcing supervision in shallow layers, the resulting inlier distribution under the strong augmentation becomes more compact and more distinguishable from outliers.

During testing, we remove all four projection heads $g_{i}$. While the existing methods [51, 48] depend on specially designed detection score functions to obtain decent results, we observe that using the simple cosine similarity with the nearest one in the learned feature space is sufficiently effective. The detection score $s_{i}$ for a test example $x_{i}$ is given as:

$$s_{i}(x_{i};\{x_{m}\})={\rm\mathop{max}}_{m}\ {\rm cosine}(f(x_{i}),f(x_{m})),$$

(10)

where $\{x_{m}\}$ denotes the set of training samples and $f(\cdot)$ extracts the $\ell_{2}$ normalized representation at the end of $res_{4}$. Following [51, 3], we observe improved performance using an ensemble of representations by test-time augmentation.

We use ResNet-18 [24] for all experiments to ensure fair comparison with [51, 49]. Our models are trained from scratch using SGD for 2,048 epochs, and the learning rate is set to 0.01 with a single cycle of cosine learning rate decay. Following [10], we employ a combination of random resized crop, color jittering, horizontal flip and gray-scale with increasing augmentation strengths for $T_{1}\sim T_{4}$ to generate positive views while use rotation $\{90^{\circ},180^{\circ},270^{\circ}\}$ as the default $\cal S$ to create virtual outliers. Thus, ${\cal D}_{\rm in}$ comprises all original training samples and $\lvert{\cal D}_{\rm out}\rvert$ triples $\lvert{\cal D}_{\rm in}\rvert$ by applying $s\in\cal S$ on $x\in{\cal D}_{\rm in}$. We maintain a 1:3 ratio of inliers to virtual outliers during mini-batch training. Detailed augmentation configurations are available in the supplementary material. We exclusively apply SA at the last residual stage, i.e. $res_{4}$ where the strongest augmentations are employed. We set $\tau$ and $\tau_{\omega}$ as 0.5. For OE [25], we use 80 Million Tiny Images [53] as the auxiliary dataset, excluding images from CIFAR-10.

Unlabeled One-class. In this setting, a single class serves as the inlier, while the remaining classes act as outliers. Following [26, 51, 49, 3], the experiments are performed on CIFAR-10 [31], CIFAR-100 (20 super-classes) [31] and ImageNet-30 [26]. In Tab. 1, we present a comprehensive comparison of our method with a range of alternatives including one-class classifiers, reconstruction-based methods and SSL methods. Notably, SSL methods using virtual outliers generated by shifting transformations such as rotation and translation, yield favorable results compared to those specifically tailored for one-class learning. Thanks to UniCLR with HA, we achieve enhanced performance across all the three datasets. Moreover, introducing supervision through Outlier Exposure (OE) [25] nearly solves the CIFAR-10 task, which is previously regarded as less effective in the contrastive based AD method [51]. We attribute the success to our contrastive aggregation strategy, which shapes a more focused inlier distribution when more outliers introduced.

Unlabeled Multi-class. This setting expands the one-class dataset to a multi-class scenario, wherein images from different datasets are treated as outliers. In the case of CIFAR-10 as the inlier dataset, we consider SVHN [35], CIFAR-100 [31], ImageNet [34], LSUN [64], ImageNet (Fix), LSUN (Fix) and linearly-interpolated samples of CIFAR-10 (Interp.) [14] as potential outliers. ImageNet (Fix) and LSUN (Fix) are the modified versions of ImageNet and LSUN, designed to address easily detectable artifacts resulting from resizing operations. For ImageNet-30, we consider CUB-200 [55], Dogs [28], Pets [38], Flowers [36], Food-101 [6], Places-365 [68], Caltech256 [20] and DTD [11] as outlier. Tab. 2 shows that our UniCon-HA outperforms other counterparts on most benchmarks. Though the training set follows a multi-center distribution, the straightforward aggregation of all data into a single center proves remarkably effective in AD.

Labeled Multi-class. In the multi-class setting with labeled data, rather than treating all inliers as a single class, as seen in the previous scenarios, we designate inliers sharing identical labels as positives. Conversely, inliers with differing labels or those generated by distributionally-shifted augmentations are negatives. From Tab. 4, by incorporating labels into the UniCLR loss, our method not only improves the performance in unlabeled multi-class setting but also consistently surpasses other competitors that employ virtual outliers, i.e. RotNet [26], GEOM [18], CSI [51] and DROC [49]. It suggests that our method generalizes well to labeled multi-class inliers.

Realistic Dataset. Following DROC [49], we learn patch representations of 32$\times$32. Tab. 3 shows that our method outperforms the counterparts that also incorporate rotation augmentation. Though CutPaste [33] exhibits better performance than ours, it is crucial to understand that CutPaste is specially designed for industrial anomaly localization, making it unsuitable for our settings. For instance, CutPaste only achieves 69.4% while ours reaches 95.4% in the one-class CIFAR-10 scenario.

We conduct ablation studies on (a) various shifting transformations and (b) aggregation strategies: SA and HA.

Shifting Transformation. In contrast to CSI [51], we investigate various shifting transformations beyond rotation, including CutPerm [51], Gaussian blur, Gaussian noise and Sobel filtering. From Tab. 5, our method consistently outperforms CSI under different shifting transformations, with rotation being the most effective. One plausible explanation is that rotation creates more distinguishable negative samples from the original ones, facilitating the learning process. Notably, different from CSI [51] and RotNet [26], we do not learn to differentiate specific transformation types. It encourages the community to rethink the necessity of transformation prediction through a classifier, such as the task of 4-way rotation prediction. Please refer to the supplementary material for further analysis on additional transformations.

Aggregation Strategy. Beyond using the UniCLR loss, we also employ the SA and HA strategies to prompt a more purified and compact concentration of inliers, respectively. To assess the efficacy of each strategy, we establish two baselines: one with a rotation classifier and one without, conducting vanilla contrastive learning on the union of inliers and virtual outliers. The results in Tab. 6 indicate that both single-stage and multi-stage aggregations yield improved outcomes through SA. This underscores the benefits of mitigating the impact of the outliers generated by unexpected data augmentation, thereby purifying the inlier distribution. Tab. 6 reveals two pivotal observations regarding HA: firstly, the presence of UniCLR diminishes the impact of a rotation classifier (2,3,4 vs. 6,8,9), thanks to promoting inlier concentration and outlier dispersion. Secondly, enabling solely $res_{4}$ for contrastive aggregation significantly improves baselines (1 vs. 2, or 5 vs. 6). Broadly, HA leads to a noteworthy and consistent gain when applied across more stages (2 vs. 3, or 6 vs. 8).