SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Taoli Cheng
University of Montreal
chengtaoli.1990@gmail.com
aaron.courville@umontreal.ca
As a classical generative modeling approach, energy-based models have the natural advantage of flexibility in the form of the energy function. Recently, energy-based models have achieved great success in modeling high-dimensional data in computer vision and natural language processing. In line with these advancements, we build a multi-purpose energy-based probabilistic model for High Energy Physics events at the Large Hadron Collider. This framework builds on a powerful generative model and describes higher-order inter-particle interactions. It suits different encoding architectures and builds on implicit generation. As for applicative aspects, it can serve as a powerful parameterized event generator for physics simulation, a generic anomalous signal detector free from spurious correlations, and an augmented event classifier for particle identification.
The Large Hadron Collider (LHC) [1], the most energetic particle collider in human history, collides highly energetic protons to examine the underlying physics of subatomic particles and extend our current understanding of the fundamental forces of nature, summarized by the Standard Model of particle physics. After the great success in observing the Higgs boson [2, 28], the most critical task of searching for new physics signals remains challenging. High Energy Physics (HEP) events produced at the LHC have the properties of high dimensionality, high complexity, and enormous data size. To detect rare signals from enormous background events, physics observables describing these patterns have been used to identify different types of radiation patterns. However, it's not possible to utilize all the information recorded by the detectors with a few expert-designed features. Thus deep neural classifiers and generative models, which can easily process high-dimensional data meet the needs for more precise data analysis and signal detection.
Energy-based Models (EBMs) [38, 3, 46], as a classical generative framework, leverage the energy function for learning dependencies between input variables. With an energy function  $E(\mathbf{x})$  and constructing the un-normalized probabilities through the exponential  $\tilde{p}(\mathbf{x}) = \exp(-E(\mathbf{x}))$ , the energy model naturally yields a probability distribution. Despite the flexibility in the modeling, the training of EBMs has been cumbersome and unstable due to the intractable partition function and the corresponding Monte Carlo sampling involved. More recently, EBMs have been succeeding in high-dimensional modeling [55, 56, 25, 23, 62, 19, 51] for computer vision and natural language processing. At the same time, it has been revealed that neural classifiers are naturally connected with EBMs [67, 33, 32], combining the discriminative and generative learning processes in a common learning regime. More interestingly, compositionality [30] can be easily incorporated within the framework of EBMs by simply summing up the energy functions [21, 22]. On the other hand, statistical physics originally inspired the invention of EBMs. This natural connection in formalism makes EBMs appealing in modeling physical systems. In physical sciences, EBMs have been used to
Topic	Practice
Generative modeling	Parameterized event generation
Out-of-distribution detection	Model-independent new physics search
Hybrid modeling	Classifier combined with EBMs
Table 1: Application aspects of Energy-based Models for High Energy Physics.
simulate condensed-matter systems and protein molecules [57]. They have also been shown great potential in structure biology [24], in a use case of protein conformation.
Motivated by the flexibility in the architecture and the compatibility with different tasks, we explore the potential of EBMs in modeling radiation patterns of elementary particles at high energy. The energy function is flexible enough to incorporate sophisticated architectures. Thus EBMs provide a convenient mechanism to simulate complex correlations in input features or high-order interactions between particles. Aside from image generation, applications for point cloud data [68], graph neural networks [48] for molecule generation are also explored. In particle physics, we leverage the self-attention mechanism [6, 65], to mimic the complex interactions between elementary particles.
As one important practical application, neural net-based unsupervised learning of physics events [58, 42, 8, 64] have been explored in the usual generative modeling methods including Generative Adversarial Networks (GANs) [31] and Variational Autoencoders (VAEs) [45]. However, GANs employ separate networks, which need to be carefully tuned, for the generation process. They usually suffer from unstable training, high computation demands, and mode collapse. In comparison, VAEs need a well-designed reconstruction loss, which could be difficult for sophisticated network architectures and complex input features. EBMs thus provide a strong alternative generative modeling framework of LHC events, by easily incorporating sophisticated physics-inspired neural net architectures.
At the same time, EBMs can serve as generic signal detectors, since out-of-distribution (OOD) detection comes naturally in the form of energy comparison. More importantly, EBMs incur fewer spurious correlations in OOD detection. This plays a slightly different role in the context of signal searches at the LHC. There are correlations that are real and useful but at the same time hinder effective signal detection. As we will see in Section 4.2, the EBMs are free from the notorious correlation observed in many anomaly detection methods in HEP, in both the generative and the discriminative approaches.
As summarized in Table 1, we build a multi-tasking framework for High Energy Physics. To that end, we construct an energy-based model of the fundamental interactions of elementary particles to simulate the resulting radiation patterns. We especially employ the short-run Markov Chain Monte Carlo for the EBM training. We show that EBMs are able to generate realistic event patterns and can be used as generic anomaly detectors free from spurious correlations. We also explore EBM-based hybrid modeling combining generative and discriminative models for HEP events. This unified learning framework paves for future-generation event simulators and automated new physics search strategies. It opens a door for combining different methods and components towards a powerful multi-tasking engine for scientific discovery at the Large Hadron Collider.
Describing HEP Events Most particle interactions happening at the LHC are governed by Quantum Chromodynamics (QCD), due to the hadronic nature of the proton-proton collision. Thus jets are enormously produced by these interactions. A jet is formed by collimated radiations originating from highly energetic elementary particles (e.g., quarks, gluons, and sometimes highly-boosted electro-weak bosons). The tracks and energy deposits of the jets left in the particle detectors reveal the underlying physics in complex patterns. A jet may have hundreds of jet constituents, resulting in high-dimensional and complex radiation patterns and bringing difficulties in encoding all the information with expert features. Reconstructing, identifying, and classifying these elementary particles manifest in raw data is thus critical for ongoing physics analysis at the LHC.
Specifically, each constituent particle within a jet has  $(\log p_T,\eta ,\phi)$  as the descriptive coordinates in the detector's reference frame. Here,  $p_T$  represents the transverse momentum perpendicular to the
beam axis, and  $(\eta ,\phi)$  denotes the spatial coordinates within the cylindrical detector. (More details about the datasets can be found in Appendix A.) Thus a jet can be described by the high-dimensional vector  $\mathbf{x} = \{(\log p_T,\eta ,\phi)_i\}_i^N$  , supposing the jet has N constituents. Our goal is thus to model the data distribution  $p(\mathbf{x})$  precisely, and  $p(y|\mathbf{x})$  in the case of supervised classification with y denoting the corresponding jet type.
Parameterized Generative Modeling At the LHC, event simulation serves as an important handle for background estimation and data analysis. For many years, physics event simulators [10] have been built on Monte Carlo methods based on physics rules. These generators are slow and need to be tuned to the data frequently. Deep neural networks provide us with an efficient parameterized generative approach to event simulation for the coming decades. Generative modeling in LHC physics has been experimented with GANs (image-based [58, 27] and point-cloud-based [43]) and VAEs [64]. There are also GANs [8] working on high-level features for event selection. However, as mentioned in the introduction, these models all require explicit generators, which brings obstacles to situations with complex data formats and sophisticated neural architectures.
OOD Detection for New Physics Searches Despite the great efforts in searching for new physics signals at the LHC, there is no hint of beyond-Standard-Model physics. Given the large amount of data produced at the LHC, it has been increasingly challenging to cover all the possible search channels. Experiments at the LHC over the past decades have been focused on model-oriented searches, as in searching for the Higgs boson [26, 36]. The null results up to now from model-driven searches call for novel solutions. We thus shift to model-independent and data-driven anomaly detection strategies, which are data-oriented rather than theory-guided, for new physics searches.
Neural nets-based anomaly detection in HEP takes different formats. Especially, unsupervised generative models trained on background events can be used to detect potential unseen signals. More specifically, Variational Autoencoders [15, 41, 5] have been employed to detect novel signals, directly based on detector-level features. High-level observables-based VAEs [12, 39] have also been explored. However, naively trained VAEs (and Autoencoders) are usually subject to spurious correlations which result in failure modes in anomaly detection, partially due to the format of the reconstruction loss involved. Usually, one needs to explicitly mitigate these spurious correlations with auxiliary tasks such as outlier exposure [35, 15], or modified model components such as autoencoders with smeared pixel intensities [29] and autoencoders augmented with energy-based density estimation [69, 20].
At the same time, EBM naturally has the handle for discriminating between in-distribution and out-of-distribution examples and accommodates tolerance for spurious correlations. While in-distribution data points are trained to have lower energy, energies of OOD examples are pushed up in the learning process. This property has been used for OOD detection in computer vision [25, 33, 23, 71]. This thus indicates great potential for EBM-based new physics detection at the LHC.
3.1 Energy Based Models
Energy-based models are constructed to model the unnormalized data probabilities. They leverage the property that any exponential term  $\exp(-E(\mathbf{x}))$  is non-negative and thus can represent unnormalized probabilities naturally. The data distribution is modeled through the Boltzmann distribution:  $p_{\theta}(\mathbf{x}) = \exp(-E_{\theta}(\mathbf{x})) / Z(\theta)$  with the energy model  $E_{\theta}(\mathbf{x})^{1}: \mathcal{X} \to \mathbb{R}$  mapping  $\mathbf{x} \in \mathcal{X}$  to a scalar. And the partition function  $Z = \int \tilde{p}(\mathbf{x}) d\mathbf{x} = \int \exp(-E_{\theta}(\mathbf{x})) d\mathbf{x}$  integrates over all the possible states.
EBMs can be learned through maximum likelihood  $\mathbb{E}_{p_D}(\log p_\theta (\mathbf{x}))$ . However, the training of EBMs can pose challenges due to the intractable partition function in  $\log p_{\theta}(\mathbf{x}) = -E_{\theta}(\mathbf{x}) - \log Z(\theta)$ . Though the partition function is intractable, the gradients of the log-likelihood do not involve the partition function directly. Thus when taking gradients w.r.t. the model parameters  $\theta$ , the partition
<sup>1</sup>We note that this energy function is not to be confused with the physical energy of the HEP events/objects. Here the energy is essentially an abstract quantity used for probabilistic modelling of the radiation patterns and substructures. It can approximately serve as a measure of radiation substructure or complexity when describing HEP events.
Figure 1: Schematic of the EBM model. The energy function  $E(\mathbf{x},y)$  is estimated with a transformer. The training procedure is governed by Contrastive Divergence (the vertical dimension), for which the model distribution estimation  $q_{\theta}(\mathbf{x})$  is obtained with Langevin Dynamics (the horizontal dimension), evolving samples from random noises  $\mathbf{x}_0^-$ .
function is canceled out. The gradient of the maximum likelihood loss function can be written as:
$$
\begin{array}{l} \nabla_ {\theta} \mathcal {L} (\theta) = - \mathbb {E} _ {p _ {D} (\mathbf {x})} [ \nabla_ {\theta} \log p _ {\theta} (\mathbf {x}) ] (1) \\ = \mathbb {E} _ {\mathbf {x} ^ {+} \sim p _ {D} (\mathbf {x})} \left[ \nabla_ {\theta} E _ {\theta} \left(\mathbf {x} ^ {+}\right) \right] - \mathbb {E} _ {\mathbf {x} ^ {-} \sim p _ {\theta} (\mathbf {x})} \left[ \nabla_ {\theta} E _ {\theta} \left(\mathbf {x} ^ {-}\right) \right], (2) \\ \end{array}
$$
where  $p_D(\mathbf{x})$  is the data distribution and  $p_{\theta}(\mathbf{x})$  is the model distribution. The training objective is thus composed of two terms, corresponding to two different learning phases (i.e., the positive phase to fit the data  $\mathbf{x}^{+}$ , and the negative phase to fit the model distribution  $\mathbf{x}^{-}$ ). When parameterizing the energy function with feed-forward neural networks [54], the positive phase is straightforward. However, the negative phase requires sampling over the model distribution. This leads to various Monte Carlo sampling strategies for estimating the maximum likelihood.
Contrasting the energies of the data and the model samples as proposed Contrastive Divergence (CD) [37, 11] leads to an effective strategy to train EBMs with the following CD objective:
$$
D _ {\mathrm {K L}} \left(p _ {D} (\mathbf {x}) \| p _ {\theta} (\mathbf {x})\right) - D _ {\mathrm {K L}} \left(T p _ {D} (\mathbf {x}) \| p _ {\theta} (\mathbf {x})\right), \tag {3}
$$
where  $T$  denotes the one-step Monte Carlo Markov Chain (MCMC) kernel imposed on the data distribution. In more recent approaches for high-dimensional modeling [56, 25], we can directly initialize from random noises to generate MCMC samples instead of initializing from the data distribution as in the original CD approach. This strategy also helps with exploring and mixing between modes.
Negative Sampling The most critical component in the negative phase is the sampling to estimate the model distribution. We employ gradient-based MCMC generation for the negative sampling, which is handled by Langevin Dynamics [52, 66]. As written in Eq. 4, Langevin dynamics uses gradients w.r.t. the input dimensions, associated with a diffusion term to inject stochasticity, to generate a sequence of negative samples  $\{\mathbf{x}_k^-\}_{k=1}^K$ .
$$
\mathbf {x} _ {k + 1} ^ {-} = \mathbf {x} _ {k} ^ {-} - \frac {\lambda^ {2}}{2} \nabla_ {\mathbf {x}} E _ {\theta} \left(\mathbf {x} _ {k} ^ {-}\right) + \lambda \cdot \epsilon , \text {w i t h} \epsilon \sim \mathcal {N} (0, 1) \tag {4}
$$
MC Convergence The training anatomy [55, 56] for short-run non-convergent MCMC and long-run convergent MCMC shows that short-run ( $\sim$  5-100 steps) MCMC initialized from random distributions
is able to generate realistic samples, while long-run MCMC might be oversaturated with lower-energy states.
To improve mode coverage, we use random noise to initialize MCMC chains. To accelerate training, we employ a relatively small number of MCMC steps. In practice, we can reuse the generated samples as initial samples of the following MCMC chains to accelerate mixing, similar to Persistent Contrastive Divergence [63]. Following the procedure in [25], we use a randomly initialized buffer that is consistently updated from previous MCMC runs as the initial samples. (As empirically shown [13, 66], a Metropolis-Hastings step is not necessary. So we ignore this rejection update in our experiments.)
Energy Function Since there is no explicit generator in EBMs, we have much freedom in designing the architectures of the energy function. This also connects with the fast-paced development of supervised neural classifier architectures for particle physics. We can directly reuse the architectures from supervised classifiers in the generative modeling of EBMs. We use a self-attention-based transformer to parameterize the energy function  $E_{\theta}(\cdot)$ . We defer the detailed description to Sec. 3.3.
The full algorithm for training EBMs is described in Algorithm 1.
Algorithm 1 EBM training with MCMC by Langevin Dynamics
Input: training samples  $\{\mathbf{x}_i^+\}_{i=1}^N$  from  $p_{\mathrm{D}}(\mathbf{x})$ , parameterized energy function  $E_{\theta}(\cdot)$ , initial buffer  $\mathcal{B} \gets \emptyset$ , Langevin dynamics step size  $\lambda_{\mathbf{x}}$ , number of MCMC steps K, model parameter learning rate  $\lambda_{\theta}$ , regularization strength  $\alpha$
for Gradient descent step  $1 = 0\ldots \mathrm{L - 1}$  do
$$
\mathbf {x} _ {i} ^ {+} \sim p _ {\mathrm {D}} (\mathbf {x})
$$
$\mathbf{x}_{i,0}^{-} \sim 0.95 * \mathcal{B} + 0.05 * \mathcal{U} \triangleright$  Reinitialize the samples in the buffer with random noise in the probability of 0.05
for Langevin dynamics step  $\mathrm{k} = 0\ldots \mathrm{K - 1}$  do
$\mathbf{x}_{i,k + 1}^{-} = \mathbf{x}_{i,k}^{-} - \lambda_{\mathbf{x}}\nabla_{\mathbf{x}}E_{\theta}(\mathbf{x}_{i,k}^{-}) + 0.005\cdot \epsilon_{k},$ $\epsilon_{k}\sim \mathcal{N}(0,1)\triangleright$  Langevin Dynamics taking limits w.r.t. input dimensions
$$
\begin{array}{l} \mathbf {x} _ {i} ^ {-} \leftarrow \mathbf {x} _ {i, K} ^ {-} \\ \mathcal {L} _ {\mathrm {C D}} = \frac {1}{N} \sum_ {i} \left(E _ {\theta} \left(\mathbf {x} _ {i} ^ {+}\right) - E _ {\theta} \left(\mathbf {x} _ {i} ^ {-}\right)\right) \\ \mathcal {L} _ {\text {r e g}} = \frac {1}{N} \sum_ {i} \left(E _ {\theta} \left(\mathbf {x} _ {i} ^ {+}\right) ^ {2} + E _ {\theta} \left(\mathbf {x} _ {i} ^ {-}\right)\right) ^ {2} \\ \theta \leftarrow \theta - \lambda_ {\theta} \nabla_ {\theta} \left(\mathcal {L} _ {\mathrm {C D}} + \alpha \mathcal {L} _ {\mathrm {r e g}}\right) \\ \mathcal {B} \leftarrow \mathbf {x} _ {i, K} ^ {-} \cup \mathcal {B} \\ \end{array}
$$
$\triangleright L_{2}$  Regularization
> Update model parameters with gradient descent
> Update the buffer with generated samples
3.2 Hybrid Modeling
Neural Classifier as an EBM A classical classifier can be re-interpreted in the framework of EBMs [46, 33, 67, 17, 40], with the logits  $\mathbf{g}(\mathbf{x})$  corresponding to negative energies of the joint distribution  $p(\mathbf{x},y) = \frac{\exp(\mathbf{g}(\mathbf{x})_y)}{Z}$ , where  $\mathbf{g}(\mathbf{x})_y$  denotes the logit corresponding to the label  $\mathbf{y}$ . Thus the probability marginalized over  $\mathbf{y}$  can be written as  $p(\mathbf{x}) = \frac{\sum_y\exp(\mathbf{g}(\mathbf{x})_y)}{Z}$ , with the energy of  $\mathbf{x}$  as  $-\log \sum_y\exp (\mathbf{g}(\mathbf{x})_y)$ . We are then brought back to the classical softmax probability  $\frac{\exp(\mathbf{g}(\mathbf{x})_y)}{\sum_y\exp(\mathbf{g}(\mathbf{x})_y)}$  when calculating  $p(y|\mathbf{x})$ .
This viewpoint provides a novel method for jointly training a supervised classifier and an unsupervised generative model. Specifically, [33] successfully incorporated EBM-based generative modeling into a classifier in image generation and classification. We follow their proposal to train the hybrid model as follows to ensure the classifier  $p(y|\mathbf{x})$  is unbiased. The joint log-likelihood is decomposed into two terms:
$$
\log p (\mathbf {x}, y) = \log p (\mathbf {x}) + \log p (y | \mathbf {x}). \tag {5}
$$
Thus one can maximize  $\log p(\mathbf{x})$  with the contrastive divergence of the EBM with the energy function  $E(\mathbf{x}) = -\log \sum_{y}\exp (\mathbf{g}(\mathbf{x})_{y})$ , and maximize  $\log p(y|\mathbf{x})$  with the usual cross-entropy of the classification.
3.3 Energy-based Models for Elementary Particles
We would like to construct an energy-based model for describing jets and their inner structures. In conceiving the energy function for these elementary particles, we consider the following constraints and characteristics: 1) permutation invariance – the energy function should be invariant to jet constituent permutations, and 2) higher-order interactions – we would like the energy function to be powerful enough to simulate the complex inter-particle interactions.
Thus, we leverage the self-attention-based transformer [65] to approximate the energy function, which takes into account the higher-order interactions between the component particles. As indicated in Eq. 6b, the encoding vector of each constituent  $W$  is connected with all other constituents through the self-attention weights  $A$  in Eq. 6a, which are already products of particle representations  $Q$  and  $K$ .
$$
A = \operatorname {s o f t m a x} \left(Q \cdot K ^ {T} / \sqrt {d _ {\text {m o d e l}}}\right) \tag {6a}
$$
$$
W = A \cdot V \tag {6b}
$$
Moreover, we can easily incorporate particle permutation invariance [70] in the transformer, by summing up the encodings of each jet constituent. The transformer architecture is shown in Fig. 1. The coordinates  $(\log p_T,\eta ,\phi)$  of each jet constituent are first embedded into a  $d_{\mathrm{model}}$  -dimensional space through a linear layer, then fed into  $N_{L}$  self-attention blocks sequentially. After that, a sum-pooling layer is used to sum up the features of the jet constituents. Finally, a multi-layer-perceptron projector maps the features into the energy score. Model parameters are recorded in Table 4 of Appendix A.
Training Details The training set consists of 300,000 QCD jets. We have 10,000 samples in the buffer and reinitialize the random samples with a probability of 0.05 in each iteration. We use a relatively small number of steps (e.g., 24) for the MCMC chains. The step size  $\lambda_{\mathbf{x}}$  is set to 0.1 according to standard deviations of the input features. The diffusion magnitude within the Langevin dynamics is set to 0.005. The number of steps used in validation steps is set to 128 for better mixing.
We use Adam [44] for optimization, with the momenta  $\beta_{1} = 0.0$  and  $\beta_{2} = 0.999$ . The initial learning rate is set to 1e-4, with a decay rate of 0.98 for each epoch. We use a batch size of 128, and train the model for 50 epochs. More details can be found in Appendix A.
Model Validation In monitoring the likelihood, the partition function can be estimated with Annealed Importance Sampling (AIS) [53]. However, these estimates can be erroneous and consume a lot of computing resources. Fortunately for physics events, we have well-designed high-level features as a handle for monitoring the generation quality. Especially, we employ the boost-invariant jet transverse
momentum  $p_{T} = \sum_{i=1}^{N} p_{Ti}$  and the Lorentz-invariant jet mass  $M = \sqrt{\left(\sum_{i=1}^{N} E_{i}\right)^{2} - \left(\sum_{i=1}^{N} \mathbf{p}_{i}\right)^{2}}$  as the validation observables. And we calculate the Jensen-Shannon Divergence (JSD), between these high-level observable distributions of the data and the model generation, as the metric. In contrast to the short-run MCMC in the training steps, we instead use longer MCMC chains for generating the validation samples.
When we focus on downstream tasks such as OOD detection, it's reasonable to employ the Area Under the Receiver Operating Characteristic (ROC) Curve (AUC) as the validation metric, with Standard Model top jets serving as the benchmark signal.