SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Generation Test-time generation is achieved in MCMC transition steps from the proposal random (Gaussian) distribution. We use a colder model and a smaller step size at test time which is annealed from 0.1 and decreases by a factor of 0.8 in every 40 steps to encourage stable generation. The diffusion noise magnitude is set to 0.001. Longer MCMC chains with more steps (e.g., 200) are taken to achieve realistic generation. A consistent observation across various considered methods is that the step size stands out as the crucial parameter predominantly influencing the quality of generation.
In practice, MCMC-based generation could be slow, especially in cases with very long MCMC chains. Though in our case, 200 steps are enough to produce very realistic samples, it is worth exploring methods to accelerate the generation process.
Anomaly Detection In contrast to prevalent practices in computer vision, EBM-based OOD detection in LHC physics exhibits specificity. The straightforward approach of comparing two distinct datasets, such as CIFAR10 and SVHN, overlooks the intricate real-world application environments. Adapting OOD detection to scientific discovery at the LHC, we reformulate and tailor the decision process as follows: if we train on multiple known Standard Model jet classes, we focus on class-conditional model evaluation for discriminating between the unseen test signals and the most copious background events (i.e., QCD jets rather than all the Standard Model jets).
4.1 Generative Modeling - Energy-based Event Generator
We present the generated jets transformed from initial random noises with the Langevin dynamics MCMC. Due to the non-human-readable nature of physics events (e.g., low-level raw records at the particle detectors), we are not able to examine the generation quality through formats such as images directly. However, it has a long history that expert-designed high-level observables can serve as strong discriminating features. In the first row of Fig. 2, we first show the distributions of input features for the data and the model generation. Meanwhile, in the second row, we plot the distributions of high-level expert observables including the jet transverse momentum  $p_{\mathrm{T}}$  and the jet mass M. Through modeling low-level features in the detector space, we achieve precise recovery of the high-level physics observables in the theoretical framework. For better visualization, we can also map the jets onto the  $(\eta ,\phi)$  plane, with pixel intensities associated with the corresponding energy deposits. We show the average generated jet images in Fig. 7 of Appendix B, comparing to the real jet images (right-most) in the  $(\eta ,\phi)$  plane. In Table 2, we present the Jensen-Shannon Divergence of the high-level observables  $p_T$  and  $M$  distributions between real data and model generation, as the quantitative measure of the generation performance.
Figure 2: Top: Input feature distributions of jet constituents for the data and the model generation. Bottom: High-level feature distributions for the data and the model generation.
Model	JSD (pT) / 10-4	JSD (M)/ 10-4	JSD (M/pT)/ 10-4
β-VAE [15]	3.7	11.0	-
EBM	2.4 ± 1.9	3.2 ± 2.0	6.3 ± 4.6
Table 2: Model comparison in terms of generation quality measured in Jensen-Shannon Divergence of high-level observables  $p_T$  and  $M$ . For EBMs, we present the means and standard deviations for the JSDs obtained from 5 independent runs. We also present a  $\beta$ -VAE [15] model as a baseline.
4.2 Anomaly Detection - Anomalous Jet Tagging
Since EBMs naturally provide an energy score for each jet, for which the in-distribution samples should have lower scores while OOD samples are expected to incur higher energies. Furthermore, a classifier, when interpreted as an energy-based model, the transformed energy score can also serve as an OOD identifier [33, 49].
EBM In HEP, the in-situ energy score can be used to identify potential new physics signals. With an energy-based model, which is trained on the QCD background events or directly on the slightly signal-contaminated data, we expect unseen signals (i.e., non-QCD jets) to have higher energies and correspondingly lower likelihoods.
In Fig. 3, we compare the energy distributions of in-distribution QCD samples, out-of-distribution signal examples (hypothesized heavy Higgs boson [7] which decays into four QCD sub-jets), and random samples drawn from the proposal Gaussian distribution. We observe that random samples unusually have the highest energies. Signal jets have relatively higher energies compared with the QCD background jets, making model-independent new physics searches possible.
Spurious Correlation A more intriguing property of EBMs is that spurious correlations can be better handled. Spurious correlations in jet tagging might result in distorted background distributions and obscure effective signal detection. For instance, VAEs in OOD detection can be highly correlated with the jet masses [15], similar to the spurious correlation with image pixel numbers in computer vision [50, 60]. In the right panel of Fig. 3, we plot the correlation between the energy score and the jet mass. Unlike other generative strategies for model-independent anomaly detection, EBMs are largely free from the spurious correlation between the energy  $E(\mathbf{x})$  and the jet mass M. The underlying reason for EBMs not presenting mass correlation could be the negative sampling involved in the training process. Larger mass modes are already covered during the negative sampling process. This makes EBMs a promising candidate for model-independent new physics searches.
Figure 3: Left: Energy distributions for random samples, background QCD jets, and novel signals. Right: Correlation between the jet mass  $M_J$  and the energy  $E$ .
EBM-CLF The task of classifying/identifying different Standard Model jet types and the task of searching for beyond the Standard Model signals actually can be unified in a single approach with neural classifiers distinguishing different Standard Model particle types [16]. Compared with naive generative model based OOD, training neural classifiers on different known Standard Model jets helps the model with learning more meaningful and robust representations for effectively detecting new signals. Further combining generative models and discriminative classifiers can be achieved in the EBM framework elegantly. We employ the hybrid learning scheme [54, 33] combining the discriminate and the generative approaches. It links with many OOD detection techniques and observations. For instance, the maximum logit for OOD detection [34] comes in the format of the lowest energy  $\min_y E(\mathbf{x}, y)$  in this framework.
We train an EBM-based multi-class classifier (EBM-CLF) according to Eq. 5, for both discriminating different Standard Model jets (QCD jets, boosted W jets, and boosted top jets) and generating real-like jet samples. The results of generative sampling are shown in Appendix B. The jointly trained EBM and jet classifier maintain the classification accuracy (see Appendix B). The associated EBM is augmented by the discriminative task, and thus assisted with better inductive biases and domain knowledge contained in the in-distribution classes. EBM-CLF is an example of how we unify different physics tasks (jet classification, anomaly detection, and generative modeling) in a unified framework.
We measure the OOD detection performance in the ROC curve and the AUC of the binary classification between the background QCD samples and the signal jets. Table 3 records the AUCs of different models (and different anomaly scoring functions) in tagging Standard Model Top jets and hypothesized Higgs bosons (OOD  $H$ ). The jointly trained model generally has better anomaly tagging performance compared with the naive EBM. We also explore the norm of the score function of  $p_{\theta}(\mathbf{x})$ :  $\| \nabla_{\mathbf{x}}\log p_{\theta}(\mathbf{x})\| = \| \nabla_{\mathbf{x}}E(\mathbf{x})\|$  serving as the anomaly score (similar to the approximate mass in [33]). Constructed from the energy function, the score function approximately inherits the nice property of mass decorrelation. However, they have slightly worse performance compared with  $E(\mathbf{x})$ . The corresponding ROC curves are shown in the left panel of Fig. 4, in terms of the signal efficiency  $\epsilon_S$  (i.e., true positive rate) and the background rejection rate  $1 / \epsilon_B$  (i.e., 1/false positive rate). In the right panel, we plot the background mass distributions under different cuts on the energy scores. We observe excellent jet mass decorrelation/invariance for energy score-based anomalous jet tagging. Additionally for EBM-CLF, we also record the AUCs for anomaly scores of the class-conditional softmax probability  $p(y|\mathbf{x})$  and the logit  $\mathbf{g}(\mathbf{x})_y = -E(\mathbf{x},y)$  corresponding to the background QCD class, in Table 3. However, without further decorrelation strategies (auxiliary tasks or assistant training strategies), these two anomaly scores are usually strongly correlated with the masses of the in-distribution classes and distort the background distributions. Thus we list the results here only for reference.
Model	AUC (Top)	AUC (OOD H)
DisCo-VAE (κ = 1000) [15]	0.593	0.481
KL-OE-VAE [15]	0.744	0.625
Mass-decorrelated
EBM (E(x))	0.679 ± 0.009	0.794 ± 0.032
EBM (||∇x log pθ(x)||)	0.628 ± 0.017	0.738 ± 0.033
EBM-CLF (E(x))	0.703 ± 0.020	0.815 ± 0.018
EBM-CLF (||∇x log pθ(x)||)	0.679 ± 0.040	0.722 ± 0.028
Mass-correlated
EBM-CLF (g(x)y)	0.920 ± 0.002	0.877 ± 0.008
EBM-CLF (p(y|x))	0.940 ± 0.002	0.865 ± 0.012
Table 3: Anomaly detection performance for different models (anomaly scores in parentheses) measured in AUCs. We present the averaged AUCs over 5 independent models with different random seeds and the associated standard deviations. We list a few baseline models with mass decorrelation achieved either by a regularization term (DisCo-VAE ( $\kappa = 1000$ ) or an auxiliary classification task contrasting in-distribution data and outlier data (KL-OE-VAE).
We present a novel energy-based generative framework for modeling the behavior of elementary particles. By mimicking the inter-particle interactions with a self-attention-based transformer, we map the correlations in the detector space to a probabilistic space with an energy function. The energy model is used for the implicit generation of physics events. Despite the difficulty in training EBMs, we adapted the training strategies to balance learning efficiency and training stability. We adopt short-run MCMC sampling at training, while at test time we instead use dynamic generation to obtain prominent generative quality. Additionally, the framework supports the construction of an augmented classifier with integrated generative modeling. This unified approach provides us with flexible tools for high-performance parameterized physics event simulation and spurious-correlation-free model-independent signal detection. It marks a significant stride in next-generation multitasking machine learning models for high-energy physics.
Figure 4: Left: ROC Curves for EBM and EBM-CLF with the energy  $E(\mathbf{x})$  as the anomaly score. The grey line denotes the case of random guessing. Right: Background mass distributions under different acceptance rates  $\epsilon$  after cutting on the energy score from the EBM-CLF.
This work is partially supported by the IVADO Postdoctoral Research Funding. We acknowledge the computing resources provided by Mila (mila.quebec).
[1] LHC Machine. JINST, 3:S08001, 2008.
[2] Georges Aad et al. Observation of a new particle in the search for the Standard Model Higgs boson with the ATLAS detector at the LHC. Phys. Lett. B, 716:1-29, 2012.
[3] David H. Ackley, Geoffrey E. Hinton, and Terrence J. Sejnowski. A learning algorithm for boltzmann machines. Cognitive Science, 9(1):147-169, 1985.
[4] Johan Alwall, Michel Herquet, Fabio Maltoni, Olivier Mattelaer, and Tim Stelzer. Madgraph 5: going beyond. Journal of High Energy Physics, 2011(6), Jun 2011.
[5] Oliver Atkinson, Akanksha Bhardwaj, Christoph Englert, Partha Konar, Vishal S. Ngairangbam, and Michael Spannowsky. IRC-Safe Graph Autoencoder for Unsupervised Anomaly Detection. Front. Artif. Intell., 5:943135, 2022.
[6] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. CoRR, abs/1409.0473, 2015.
[7] G.C. Branco, P.M. Ferreira, L. Lavoura, M.N. Rebelo, Marc Sher, and Joao P. Silva. Theory and phenomenology of two-Higgs-doublet models. Phys. Rept., 516:1-102, 2012.
[8] Anja Butter, Tilman Plehn, and Ramon Winterhalder. How to GAN LHC Events. SciPost Phys., 7(6):075, 2019.
[9] Matteo Cacciari, Gavin P Salam, and Gregory Soyez. The anti-kt jet clustering algorithm. Journal of High Energy Physics, 2008(04):063–063, Apr 2008.
[10] J. M. Campbell et al. Event Generators for High-Energy Physics Experiments. In 2022 Snowmass Summer Study, 3 2022.
[11] Miguel Á. Carreira-Perpínán and Geoffrey Hinton. On contrastive divergence learning. In Robert G. Cowell and Zoubin Ghahramani, editors, Proceedings of the Tenth International Workshop on Artificial Intelligence and Statistics, volume R5 of Proceedings of Machine Learning Research, pages 33–40. PMLR, 06–08 Jan 2005. Reissued by PMLR on 30 March 2021.
[58] Michela Paganini, Luke de Oliveira, and Benjamin Nachman. Accelerating Science with Generative Adversarial Networks: An Application to 3D Particle Showers in Multilayer Calorimeters. Phys. Rev. Lett., 120(4):042003, 2018.
[59] Huilin Qu and Loukas Gouskos. ParticleNet: Jet Tagging via Particle Clouds. Phys. Rev. D, 101(5):056019, 2020.
[60] J. Ren, Peter J. Liu, Emily Fertig, Jasper Snoek, Ryan Poplin, Mark A. DePristo, Joshua V. Dillon, and Balaji Lakshminarayanan. Likelihood ratios for out-of-distribution detection. *ArXiv*, abs/1906.02845, 2019.
[61] Torbjörn Sjöstrand, Stephen Mrenna, and Peter Skands. A brief introduction to pythia 8.1. Computer Physics Communications, 178(11):852-867, Jun 2008.
[62] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 12438-12448. Curran Associates, Inc., 2020.
[63] Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In Proceedings of the 25th International Conference on Machine Learning, ICML '08, page 1064-1071, New York, NY, USA, 2008. Association for Computing Machinery.
[64] Mary Touranakou, Nadezda Chernyavskaya, Javier Duarte, Dimitrios Gunopulos, Raghav Kansal, Brento Orzari, Maurizio Pierini, Thiago Tomei, and Jean-Roch Vlimant. Particle-based fast jet simulation at the lhc with variational autoencoders, 2022.
[65] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. ArXiv, abs/1706.03762, 2017.
[66] Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28th International Conference on International Conference on Machine Learning, ICML'11, page 681-688, Madison, WI, USA, 2011. Omnipress.
[67] Jianwen Xie, Yang Lu, Song-Chun Zhu, and Yingnian Wu. A theory of generative convnet. In International Conference on Machine Learning, pages 2635-2644. PMLR, 2016.
[68] Jianwen Xie, Yifei Xu, Zilong Zheng, Song-Chun Zhu, and Ying Nian Wu. Generative pointnet: Deep energy-based learning on unordered point sets for 3d generation, reconstruction and classification. 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 14971-14980, 2021.
[69] Sangwoong Yoon, Yung-Kyun Noh, and Frank Chongwoo Park. Autoencoding under normalization constraints. In International Conference on Machine Learning, 2021.
[70] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póczos, Ruslan Salakhutdinov, and Alex Smola. Deep sets. In NIPS, 2017.
[71] Shuangfei Zhai, Yu Cheng, Weining Lu, and Zhongfei Zhang. Deep structured energy based models for anomaly detection. In International Conference on Machine Learning, 2016.
Datasets For the standard EBM, we train on 300,000 simulated QCD jets. For the hybrid model EBM-CLF, we train on 300,000 simulated Standard Model jets (100,000 QCD jets, 100,000 boosted jets originating from the W boson, and 100,000 boosted jets originating from the top quark). For OOD detection test sets, we employ the hypothetical Higgs boson (in the decay mode of  $H \rightarrow hh \rightarrow (bb)(bb)$ ) with a mass of 174 GeV, which decays into two lighter Higgs bosons of 80 GeV. The intermediate light Higgs boson decays into two b quarks. Each test set consists of 10,000 samples that are  $p_T$ -refined to the range of [550, 650] GeV. All the jet samples are generated with a pipeline of physics simulators.
Event Generation QCD jets are extracted from QCD di-jet events that are generated with MadGraph [4] for LHC 13 TeV, followed by Pythia8 [61] and Delphes [18] for parton shower and fast detector simulation. All jets are clustered using the anti- $k_{T}$  algorithm [9] with cone size  $R = 1.0$  and a selection cut in the jet transverse momentum  $p_{T} > 450\mathrm{GeV}$ . We use the particle flow objects for jet clustering.
Input Preprocessing Jets are preprocessed before being fed into the neural models. Jets are longitudinally boosted and centered at  $(0,0)$  in the  $(\eta ,\phi)$  plane. The centered jets are then rotated so that the jet principal axis  $\left(\sum_{i}\frac{\eta_{i}E_{i}}{R_{i}},\sum_{i}\frac{\phi_{i}E_{i}}{R_{i}}\right)$  (with  $R_{i} = \sqrt{\eta_{i}^{2} + \phi_{i}^{2}}$  ) and  $E_{i}$  is the constituent energy) is vertically aligned on the  $(\eta ,\phi)$  plane.
Hyper-parameters Hyper-parameters are recorded in Table 4. Hyper-parameters for the transformer are chosen according to a jet classification task. We scan over the following hyper-parameter ranges:
$$
\text {s t e p} \quad \text {s i z e} \in \{0. 0 1, 0. 1, 1. 0, 1 0 \}
$$
$$
s t e p s \in \{5, 1 0, 2 4, 6 0 \}
$$
$$
1 \mathbf {r} \in \{1 \mathrm {e} - 3, 1 \mathrm {e} - 4 \}
$$
We found that fewer steps  $\{5, 10\}$  make training unstable. Thus the model prefers a relatively larger number of steps and a smaller step size. However, 60 steps MCMC takes much longer time to train and no significant improvement was observed for even longer chains. We balance training efficiency and effective learning by choosing 24 steps.
The MCMC chains are initialized with Gaussian noises, where the constituent features are sampled from the following distributions:  $\log p_{Ti} \sim \mathcal{N}(2,1)$ ,  $\eta_i \sim \mathcal{N}(0,0.1^2)$ , and  $\phi_i \sim \mathcal{N}(0,0.2^2)$ .
Data
input features	{ (log(pT), η, φ)i}N i=1
input length	N=40 with zero-padding
Energy Function
Number of layers NL	8
Model dimension dmodel	128
Number of heads	16
Feed-forward dimension	1024
Dropout rate	0.1
Normalization	None
MCMC
Number of steps	24
Step size	0.1
Buffer size	10000
Resample rate	0.05
Noise	ε = 0.005
Regularization
L2 Regularization	0.1
Training
Batch size	128
Optimizer	Adam (β1 = 0.0, β2 = 0.999)
Learning rate	1e-4 (decay rate γ = 0.98)
Table 4: Model settings.
We explore the most crucial aspects of the model to test their functionality:
- With an energy function approximated with a Multi-Layer-Perceptron (MLP) net, we were not able to achieve quality generation.
- We also tried out the Hamiltonian Monte Carlo (HMC) [52] for the MCMC procedure. However, we were not able to achieve good performance in these experiments.
Results measured in the Jensen-Shannon Divergence of the high-level observables  $(p_T$  and  $M)$  are recorded in Table 5.
Ablation	JSD (PT) / 10-4	JSD (M) / 10-4
MLP	Energy Function
Fig. 5	Fig. 5
MCMC Dynamics
HMC	24	30
Table 5: Ablation study on different components of the model, training strategies, and training techniques. Since the MLP-based model is not able to produce high-quality samples, we instead show the observable distributions (Fig. 5) to visually show the failure patterns.
Figure 5: Typical high-level observable distributions for MLP-based models.
B.2 Classification Performance of EBM-CLF
The classification accuracies for EBM-CLF (QCD/W/Top 3-way classification) are recorded in Table 6. EBM-CLF performs on par with the fully supervised classifier ParticleNet [59], while EBM-CLF is trained on a much smaller dataset. The corresponding confusion matrices are displayed in Fig. 6. When we empirically down-weight the term  $\log p(y|\mathbf{x})$  (decrease  $\kappa$  in Eq. 7), the classification performance drops correspondingly.
$$
\log p (\mathbf {x}, y) = \log p (\mathbf {x}) + \kappa \log p (y | \mathbf {x}). \tag {7}
$$
B.3 Additional Plots
In Fig. 7, we show the generated jet samples displayed in images on the  $(\eta, \phi)$  plane. In Fig. 8, we show the high-level observable distributions of generated jet samples of EBM-CLF. In Fig. 9, we show the background mass distributions under different acceptance rates  $\epsilon$  after cutting on the energy score of the standard EBM.
Model	Top-1 Accuracy	Top-2 Accuracy
ParticleNet[59]	0.871	0.976
EBM-CLF (κ = 1.0)	0.850	0.969
EBM-CLF (κ = 0.5)	0.708	0.906
EBM-CLF (κ = 0.1)	0.679	0.852
Table 6: Classification performance of EBM-CLF on QCD/W/Top 3-way classification.  $\kappa$  denotes the weight of the discriminative log-likelihood.
Figure 6: Confusion matrices for EBM-CLF(left) and ParticleNet(right).
Figure 7: Jet images averaged over 10000 jet samples. From left to right, we show the initial random noises (left-most), EBM-generated jet samples by the MCMC chains in intervals (middle), and Real jets (right-most).
To ensure reproducibility and encourage open-source practice in the HEP community, we release the code implementation in https://github.com/taolicheng/EBM-HEP. The training sets and test sets are accessible at [47, 14]. Due to difficulties in aligning model comparison protocols for different research groups, we thus focus on methods with code publicly available, that serve as credible baselines, for model comparison.
Figure 8: High-level observable distributions for the generated samples of EBM-CLF (orange) and the data (blue).
Figure 9: Background mass distributions under different acceptance rates  $\epsilon$  after cutting on the energy score from the EBM.