Autoencoders for unsupervised anomaly detection in high energy physics

Throughout this work, we use top tagging as an example, which is a well established benchmark in collider physics: Simulated QCD jets (initiated by light-flavor quarks or gluons) have to be distinguished from hadronically decaying boosted top-quarks which form single jets. We use the benchmark dataset from Ref. Butter_2018 , which is publicly available at kasieczka_gregor_2019_2603256 . A subset of 100k and 40k jets of one of the two classes is used for training and validation, respectively, and two 40k subsets of each type for testing.

In the following, we describe our default setup which will be modified and improved in Section 3. We employ the preprocessing introduced in Ref. Macaluso:2018tck . The four momenta of the jet constituents are converted into two-dimensional jet images with $40\times 40$ pixels. The details of the jet generation and the preprocessing are discussed in Appendix A.

Fig. 1 shows the average images of both classes. For QCD jets, most of the intensity is concentrated in the central pixels. For top jets, there is a clearly visible three-prong structure (as expected for top-quark decays after preprocessing). Of course, individual jets are harder to distinguish than their average images may indicate.

As our anomaly detection algorithm, we use a convolutional autoencoder with an architecture similar to the one in Ref. Heimel_2019 . We implement our AE with Tensorflow 2.4.1 tensorflow2015-whitepaper , relying on the built in version of Keras chollet2015keras . Several convolution layers with $4\times 4$ kernel and average pooling layers with $2\times 2$ kernel are applied before the image is flattened and a fully connected network reduces the input further into the bottleneck latent space with 32 nodes. The Parametric ReLU activation function is used in all layers. The described encoder structure is inverted to form the corresponding decoder which is used to reconstruct the original image from its latent space description. Our architecture is defined in Fig. 2; the hyperparameter settings are described in more detail in Appendix A.

Following Ref. Heimel_2019 , to evaluate the reconstruction of the input picture we use the mean squared error (MSE), i.e. the average of the squared error of each reconstructed pixel with respect to its input value, as a loss function. During testing the value of the loss function is also used as the discriminator between signal and background. An event is tagged as signal/anomaly if the value of the loss function is larger than a given threshold. Changing the threshold value, one obtains the usual receiver operating characteristic (ROC) curve.

We first investigate what the AE is actually learning, as it is trained for the reconstruction of the input and not as an anomaly tagger. Fig. 3 shows the learning history of an exemplary top jet in terms of its reconstruction after training on top jets for a given number of epochs. To guide the eye, we also show the evolution of the squared error per pixel and highlight the reconstruction of the intensity of the brightest pixels in detail. Moreover, the MSE of the reconstruction, i.e. the loss function of the AE, is given as a measure for the overall reconstruction improvement.

During the first epoch, the AE learns to reconstruct the average top jet image, as can be seen by comparing the reconstructed image to the average image in Fig. 1. (Note that the images in Fig. 1 are shown on a logarithmic scale.) In the following epochs, the squared error is dominated by the brightest pixels. The AE improves the loss by improving their reconstruction, so its trainable weights are updated accordingly. After 10 epochs, the AE recovers a smeared reconstruction of the brightest pixels. After roughly 25 epochs, the weights are learned to reconstruct the brightest pixels so well that the corresponding error becomes small compared to the error of the remaining pixels. However, the AE further improves the reconstruction of the brightest pixels (and the surrounding zero-intensity pixels) instead of changing its focus to the dimmer pixels. The latter would probably harm the previously learned reconstruction of the brightest pixels and therefore increase the loss. After around 100 epochs the four brightest pixels are very well reconstructed, while most of the dimmer pixels are still ignored. The dimmer pixels dominate the total error, as can be seen in the right column of Fig. 3. The AE is apparently trapped in a local minimum of the loss function, and training longer does not change the picture significantly.

