SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Irvine, CA 92697
Kevin Greif
Department of Physics and Astronomy
kggreif@uci.edu
Michael Fenton
mjfenton@uci.edu
Aishik Ghosh
Physics Division
Lawrence Berkeley National Laboratory
Daniel Whiteson
High-energy collisions at the Large Hadron Collider (LHC) provide valuable insights into open questions in particle physics. However, detector effects must be corrected before measurements can be compared to certain theoretical predictions or measurements from other detectors. Methods to solve this inverse problem of mapping detector observations to theoretical quantities of the underlying collision are essential parts of many physics analyses at the LHC. We investigate and compare various generative deep learning methods to approximate this inverse mapping. We introduce a novel unified architecture, termed latent variational diffusion models, which combines the latent learning of cutting-edge generative art approaches with an end-to-end variational framework. We demonstrate the effectiveness of this approach for reconstructing global distributions of theoretical kinematic quantities, as well as for ensuring the adherence of the learned posterior distributions to known physics constraints. Our unified approach achieves a distribution-free distance to the truth of over 20 times smaller than non-latent state-of-the-art baseline and 3 times smaller than traditional latent diffusion models.
Particle physics experiments at the Large Hadron Collider study the interactions of particles at high energies, which can reveal clues about the fundamental nature of matter and forces. However, the properties of particles which result from the collisions must be inferred from signals in the detectors which surround the collision. Though detectors are designed to reconstruct the properties
of particles with high fidelity, no detector has perfect efficiency and resolution. A common strategy to account for these effects is simulation-based inference [1], in which the detector resolution and inefficiency are modeled by a simulator. Samples of simulated events can then be compared to observed data to perform inference on theoretical parameters. However, simulators with high fidelity are computationally expensive and not widely accessible outside of experimental collaborations.
An alternative approach is the reverse, mapping the observed detector signatures directly to the unobserved truth-level information. In a particle physics context, this procedure is referred to as "unfolding". In practice, the quantum mechanical nature of particle interactions makes the forward map from the true particle properties to observed data not one-to-one. As a result, there is no true inverse function which can map a given detector observation to a single point in the truth-level space. Such inverse problems are challenging, but unfolded data allows for direct comparisons with theoretical predictions and across experiments, without requiring access to detector simulation tools which may not be maintained long-term.
Unfolding methods such as Iterative D'Agostini [2], Singular Value Decomposition [3], and TUnfold [4] have seen frequent use by experimental collaborations like ATLAS [5] and CMS [6]. However, these techniques are limited to unfolding only a few dimensions, and require binning the data, which significantly constrains later use of the unfolded distributions. The application of machine learning techniques [7] has allowed for the development of un-binned unfolding with the capacity to handle higher-dimensional data. One approach is to use conditional generative models, which learn to sample from the truth-level distributions when conditioned on the detector-level data; examples include applications of generative adversarial networks [8, 9], invertible networks [10, 11], and variational auto-encoders [12]. An alternative approach uses classification models as density estimators which learn to correct imprecise truth-level distributions with re-weighting [13-15]. Generative methods naturally produce unweighted events, an advantage over classification methods which may generate very large weights or even fail if the original distributions do not sufficiently cover the entire support of the true distribution. However, generative models are not always guaranteed to produce samples which respect the important physical constraints of the original sample. While making important strides, none of these methods have cracked the ultimate goal, full-event unfolding, where the full high-dimensional detector-level observations are mapped to truth-level objects.
This paper introduces a novel generative unfolding method utilizing a diffusion model [16-18] to map detector to truth-level distributions. Diffusion models are a class of generative models which learn to approximate a reverse noise diffusion process and have proven successful in natural image generation [19, 20] and recently scientific applications such as molecular link design [21]. Diffusion models excel in learning high-dimensional probability distributions at higher fidelity than normalizing flows and without the adversarial min-max loss of GANs. In HEP, they have already found use for approximating calorimeter simulations [22-25]. Latent diffusion models (LDMs), a specific class of diffusion models, perform the denoising in an abstract latent space [26] and excel in image generation tasks. These latent embeddings are often pre-trained on secondary objectives, such as VAE reconstruction tasks or CLIP [27], to limit computational and memory requirements. We unify the abstract embedding space of latent diffusion with the recently formalized variational diffusion approach [28] to develop an end-to-end variational latent diffusion model (VLD) achieving state-of-the-art performance in complex HEP generative tasks.
2.1 Unfolding
Let  $f_{\mathrm{det}}(y)$  be the distribution which governs an observed detector-level data set  $y = \{y_i\}$ . An unfolding method aims to sample from a pre-detector distribution  $f_{\mathrm{parton}}(x)$ , where parton refers to an unobserved state of interest to physicists.  $f_{\mathrm{parton}}(x)$  is related to  $f_{\mathrm{det}}$  via convolution with a "response" function  $p(y|x)$  over the possible true values  $x$ . The response function describes the decay of the initial, unstable particles into stable particles and their interaction with the detector.
$$
f _ {\det } (y) = \int d x p (y | x) f _ {\text {p a r t o n}} (x) \tag {1}
$$
In other fields, this kind of problem is often referred to as "deconvolution".
Figure 1: Visual representations of the different perspectives captured by the parton and detector level data. The parton-level data represents a fundamental theoretical description of the decay, whereas the detector-level data represents the real measurements observed after the decay. The primary challenge in unfolding is to infer the theoretical parton representation from the observed data.
No closed form expression exists for  $p(y|x)$ , but Monte-Carlo-based simulation can sample from parton values  $x$  and produce the corresponding sample  $y$ . The parton distribution can be recovered via the corresponding inverse process if one has access to a pseudo-inversion of the response function  $p(x|y)$ , also known as the posterior.
$$
f _ {\text {p a r t o n}} (x) = \int d y p (x | y) f _ {\det } (y) \tag {2}
$$
Generative unfolding methods build the posterior as a generative model, which can be used to sample from  $p(x|y)$ . The desired parton distribution is then obtained by Equation 2. Simulated pairs of parton-detector data,  $(x,y)$ , may be used to train the generative model.
An important issue when choosing to directly model the posterior is that this quantity is itself dependent on the desired distribution  $f_{\mathrm{parton}}(x)$ , the prior in Bayes' theorem:
$$
p (x \mid y) = \frac {p (y \mid x) f _ {\text {p a r t o n}} (x)}{f _ {\det} (y)} \tag {3}
$$
Producing the data set used to train the generative model requires choosing a specific  $f_{\mathrm{parton}}(x)$ , which influences the learned posterior. In application to new datasets, this will lead to an unreliable estimate of the posterior density if the assumed prior is far enough from the truth distribution. A common method to overcome this challenge is to apply an iterative procedure, in which the assumed prior is re-weighted to match the approximation to the truth distribution provided by the unfolding algorithm [2]. Though application of this iterative procedure is not shown in this paper, the principle has been demonstrated with other generative unfolding methods [29], for which the conditions are similar.
2.2 Semi-Leptonic Top Quark Pair Production
Collisions at the LHC which result in a pair of top quarks allow for sensitive probes of new theories of physics, which makes measurement of the top quark properties an important task. Top quarks are unstable, decaying almost immediately to a  $W$  boson and a bottom quark; the  $W$  boson can then decay hadronically to two quarks or leptonically to a charged lepton and neutrino. The case where one of the produced top quarks decays hadronically and the other decays leptonically is known as the semi-leptonic decay mode, see Fig. 1a. The 4-momenta (three momentum components, one mass) of
these six objects (four quarks, the charged lepton, and the neutrino) constitute the parton-level space in this context.
The four quarks each produce a shower of particles (jets) which interact with the detector, while the neutrino passes through without leaving a trace. The resulting observed detector signature which defines the detector-level space is then quite complex, see Fig. 1b.
The semi-leptonic  $t\bar{t}$  process has been studied by the ATLAS and CMS collaborations to measure various properties of the top quark and to search for new particles and interactions [30-35]. Many of these measurements use existing unfolding techniques, which limit the unfolded measurements to one or two dimensions. An un-binned and high dimensional unfolding technique would allow physicists to use the full power of their data.
2.3 Variational Autoencoders
Variational Autoencoders (VAEs) are a class of generative models combining an autoencoder architecture with probabilistic modeling [36, 37]. VAEs learn a non-linear latent representation of input data through an encoder and decoder network while incorporating probabilistic methods and sampling through the reparameterization trick [36]. VAEs have been applied to numerous applications, such as image synthesis [38] and natural language processing [39], among many others.
The VAE encoder network is parameterized as a probabilistic function, approximating the posterior distribution of the latent variables  $z$  conditioned on the input data:  $q(z|x)$ . The decoder network likewise models the generative distribution conditioned on the latent variables  $p(x|z)$ . VAEs are trained by maximizing the evidence lower bound (ELBO), which is a lower bound on the log-likelihood of the data under the generative model [36]. The ELBO includes a reconstruction loss for training the decoder and a KL-divergence objective which enforces a regularization constraint on the learned latent posterior to a prior distribution  $p(z)$ .
$$
\mathcal {L} _ {\mathrm {V A E}} = \mathbb {E} _ {z \sim q (z | x)} [ - \log p (x | z) + D _ {K L} (q (z | x) \| p (z)) ] \tag {4}
$$
Conditional VAEs (CVAEs) [40] extend the VAE framework by conditioning both the encoder and decoder networks on additional information, such as class labels, via an arbitrary conditioning vector  $y$ . This allows CVAEs to generate samples with specific desired properties, providing more control over the generated outputs.
$$
\mathcal {L} _ {\mathrm {C V A E}} = \mathbb {E} _ {z \sim q (z | x, y)} [ - \log p (x | z, y) + D _ {K L} (q (z | x, y) \| p (z | y)) ] \tag {5}
$$
2.4 Variational Diffusion Models
Variational Diffusion Models (VDMs) define a conditional probabilistic generative model which exploits the properties of diffusion probabilistic models to generate samples by learning to reverse a stochastic flow [41]. VDMs may be seen as an extension of VAEs to a (possibly infinitely) deep hierarchical setting. The Gaussian diffusion process defines the forward stochastic flow with respect to time  $t \in [0, 1]$  over the latent space  $z_{t} \in \mathbb{Z}$  and conditioned on  $y$  as:
$$
q \left(z _ {t} | x, y\right) \sim \mathcal {N} \left(\alpha_ {t} x, \sigma_ {t} \mathbb {I}\right) \tag {6}
$$
The flow parameters,  $\sigma_{t}$  and  $\alpha_{t}$  are defined by a noise schedule. We use the continuous Variance Preserving (VP) framework throughout this work and derive these flow parameters based on a learned signal-to-noise ratio,  $e^{-\gamma_{\phi}(t)}$ , where:
$$
\sigma_ {t} = \sqrt {\operatorname {s i g m o i d} \left(\gamma_ {\phi} (t)\right)} \text {a n d} \alpha_ {t} = \sqrt {\operatorname {s i g m o i d} \left(- \gamma_ {\phi} (t)\right)}
$$
Assuming it is possible to sample from the terminal distribution  $p(z_{1})$ , we may produce samples from the data distribution by inverting the flow and sampling previous latent representations conditioned on future latent vectors. The inverse flow is modeled as  $q(z_{s}|z_{t},\hat{x}_{\theta}(z_{t},t,y))$  where  $\hat{x}_{\theta}$  is an approximate denoising of the original data at the current time-step. In practice, the data denoising is implemented using a variance-independent noise prediction network,  $\hat{\epsilon}_{\theta}$ , by the equation  $x_{\hat{\theta}}(z_{t},t,y) = \frac{(z_{t} - \sigma_{t}\hat{\epsilon}_{\theta}(z_{t},t,y))}{\alpha_{t}}$ . The noise prediction network,  $\hat{\epsilon}_{\theta}$ , is parameterized using a deep neural network. The learnable noise schedule  $\gamma_{\phi}(t)$  is also parameterized using a positive valued, monotonic neural network with learnable end-points  $\gamma_{min} = \gamma (0)$  and  $\gamma_{max} = \gamma (1)$  [41]. Following
Figure 2: A block diagram of the end-to-end VLD models with trainable components. The conditional paths are drawn in different colors depending on which model variations employ them. We use the continuous, variance preserving SDE diffusion formulation introduced in [17] and [41]. We show the equivalent ODE form of SDE equation in the diagram.
the VP framework, the noise schedule is regularized so that the terminal distribution is the unit Gaussian:  $p(z_{1}) \sim \mathcal{N}(\mathbf{0},\mathbb{I})$ . Both the noise prediction network and the noise schedule network are trained using the modified ELBO for continuous-time diffusion models [41]:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {V D M}} = D _ {K L} \left(q \left(z _ {1} | x, c\right) \| p \left(z _ {1}\right)\right) + \mathbb {E} _ {q \left(z _ {0} \mid x\right)} \left[ - \log p (x \mid z _ {0}, y) \right] \\ + \mathbb {E} _ {\epsilon \sim \mathcal {N} (\mathbf {0}, \mathbb {I}), t \sim \mathcal {U} (0, 1)} \left[ \gamma_ {\phi} ^ {\prime} (t) \| \epsilon - \hat {\epsilon} _ {\theta} \left(z _ {t}, t, y\right) \| _ {2} ^ {2} \right] \tag {7} \\ \end{array}
$$
2.5 Latent Diffusion
Latent diffusion models (LDMs)[26] are a deep generative framework that operates the diffusion process in an abstract latent space learned by a VAE to sample high-dimensional data  $p_D(x|z,y)$ , possibly conditioned on a secondary dataset  $p_C(y)$ . This approach has proven dramatically successful when employed in natural image generation applications, including text-to-image synthesis, inpainting, denoising, and style transfer [26, 20].
LDMs first train an unconditional VAE to embed the data distribution into a low dimensional latent representation using a traditional VAE approach,  $q(z_{x}|x)$  and  $p(x|z_x)$ , regularizing the latent space towards a standard normal  $p(z_{x})\sim \mathcal{N}(\mathbf{0},\mathbb{I})$ . A secondary encoder may be trained on the conditioning data  $p(z_y|y)$  along-side VAE, typically using a CLIP objective [27] to map the two datasets into a common latent space. The diffusion process is then trained to reconstruct the latents  $z_{x}$  from the flow latents  $p(z_{x}|z_{0},z_{y})$ . The diffusion model training remains otherwise identical to the standard diffusion framework.
Critically, the most successful methods train the VAE, the conditional encoder, and the diffusion process individually. While computationally efficient, this independence limits the models' generative power as each component is trained on subsets of the overall conditional generative objective. It may be possible to recover additional fidelity by instead training all components using a unified conditional generation objective. While several methods allow for training a VAE along-side diffusion [42, 43], these approaches either cannot train diffusion in the latent space or cannot account for a conditional, fully variational model. We construct a unified variational framework to allow for a conditional, probabilistic, end-to-end diffusion model.
3 Variational Latent Diffusion
This work integrates the learning capabilities of latent diffusion models with the theoretical framework of variational diffusion models in a unified conditional variational approach. This unified variational model combines the conditioning encoder, data VAE, and diffusion process into a single loss function. This framework enables further enhancement of these methods through a conditional data encoder or decoder, and an auxiliary physics-informed consistency loss which may be enforced throughout the network. We refer to this combined method as Variational Latent Diffusion (VLD), see Fig 2. The primary contributions of this paper are to define this unified model and derive the appropriate loss function to train such a model.
Conditioning (Detector) Encoder In traditional LDMs, the conditioning encoder,  $p(z_{y}|y)$ , is pretrained through an auxiliary loss term, such as CLIP [27], which aims to unify the latent space of the conditioning and data. While this approach is efficient, it may not be optimal: the encoder is trained
on one objective, and then repurposed to act as a conditioning encoder for a separate generative model. With the end-to-end framework, we simultaneously learn this encoder alongside other generative terms, enabling us to efficiently train a variable-length, high-dimensional encoder fine-tuned for the generative objective. In this work, we simplify this conditioning encoder by restricting it to a deterministic mapping,  $z_{y} = f_{\theta}(y)$ . This decision is based on prior work such as the CVAE which opts for a deterministic conditional encoder, as well as for simplicity as there is a lack of motivating benefits from a stochastic encoder.
Data (Parton) VAE The traditional LDM VAE is unconditional, as this allows it to be easily pre-trained and reused for different diffusion models. As we are training a unified conditional generative model in an end-to-end fashion, we have the option to extend the encoder and decoder with conditional probabilistic models:  $q_{\mathrm{C - VLD}}(z_x|x,z_y)$  and  $p_{\mathrm{C - VLD}}(x|z_x,z_y)$ . We experiment with both a conditional and unconditional VAE. Additionally, we explore an intermediate method that uses a conditioned encoder to estimate the VAE posterior,  $q_{\mathrm{UC - VLD}}(z_x|x,z_y)$ , but employs an unconditional decoder during generation  $p_{\mathrm{UC - VLD}}(z_x|x)$ .
VLD ELBO We interpret the continuous VDM as an infinitely deep hierarchical VAE as presented by Kingma et al. [41]. This interpretation allows us to seamlessly integrate the VAE into a unified diffusion framework by incorporating the VAE as an additional component in the hierarchy. Consequently, the hierarchical variational ELBO incorporates an extra KL divergence term, which serves to regularize the encoder posterior distribution [44]. We combine this hierarchical objective with the denoising loss term derived in [41] to define a combined ELBO for the entire generative model.
$$
\begin{array}{l} \mathcal {L} _ {V L D} = D _ {K L} \left(q \left(z _ {1} \mid x, z _ {y}\right) \| p \left(z _ {1}\right)\right) + \mathbb {E} _ {q \left(z _ {x} \mid x, z _ {y}\right)} \left[ - \log p \left(x \mid z _ {x}, z _ {y}\right) \right] \\ + D _ {K L} \left(q \left(z _ {x} \mid x, z _ {y}\right) \| p \left(z _ {x} \mid z _ {0}\right)\right) + \mathbb {E} _ {\epsilon \sim \mathcal {N} (\mathbf {0}, \mathbb {I}), t \sim \mathcal {U} (0, 1)} \left[ \gamma_ {\phi} ^ {\prime} (t) \| \epsilon - \hat {\epsilon} _ {\theta} \left(z _ {t}, t, z _ {y}\right) \| _ {2} ^ {2} \right] \tag {8} \\ \end{array}
$$
The additional KL term may be derived explicitly if we assume a Gaussian VAE and a Gaussian diffusion process. The posterior is parameterized using a learned Gaussian, as in a standard VAE:  $q(z_{x}|x,z_{y}) \sim \mathcal{N}(\mu_{\theta}(x,z_{y}),\sigma_{\theta}(x,y))$ . The prior can be reformulated using the definition of the forward flow from Equation 6. Employing the reparameterization trick, we can rewrite the expression of  $z_{0}$  in terms of  $z_{x}$  as  $z_{0} = \alpha_{0}z_{x} + \sigma_{0}\epsilon$ , where  $\epsilon \sim \mathcal{N}(\mathbf{0},\mathbb{I})$ . Solving this equation for  $z_{x}$  yields another reparameterized Gaussian, which allows us to define the prior over  $z_{x}$  as:
$$
p \left(z _ {x} \mid z _ {0}\right) \sim \mathcal {N} \left(\frac {1}{\alpha_ {0}} z _ {0}, \frac {\sigma_ {0}}{\alpha_ {0}} \mathbb {I}\right) \tag {9}
$$
Physics-Informed Consistency Loss Reconstructing the mass of truth-level physics objects is challenging due to their highly peaked, low-variance distributions. For certain particles like leptons, the mass distribution exhibits a two-valued delta distribution, while for light quarks, it is consistently set to zero. Predicting these distributions is more difficult than predicting the energy of truth-level physics objects, which have a broader range. In special relativity, the mass  $(M)$ , energy  $(M)$ , and momentum  $\mathbf{p}$  of a particle are related by  $c^4 M^2 = E^2 - (c\|\mathbf{p}\|)^2$ . We operate in natural units with a unit speed of light  $c = 1$ . Forcing the predicted mass, energy, and momentum to satisfy this equality improves stability and accuracy by capturing this underlying physical relationship between these quantities. We introduce a consistency loss,  $\mathcal{L}_C$ , in addition to the regular reconstruction loss, weighted by a hyper-parameter  $\lambda_C$ . Similar physics-informed constraints have previously been used for generative models in HEP [45-48]. The consistency loss minimizes the mean absolute error (MAE) between the predicted mass term and the corresponding predicted energy and momentum terms, encouraging the model to learn a more physically consistent representation.
$$
\mathcal {L} _ {C} = \lambda_ {C} \left| \hat {M} ^ {2} - \left(\hat {E} ^ {2} - \| \hat {p} \| ^ {2}\right) \right| \tag {10}
$$
4 Unfolding Semi-Leptonic  $t\bar{t}$  Events
Generative models can be trained to estimate a conditional density given any set of paired data. In the unfolding context, a Monte Carlo simulation can be used to generate pairs of events at detector and parton level. The density of parton level events  $f_{\mathrm{parton}}(x)$  can be taken as the data distribution, and the density of detector level events  $f_{\mathrm{det}}(y)$  can be taken as the conditioning distribution. A generative model can then be used to unfold a set of observed events to the corresponding parton level events with the following procedure:
1. Sample a parton configuration from the distribution governing the process of interest:  $x \sim p_D(x)$ . This can be done using a matrix element solver such as MADGRAPH [49].
2. Sample a possible detector observation  $y \sim p_{C}(y|x)$  using the tools PYTHIA8 [50] and DELPHES [51], which simulate the interactions of particles in flight and the subsequent interactions with a detector.
3. Train a generative model to approximate the inverse distribution  $p_{\theta}(x|y)$ .
4. Produce new posterior samples for inference data with unknown parton configurations.
4.1 Generative Models