The previous discussion applies to both training on QCD and top images. To show that the example presented in Fig. 3 is rather generic, the reconstruction capabilities of the AE for the whole test dataset are summarized in Fig. 4, where we quantify the quality of the reconstruction for individual pixels of a jet image. Specifically, for each of the 40k jets of the test data sample, we determine the ratio of the reconstructed and input intensities, $r$, for each pixel, and show the fraction of test jet images where the brightest, next-to-brightest etc. pixel (from left to right on the horizontal axis) is reconstructed with a certain quality $r$. In Fig. 4 (left) this fraction is shown for the AE trained and tested on QCD jets. For the majority of the jet images, the brightest pixels are reconstructed well (blue histogram, corresponding to a ratio $80\%\leq r\leq 120\%$). On the other hand, the dimmer a pixel the more likely it will be reconstructed insufficiently, e.g. with an intensity of less than 10% of the input intensity (red histogram). Note that the overall number of jets in the training data, and thus the fraction of jet images, decreases as one requires an increasing number of non-zero pixels. A qualitatively similar picture emerges for the AE trained and tested on top jets (see Fig. 4, right). In the next section, we will explore why such a limited AE can nevertheless tag top jets as anomalies in a QCD background.

We can gain additional insight into the AE performance by investigating the reconstruction loss as a function of the number of non-zero pixels, $N_{p}$, of the jet images. We show the $N_{p}$ distribution for QCD and top images in Fig. 5. When training on QCD images, the reconstruction loss of QCD images on average increases strongly with $N_{p}$, as shown in Fig. 6, left. QCD jets with small $N_{p}$ are simply easier to reconstruct. When training on top images, top jets are also harder to reconstruct for larger $N_{p}$, see Fig. 6, right. However, in this case the correlation is less pronounced. Since, on average, top images have more non-zero pixels than QCD images, see Fig. 5, the top-trained AE sees less training data with small $N_{p}$. Given this correlation between the reconstruction loss and the number of non-zero pixels $N_{p}$, we conclude that $N_{p}$ at least partially describes the complexity of the images.

However, complexity is not only related to $N_{p}$. For fixed $N_{p}$, QCD jets are on average also easier to reconstruct than top jets. This can be understood from the underlying physics. A top jet is naturally composed of three sub-jets initiated by the top-quark decay products. Hence, this intrinsic three-prong structure leads to more complex structures in the jet images.

These results show that the AE has a strong complexity bias: images which would intuitively be labeled as simpler, are reconstructed better. It has been discussed in the literature, especially for natural images, that simpler images may be harder to identify as anomalous nalisnick2019deep ; zong2018deep ; gong2019memorizing ; ren2019likelihood ; serra2019input ; schirrmeister2020understanding ; kirichenko2020normalizing ; tong2020fixing . In the context of natural images, it was noted that the algorithm tends to learn features that are not representative of the specific training set, like local pixel correlations.

The strong complexity bias in our application might be mainly due to the limited reconstruction performance of the AE discussed in Section 2.2. It is not surprising that the AE is not able to reconstruct the complex structure of top images when ignoring all but a few pixels. Limitations to the reconstruction of structures in high energy physics have also been noted recently in batson2021topological ; Collins:2021nxn .

The correlation between the complexity of an image and the reconstruction loss is dangerous if the AE is to be used as a model-independent anomaly tagger. One should not assume that a new physics signal looks more complex than the SM background. An anomaly tagger should only reconstruct well those images it has been trained on. On the other hand, our AE has the tendency to better reconstruct simple images. For a functional model-independent anomaly tagger the first effect should be much more pronounced than the second. For the case at hand, these two effects go in the same direction for the direct tagger, since it is trained on simple QCD images, but go in opposite directions for the inverse tagger, which is trained on more complex top images. This can be seen already in Fig. 6, where we also show the loss of top jets in the test set for the AE trained on QCD jets and vice versa. Indeed, when training on top jets, even QCD images with the same number of non-zero pixels are better reconstructed than top jets, due to the complexity issue. Although the inverse tagger has improved the top reconstruction through its training on top jet images, it cannot overcome the complexity bias. Given the limited reconstruction capabilities, after learning to reproduce top images, the inverse tagger has also learned to interpolate QCD images. On the other hand, the direct tagger is never exposed to the complex features characterizing top images during training and thus fails to extrapolate to the corresponding top jet structures.

Having investigated the limitations in the AE reconstruction power and the complexity bias, we show the performance of the direct and inverse taggers in Fig. 7. The ROC curves are constructed by varying the loss threshold for tagging a jet as anomalous. We also show the ROC curve for a supervised CNN tagger inspired by Ref. Kasieczka_2019 (see Appendix A). Our direct AE tagger reproduces the results for the direct unsupervised AE tagger in Ref. Heimel_2019 , which has used a similar setup. As expected, it performs worse than a supervised tagger, but much better than random guessing. A similar performance for the unsupervised AE tagger has also been obtained in Ref. Farina_2020 in a slightly different setup. Hence, our direct tagger works as expected.

As can be seen from Fig. 7, the inverse tagger performs worse than randomly tagging jets as anomalous. Our inverse autoencoder fails to tag an anomaly (QCD jets in this case) that is simpler than the background. Through its training on top jets, the inverse tagger learns to reconstruct top jets better than the direct tagger (see right plot in Fig. 7), and its performance for the reconstruction of QCD jets is diminished. However, this is not enough to overcome the complexity bias.

To summarize, even with a limited reconstruction capability, an autoencoder can be a good anomaly tagger if there is a bias to reconstruct the background better. However, if the bias works against anomaly detection, the learning capabilities of the AE may not be sufficient to overcome the bias. Only a powerful AE with a background specific data compression in the latent space could potentially be able to overcome such a bias. Improvements of our setup to partly achieve these goals are discussed in Section 3.

It should be noted that a perfect AE, which is always able to perfectly reconstruct the input via the identity mapping regardless of the input data, would be useless as a tagger as it would always interpolate perfectly from the learned data to the anomalies.

As shown in Section 2.2, the AE focuses on reconstructing only the brightest pixels of an image. The contrast between bright and dim pixels is apparently so large that the dim pixels are not recognized during the early learning epochs. To tackle this problem, we propose several remappings of the pixel intensities, such that the intensities of the dim pixels are enhanced compared to those of the bright pixels. Such a remapping should encourage the AE to ignore less pixels and learn more of the intrinsic structure and features of the given training images. Hence, in particular the complexity of the simplest pictures (having only one or two very bright pixels) is increased, and we expect to also reduce the complexity bias. The intensity remapping can be considered as part of an alternative preprocessing.

Due to the normalization of the input features in the standard preprocessing described in Appendix A, all pixel intensities lie in the interval between 0 and 1. To highlight dim pixels by remapping this interval onto itself, we use the functions

which are displayed in Fig. 17 in Appendix B. For the mapping $\text{R}_{3}$ we use $\alpha=1000$; the $\Theta$-function maps each non-zero pixel to one. The original images correspond to the identity remapping $\text{R}_{0}(x)=x$ and are labeled accordingly in the following. After the remapping the image is normalized again to have a pixel sum equal to one. While remappings $\text{R}_{1}$ to $\text{R}_{3}$ are invertible, information is lost through $\text{R}_{4}$, which should be thought of as a limiting case. The effects of the remappings are exemplified and quantified further in Appendix B.

As discussed in Section 2, it is difficult for the AE to learn the structure of an image. Instead, it only learns the exact position of the few brightest pixels. This is actually not surprising, since the AE is not rewarded by the loss function to reconstruct the image structure. It rather has to reproduce the exact pixel position, which seems to be an unnecessarily tough task. For example, shifting the input image by one pixel would be considered by a human as a valid reconstruction, but is strongly penalized for the standard MSE loss.

In fact, the similarity of two images is only poorly represented by the MSE loss. To compare two distributions, i.e. the intensity distribution in our case, the Wasserstein distance, also called energy or earth mover’s distance 10.1023/A:1026543900054 , would be a much better choice. However, its full calculation includes an optimization problem that would need to be solved for each image during training. Approximate estimates of this distance, e.g. the sliced Wasserstein distance bonneel2014sliced , are also computationally expensive. We thus propose a simple loss function which nevertheless provides a better notion of similarity between two pictures than the MSE loss.

To introduce a notion of neighborhood for each pixel in the loss function we convolve the whole image with a smearing kernel. Hence, neighboring pixels have a partly overlapping smeared distribution which can be recognized by the loss function. The optimal reconstruction, i.e. vanishing loss, is still achieved by reconstructing the original (unsmeared) input image.

On the discrete set of pixels in each image the convolution is defined as

where $I_{kl}$ denotes the intensity of the pixel with coordinates $(k,l)$ of the original image, and $S_{ij}$ refers to the smeared pixels. The kernel is defined as

with the additional constraint that it is set to zero for negative values. Hence, it is only non-zero for an Euclidean distance to the center pixel which is smaller than $L$ pixels and has an approximately circular shape. We choose $L=8$. The normalization is chosen such that the intensity of the central pixel is unchanged. To avoid boundary effects, the original picture is padded with zeros on each side and becomes $54\times 54$ pixels. There is an infinite amount of possible choices for this kernel. However, we do not expect the following results to depend on its details. The reconstruction loss is defined as the MSE of the smeared input and the reconstructed image. We refer to it as kernel MSE or KMSE in the following. Since the smearing is a standard matrix convolution its computational cost is negligible for training and testing.

To evaluate the effect of the intensity remapping on the autoencoder performance, we again investigate the distribution of the ratio $r$ of the reconstructed intensity and the input intensity for the leading pixels. The corresponding histogram, i.e. the equivalent of Fig. 4, is shown for the $\text{R}_{2}$-remapping in Fig. 8. When training on QCD jets (left figure), the fraction of well-reconstructed pixels is increased by the remapping. This suggests that the AE indeed learns more of the jet features. When training on top jets (right figure), the fraction of well-reconstructed jets is not increased by much, but it extends to more of the brighter pixels, also signaling that more features are potentially accessible. The decrease of performance beyond the 20th pixel (and in particular the disappearance of the peak) compared to the performance without remapping, Fig. 4 (right), is not discouraging, since these pixels correspond to very soft particles.¹¹1They are only accidentally well reconstructed for $\text{R}_{0}$ by a tiny non-vanishing intensity which is typically found in the reconstructed image in the vicinity of bright pixels. Since the $\text{R}_{2}$-remapping highlights their intensity, this mechanism does not work any more.

In Fig. 9, we compare the reconstruction quality for the different remappings. The distribution of the intensity ratio $r$ (displayed in Figs. 4 and 8 by stacked histograms) is represented by its median and the 25%/75% quantiles. As long as the distribution is dominated by values of $r$ close to one, the leading pixels are reconstructed well. For both the QCD and the top case, the median curves shift to the right for the remappings, i.e. more of the leading pixels are on average taken into account by the AE. Hence, the AE ignores less pixels and thus learns more features of the jet images. For the AE trained and tested on top jets, right figure, this is only achieved by reducing the reconstruction precision of the brightest pixels.

In Fig. 10, we show the reconstruction performance for training on $\text{R}_{0}$ images using the KMSE loss introduced in Section 3.2. Note that the KMSE loss function is only used during training. To evaluate the pixel reconstruction in the plots no smearing kernel is applied. As one expects, the exact intensity reconstruction of the leading pixels is traded for paying more attention to dimmer pixels. The strong focus to reduce the squared error of the brightest pixels as much as possible is removed. The same can be seen in Fig. 11 where we again show the distributions of the intensity ratio $r$ in terms of the median and quantiles for the different remappings. Although the brightest pixels are not very precisely reconstructed, dim pixels often have a significant part of their intensity reconstructed. Also when training on top jets using the KMSE loss functions, the remappings lead to more pixels being taken into account.

The stacked histograms for all remappings with and without KMSE loss function, which contain additional information on the $r$-distribution, are displayed in Appendix C.

The reconstruction of top jets (when training on top jets) with the $\text{R}_{0}$ and $\text{R}_{2}$ remappings using the standard MSE and the KMSE loss functions is illustrated in Fig. 12. Using the KMSE loss, the AE focuses more on the reconstruction of a continuous distribution of the intensity than on the reconstruction of individual pixels. Even relatively dim regions far from the center receive some attention. Thus, as also shown in Fig. 11, the dim pixels are partially reconstructed as part of this continuous distribution. Therefore, we expect the KMSE autoencoder to extract more useful features of the overall jet structure compared to the standard AE that focuses almost exclusively on the brightest pixels.

To understand the implications of the intensity remapping with respect to the complexity bias, we show the loss of $\text{R}_{2}$ images as a function of the number of non-zero pixels $N_{p}$ in Fig. 13 for the AE trained on QCD (left figure) or top images (right figure) using the MSE loss function. Comparing to Fig. 6, the bias for the QCD-trained AE is reduced a lot and even reversed in direction. For the AE trained on top jets, the bias is also reversed, but additionally increased. This can be understood since the inverse AE sees only few of the training images with small $N_{p}$, the complexity of which has been increased by the intensity remapping.

Training on QCD images, the AE still learns to reconstruct images with low pixel number rather efficiently. Looking at the reconstruction of top images by the direct tagger in the same plot, it still cannot extrapolate well to this unseen data, so that the direct tagger will perform well (see next section). Training on top images with $\text{R}_{2}$ remapping, the AE has lost its ability to simply interpolate the QCD images with small $N_{p}$, because their complexity has increased w.r.t. the $\text{R}_{0}$ case. Hence, this AE is a working inverse tagger. However, comparing the average loss for top and QCD jets in the right panel of Fig. 13, only in a region below 40 active pixels, top images are on average reconstructed better than QCD images. Moreover, the difference is small compared to the width of the distribution and the slope of the bias. Most of the inverse tagging performance is based on the QCD images with small $N_{p}$, due to the reversed bias and still not due to a well-performing AE.

Using the KMSE loss, Fig. 14 shows that the loss as a function of $N_{p}$ for the direct tagger remains very similar to what we have seen in the MSE case. However, we observe that QCD jets with low $N_{p}$ are not automatically learned to be reconstructed well when training the inverse tagger on top images. Moreover, for $N_{p}<35$ top images are again reconstructed on average better than QCD images. Using both a remapping and the KMSE loss function for training, the behavior gets more pronounced. The corresponding distributions are shown in Fig. 19 in Appendix C for all combinations of remappings and loss functions.

The performance for the direct and inverse taggers using the different intensity remappings introduced in Section 3.1 are summarized in terms of their ROC curves in Fig. 15. Derived performance measures are shown in Table 1 in Appendix C. For the direct tagger, QCD jets are background ($\epsilon_{B}=\epsilon_{\rm QCD}$) and top jets are signal ($\epsilon_{S}=\epsilon_{\rm top}$), so that the area under the curve plotted in Fig. 15 directly corresponds to the AUC. For the inverse tagger, signal and background interchange their role ($\epsilon_{S}=\epsilon_{\rm QCD}$ and $\epsilon_{B}=\epsilon_{\rm top}$) and the area under the plotted curves should be as small as possible, since it corresponds to 1-AUC.

Plotting the ROC curves in this way, the area between the ROC curves of the direct and the inverse tagger, called ABC in the following, can be interpreted as the background specific learning achieved by the AE. Furthermore, the difference between the two AUC values for the direct and the inverse tagger, called $\Delta{\rm AUC}$ in the following, is a measure for how biased the AE is. For an AE with small ABC or large $\Delta{\rm AUC}$ one cannot expect a model-independent tagging capability. On the other hand, a large ABC and a small $\Delta$AUC do not guarantee model-independence. It is only an indication that a step towards model-independence has been made.

Using the standard preprocessing without intensity remapping ($\text{R}_{0}$), we find a small ${\rm ABC}=0.28$ and a large $\Delta{\rm AUC}=0.55$. As discussed in Section 2.4, the direct tagger only works due to a large complexity bias, and the inverse tagger fails.

For $\text{R}_{1}$-remapping we see a minor decrease in the direct tagging performance, but a significant improvement in the inverse tagging performance with ${\rm ABC}=0.48$. This improvement is strong enough to enable inverse tagging although $\Delta{\rm AUC}=0.35$ is still large. In contrast, $\text{R}_{2}$-remapping shows little bias ($\Delta{\rm AUC}=0.06$) and an ${\rm ABC}=0.43$. Both direct and inverse tagging have a modest AUC close to 0.7. $\text{R}_{3}$ has a little more bias ($\Delta{\rm AUC}=0.18$) but the best ${\rm ABC=0.50}$. $\text{R}_{4}$-remapping shows that too much highlighting of dim pixels is also counterproductive. It leads to $\Delta{\rm AUC}=-0.3$ and a very poor ${\rm ABC}=0.11$.

Although these results are a promising first step towards a more model-independent anomaly tagging, we want to recall that the inverse tagging is mainly caused by the correlation of the AE loss with the number of non-zero pixels $N_{p}$. The AE is still not able to reconstruct a top jet on which it has been trained better than a QCD jet if both have the same $N_{p}$. Moreover, in particular the tagging performance of the inverse taggers is still very poor. To illustrate this, we also show the two taggers which only use the distribution of $N_{p}$ for both classes, as shown in Fig. 5, to tag anomalies with either a large or a small $N_{p}$. None of our inverse taggers can beat the performance of the trivial tagger looking for few non-zero pixels. This shows once more that the inverse tagging for the intensity remappings mostly relies on the correlation of the loss with the number of pixels.

Fig. 16 shows the ROC curves for direct and inverse taggers using the KMSE loss and the different intensity remappings (see also Table 2 in Appendix C for additional performance measures). The KMSE loss is also used as the anomaly score. On the one hand, the direct $\text{R}_{0}$-tagger using the KMSE loss performs similarly to the MSE one. On the other hand, also the KMSE autoencoder is not able to provide a working inverse tagger, although the performance of the inverse tagger is significantly better than the AE with MSE loss. With ${\rm ABC}=0.41$ and $\Delta{\rm AUC}=0.44$ it is less biased towards reconstructing QCD better.

For the intensity remappings, we find ${\rm ABC}=0.52$ and $\Delta{\rm AUC}=0.27$ for $\text{R}_{1}$, hence the KMSE loss improves the overall performance. The remapping $\text{R}_{2}$ results in ${\rm ABC}=0.27$ and $\Delta{\rm AUC}=-0.09$, i.e. the bias is still small but the performance is reduced. The remapping $\text{R}_{3}$ shows almost no bias ($\Delta{\rm AUC}=0.02$) but also reduced ${\rm ABC}=0.32$. The remapping $\text{R}_{4}$ still fails completely. In contrast to the MSE loss, using KMSE the inverse taggers based on $\text{R}_{1}$ to $\text{R}_{4}$ all show a performance very close to the pixel number tagger. While we have investigated the ability of the KMSE-based AE to reconstruct more jet features in Section 3.3, the tagging performance shows that these features are not necessarily correlated with only one of the two classes.