SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Multiple baseline generative models are assessed alongside the novel VLD approach, with the goal of investigating the impact of each VLD component, including the conditional VAE, the denoising model, and the variational aspects of the diffusion. Note that the network architectures of the VAEs, denoising networks, and detector encoders are identical where relevant.
CVAE A traditional conditional Variational Autoencoder [40] approach employing a conditional encoder and decoder. We use a Gaussian likelihood for the decoder and a standard normal prior for the encoder, following conventional practices for VAE models.
CINN A conditional Invertible Neural Network [52], which represents the latest deep learning approach that has demonstrated success in unfolding tasks. This model utilizes a conditional normalizing flow to train a mapping from a standard normal distribution to the parton distribution, conditioned on the detector variables. The normalizing flow incorporates an All-In-One architecture [53], following the hyperparameters detailed in the CINN paper [52], which combines a conditional affine layer with global affine and permutation transforms to create a powerful invertible block. In this work, the MMD objective defined in [52] is replaced with a MSE reconstruction objective and the physics-informed consistency loss, for fair comparison with other models.
VDM A Variational Diffusion Model (VDM) [41] that aims to denoise the parton vector directly. This model serves as a baseline for examining the impact of the VAE in latent diffusion approaches. The denoising model is trained using a Mean Squared Error loss against the generated noise.
LDM A Latent Diffusion Model (LDM) with a pre-trained VAE, popularized by recent achievements in text-to-image generative models [26]. The VAE is pre-trained using a Gaussian likelihood and a minimal prior weight  $(10^{-4})$ . This baseline is meant to highlight the importance of the unified end-to-end architecture as all other aspects of the network are identical to the proposed method.
VLD, C-VLD, UC-VLD These models are variations on the proposed unified Variational Latent Diffusion (VLD) architecture. They correspond to an unconditional VAE (VLD), a conditional encoder and decoder (C-VLD), or a conditional encoder with an unconditional decoder (UC-VLD).
4.2 Detector Encoder
All of the generative models are conditioned on detector observations, represented as a set of vectors for each jet and lepton in the event, as described in Section 5. Additionally, the missing transverse momentum (MET) from the neutrino is included as a fixed-size global variable. As there is no inherent ordering to these jets, it is crucial to use a permutation-invariant network architecture for the encoder. We use the jet transformer encoder from the SPANet (v2.1, BSD-3) [54] jet-parton reconstruction network to embed detector variables. This architecture leverages the permutation invariance of attention to contextually embed a set of momentum vectors. We extract the fixed-size event embedding vector from the central transformer, mapping the variable-length, unordered detector observations into a fixed-size real vector  $E_{C}(y) = z_{y} \in \mathbb{R}^{D}$ .
4.3 Parton Encoder-Decoder
For a given event topology, partons may be represented as a fixed-size vector storing the momentum four-vectors of each theoretical particle. We describe the detailed parton representation in Section 5,
Wasserstein	Energy	K-S	KL64	KL128	KL256
VLD	108.76	7.59	4.08	3.47	3.74	4.53
UC-VLD	73.56	6.35	3.41	5.77	7.10	8.48
C-VLD	389.62	25.39	4.65	9.54	10.09	10.79
LDM	402.32	24.09	5.91	14.71	16.34	17.92
VDM	2478.35	181.35	17.14	29.28	32.29	35.60
CVAE	484.56	32.29	6.37	7.79	9.17	10.60
CINN	3009.08	185.13	15.74	28.55	30.19	32.37
Table 1: Total distance measures across all 55 components for every model and metric. The independent sum of 1-dimensional distances for each component are summed across all the components to compute the total metrics.
which consists of a single 55-dimensional vector for each event. The encoder and decoder network employ a ConvNeXt-inspired block structure [55] for the hidden layers, described in Appendix A, which allows for complex non-linear mappings into the latent space. Unlike traditional VAE applications, our latent space may be higher dimensionality than the original space. The VAE's primary purpose therefore differs from typical compression applications, and instead solely transforms the partons into an optimized representation for generation.
The encoder uses this feed-forward block network and produces two outputs: the mean,  $\mu_{\theta}(x,z_y)$  and log-standard deviation,  $\sigma_{\theta}(x,z_y)$ , of the encoded vector, possibly conditioned on the detector observations  $z_{y}$ . The decoder similarly accepts a latent parton representation, possible conditioned on the detector, and produces a deterministic estimate of the original parton configuration  $\hat{x} = D(z_x,z_y)$ .
Dataset Each of the generative approaches is trained to unfold a simulated semi-leptonic  $t\bar{t}$  production data set. Matrix elements are evaluated at a center-of-mass energy of  $\sqrt{s} = 13$  TeV using MADGRAPH_AMC@NLO [49] (v2.7.2, NCSA license) with a top mass of  $m_{t} = 173$  GeV. The parton showering and hadronization are simulated with PYTHIA8 [50] (v8.2, GPL-2), and the detector response is simulated with DELPHES [56] (v3.4.1, GPL-3) using the default CMS detector card. The top quarks each decay to a W-boson and  $b$ -quark, with the W-bosons subsequently decaying either to a pair of light  $(u,d,s,c)$  quarks  $qq'$  or a lepton-neutrino pair  $\ell \nu$  ( $\ell = e,\mu$ ). A basic event selection is then applied on the reconstructed objects at detector-level. Electrons and muons are selected with a transverse momentum requirement of  $p_{\mathrm{T}} > 25$  GeV and absolute value of pseudorapidity  $|\eta| < 2.5$ . The  $b$  and light quarks are reconstructed with the anti- $k_{\mathrm{T}}$  jet algorithm [57] using a radius parameter  $R = 0.5$  and the same  $p_{\mathrm{T}}$  and  $|\eta|$  requirements as the leptons. Jets originating from  $b$ -quarks are identified with a "b-tagging" algorithm that incorporates a  $p_{\mathrm{T}}$  and angular ( $\eta,\phi$ ) dependent identification efficiency and mis-tagging rate. Selected events are then required to contain exactly one lepton and at least 4 jets, of which at least two must be  $b$ -tagged. Events are separated into training and testing data sets, consisting of 9,865,402 and 1,332,514 events respectively.
**Parton Data** The kinematics for the six final state partons are used as unfolding targets  $(b, q_1, q_2, \bar{b}, \nu_l, l)$ , along with the kinematics of the intermediate resonance particles ( $W_{\mathrm{lep}}$ ,  $W_{\mathrm{had}}$ ,  $t$ ,  $\bar{t}$ ), and the entire  $t\bar{t}$  system. The parton-level data consists of 11 momentum vectors, each represented by the five quantities ( $M$ ,  $\log E$ ,  $p_x$ ,  $p_y$ ,  $p_z$ ); where  $M$  is the invariant mass of the particle;  $E$  is the energy; and  $p_x$ ,  $p_y$ , and  $p_z$  are the Cartesian coordinates of the momentum. The Cartesian components of the momentum are used for regression, as they have roughly Gaussian distributions compared to the commonly employed cylindrical coordinate representation. Although regressing both the mass and energy for each parton over-defines the 4-momentum, these components exhibit different reconstruction characteristics due to sharp peaks in the mass distributions. During evaluation, either the mass or energy can be used to compute any derived quantities. In our experiments, the regressed mass is only used for the mass reconstruction, and the predicted energy is used for other kinematics.
Detector Observations The detector-level jets and leptons are used as the conditioning data. The jets are stored as variable-length sets of momentum vectors with a maximum of 20 jets in each
Figure 3: Highlighted reconstruction components. The top row presents the full global histogram while the lower plot presents the ratio between the predicted histogram and the truth. Notice the improved mass shape compared to the pre-trained and non-latent models.
event. This study is limited to semi-leptonic  $t\bar{t}$  events, so each event is guaranteed to have a single lepton. The missing transverse momentum in each event (MET) is also computed and included in the conditioning. The jets and leptons are represented using both polar,  $(M,p_{\mathrm{T}},\phi ,\eta)$ , and Cartesian,  $(E,p_{x},p_{y},p_{z})$ , representations. We also include a one-hot particle identity, encoding either  $\mu$  or  $e$  for the lepton, or  $b$  or non-  $b$  for the jets as estimated by the  $b$ -tagger, resulting in 12 dimensions for each jet.
Training Networks were trained using the MSE for the reconstruction and noise loss, along with the physics-informed consistency loss with a weight of  $\lambda_C = 0.1$ . Each model underwent training for 24 hours using four NVIDIA RTX 3090 GPUs, resulting in 500,000 to 1,000,000 gradient steps for each model. Models were trained until convergence and then fine-tuned with a smaller learning rate. Full hyperparameters are presented in Appendix B.
Diffusion Sampling Variational diffusion models dynamically adapt the noise schedule during training by minimizing the variance of the ELBO [41]. After training, however, VDMs may employ a more traditional discrete noise schedule, and this approach is preferable when sampling for inference. The PNDM [58] sampler is used for generating parton predictions.
Global Distributions Each trained model was evaluated on the testing data, sampling a single parton configuration for each detector-level event. The global distributions of the 55 reconstructed parton components were then compared to the true distributions. Complete unfolded distributions are presented in Appendix F. Several highlighted reconstruction distributions are presented in Figure 3. Additionally, each model was assessed using several distribution-free measures of distance. The bin-independent Wasserstein and Energy distances, the non-parametric Kolmogorov-Smirnov (K-S) test, as well as three different empirical KL divergence measures using 64, 128, and 256 bins, are presented in Table ??. Full details about the distance functions are presented in Appendix C, and full tables of the distances per particle and per component are presented in Appendices D and E.
Global Performance The two proposed VLD models with unconditional decoders (VLD and UC-VLD) consistently exhibited the best performance across all distance metrics. The end-to-end training procedure demonstrates improved performance over the pre-trained LDM model. It is interesting to note that UC-VLD has lower distribution-free distance whereas VLD has a lower histogram distance. This is likely because the histogram distance will soften the effect of outliers in the distribution as the bins will combine many different points into a single less noisy measurements. The conditional decoder in C-VLD and CVAE was found to worsen reconstruction. This is likely because the training procedure always employs the true encoded parton-detector pairs,  $(z_{x}, z_{c})$ , whereas the inference procedure estimates the latent parton vector while using the true encoded detector variables for conditioning,  $(\hat{z}_{x}, z_{c})$ . The lower performance may be evidence that this inference data technically falls out-of-distribution for the conditional decoder, indicating that an unconditional decoder is a more robust approach. Finally, we find that all latent models greatly outperformed the models that directly reconstructed the partons in data-space (CINN and VDM).
Figure 4: Highlighted reconstruction per-event posteriors for several events and components. We compare the LVD posteriors to an empirically brute-forced estimate of the posterior.
Posterior Predictions One key innovation of generative methods is the ability to sample from the posterior to illustrate the space of valid reconstructed partons for a detector level event. While the true posterior is not available, the unfolded distributions can be compared to a brute-force posterior distribution derived from the training data. This posterior is defined by a re-weighting of the parton level training data, where the weights are given by the inverse exponential of the  $L_{2}$  distance between the testing event's detector configuration,  $y_{T}$ , and every training event's detector configuration,  $y_{i}$ :  $w_{i} = e^{-\|y_{T} - y_{i}\|}$ . Selected posterior distributions are presented in Figure 4, and complete posterior examples for individual events are presented in Appendix G. The latent diffusion models have much smoother posteriors than the empirical estimates, with the proposed VLD model producing more density close to the true parton configuration. Critically, the brute-force posterior often matches the unconditional parton level distribution, proving it is difficult to recover the true posterior. We also note that the VLD model was also able to reproduce a bimodal neutrino  $\eta$  posterior. Neutrinos are not directly measurable at the detector, and their kinematics must be inferred from the event's missing energy. For events with a single neutrino, the missing neutrino energy defines a quadratic constraint on the term which often leads to two configurations satisfying both energy and momentum conservation. The network appears to learn this phenomenon and presents two likely  $\eta$  values for the neutrino.
This paper introduced a novel extension to variational diffusion models, incorporating elements from latent diffusion models to construct a powerful end-to-end latent variational generative model. An array of generative models were used to unfold semi-leptonic  $t\bar{t}$  events, an important inverse problem in high-energy physics. A unified model — combining latent representations, continuous variational diffusion, and detector conditioning — offered considerable advantages over the individual application of each technique. This addresses the challenge of scaling generative unfolding methods for high-dimensional inverse problems, an important step towards unfolding full collision events at particle-level. Despite being tested on a single topology, our method consistently improved baseline results, underscoring the importance of latent methods for such high-dimensional inverse problems. The framework presented may be broadly applicable to arbitrary topologies, although always limit to a single topology at a time. Future work will focus on broadening the method's applicability to different event topologies, unfolding to other stages of the event simulation chain (such as "particle level") to remove dependence on event topology, and evaluating its dependency on the simulator's prior distribution. The methods described in this study aim to provide a general end-to-end variational model applicable to numerous high-dimensional inverse problems in the physical sciences.
We would like to thank Ta-Wei Ho and Hideki Okawa for assistance in generating the  $t\bar{t}$  sample used in this study. DW, KG, AG, and MF are supported by DOE grant DE-SC0009920, and AG is also supported under contract DE-AC02-05CH11231. The work of AS and PB in part supported by ARO grant 76649-CS to PB. We thank Vinicius Mikuni for fruitful discussions on unfolding methods.
[54] Alexander Shmakov, Michael James Fenton, Ta-Wei Ho, Shih-Chieh Hsu, Daniel Whiteson, and Pierre Baldi. SPANet: Generalized Permutationless Set Assignment for Particle Physics using Symmetry Preserving Attention. SciPost Phys., 12:178, 2022. doi: 10.21468/SciPostPhys.12.5.178.
[55] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. CoRR, abs/2201.03545, 2022.
[56] J. de Favreau, C. Delaere, P. Demin, A. Giammanco, V. Lemaitre, A. Mertens, and M. Selvaggi. DELPHES 3, A modular framework for fast simulation of a generic collider experiment. JHEP, 02:057, 2014. doi: 10.1007/JHEP02(2014)057.
[57] Matteo Cacciari, Gavin P. Salam, and Gregory Soyez. The Anti- $k_{t}$  Jet Clustering Algorithm. JHEP, 04:063, 2008. doi: 10.1088/1126-6708/2008/04/063.
[58] Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. In International Conference on Learning Representations, 2022.
[59] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017.
[60] Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR, abs/1606.08415, 2016.
[61] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
[62] Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Caglar Gülçehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learning. CoRR, abs/1910.06764, 2019.
[63] Michael James Fenton, Alexander Shmakov, Ta-Wei Ho, Shih-Chieh Hsu, Daniel Whiteson, and Pierre Baldi. Permutationless Many-jet Event Reconstruction with Symmetry Preserving Attention Networks. Phys. Rev. D, 105(11):112008, 2022. doi: 10.1103/PhysRevD.105.112008.
[64] Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in neural information processing systems, 33:19667-19679, 2020.
Figure 5: Block diagram from the gated inverse-bottleneck feed forward block used for the networks.
A Network Architecture
A.1 Feed-Forward Block
A custom feed-forward block, derived from the successful ConvNeXt [55] and transformer models [59], is the foundation of the networks in this study. This block features an inverted bottleneck formed by three linear layers, a GELU activation [60], layer normalization [61] for regularization, and a gated residual connection inspired by GTRXL [62]. The block uses a GRU recurrent layer instead of a traditional skip connection, where the regular output is the input to the GRU and the skip-connection value is the hidden input. A complete block diagram of the feed-forward block is presented in Figure 5.
A.2 Detector Encoder
The detector encoder network, as defined in SPANet [63, 54], processes variable-length sets of detector observations into a fixed-size  $D$ -dimensional vector. It uses the gated transformer architecture from SPANet version 2, with  $N_{S}$  transformer encoder blocks. Instead of utilizing the tensor attention layers, an event-level representation of the detector observations is extracted from the central transformer encoder.
A.3 Parton Encoder
The encoder starts with an embedding layer, transforming the fixed-size 55 dimensional parton representation into a  $D$ -dimensional vector via a linear layer. The encoder's body comprises  $N_{E}$  feed-forward blocks arranged in series. The input can be the embedded  $D$ -dimensional parton vector, or its concatenation with the  $D$ -dimensional encoded detector data. Two independent networks, each accepting identical input and sharing the same block structure, predict the mean,  $\mu_{\theta}(x;z_y)$ , and the log standard deviation,  $\log \sigma_{\theta}(x;z_y)$ , respectively. Normalizing and then scaling the mean also helped prevent the encoder from learning very small-valued components:  $\frac{\sqrt{D}}{\|\mu_{\theta}(x;z_y)\|} \mu_{\theta}(x;z_y)$
A.4 Parton Decoder
The decoder retains the encoder's linear block structure. As a deterministic decoder, it comprises a single stack of  $N_{D}$  feed-forward blocks and a concluding linear layer mapping  $D$  dimensions back to 55. Conditional decoders also append the  $D$ -dimensional detector vector to the input before processing.
A.5 Denoising Network
The denoising network, employing the same linear block structure, consists of  $N_{\epsilon}$  feed-forward blocks. It maps the  $D$ -dimensional latent sample  $z_{t}$  to the approximate noise that produced it,  $\epsilon_{\theta}(z_t,t,z_y)$ . The time,  $t\in [0,1]$ , is encoded with a 32-dimensional sinusoidal position encoding as detailed in [59]. The latent vector, position encoding, and conditioning are concatenated and fed through the feed-forward network to produce the noise estimate.
We present a full table of hyperparameters used throughout the experiments. All models use the same set of parameters, ignoring any that do not apply to particle methods. Parameters were not tuned using rigorous search. The detector transformer parameters were extract from experiments presented in SPANet [54], and the other networks where tuned to contain a similar number of parameters as the detector encoder.
Parameter	Value
Latent Dimensionality (D)	96
Attention Heads	4
Inverse Bottleneck Expansion (k)	2
Detector Transformer Encoder Layers (NS)	8
Parton Encoder Blocks (NE)	6
Parton Decoder Blocks (ND)	6
Denoising Network Blocks (NE)	10
Primary Learning Rate	5·10-4
Fine-tuning Learning Rate	1·10-4
L2 Gradient Clipping Limit	1.0
Consistency Loss Scale (λC)	0.1
Batch Size (Per GPU)	4096
Table 2: Table of complete hyperparameters used for training all generative models
C Distance Metrics
As the parton global distributions do not have a known family of distributions to describe their components, model-free measures of distribution distance must be used to evaluate the models. Three different families of distance measures are used. These non-parametric distances are only defined for 1-dimensional distributions. As there is no commonly accepted way of measuring distance for  $N$ -dimensional distributions, the 1-dimensional distances are simply summed across the components. Although not ideal, it is enough to compare different models and rank them based on performance.
C.1 Wasserstein Distance
The Wasserstein distance, often referred to as the earth-mover distance, quantifies the amount of work it takes to reshape one distribution into another. This concept originated from the field of optimal transport and has found wide applications in many areas, including machine learning. An equivalent definition defines this distance as the minimum cost to move and transform the mass of one distribution to match another distribution. For a pair of 1-dimensional distribution samples, denoted  $u$  and  $v$ , the Wasserstein distance can be computed in a bin-independent manner. This is achieved by computing the integral of the absolute difference between their empirical cumulative distribution functions (CDFs),  $U(x)$  and  $V(x)$ .
$$
D _ {\text {W a s s e r s t e i n}} (u, v) = \int_ {- \infty} ^ {\infty} | U (x) - V (x) | d x
$$
C.2 Energy Distance
Energy distance is another statistical measure used to quantify the difference between two probability distributions based on empirical CDFs. It compares the expected distance between random variables
drawn from the same distribution (intra-distribution) with the expected distance between random variables drawn from different distributions (inter-distribution). The Energy distance may be defined as the squared variant of the Wasserstein distance.
$$
D _ {\mathrm {E n e r g y}} (u, v) = \sqrt {2 \int_ {- \infty} ^ {\infty} \left(U (x) - V (x)\right) ^ {2} d x}
$$
C.3 Kolmogorov-Smirnov Test
The two-sample Kolmogorov-Smirnov (K-S) test is a non-parametric statistical hypothesis test used to compare the underlying probability distributions of two independent samples. It is particularly useful in machine learning applications where the goal is to assess whether two datasets come from the same distribution or if they differ significantly, without making any assumptions about the underlying distribution shape. It is also based on empirical CDFs.
C.4 KL-Divergence
An alternative approach to empirical CDF approaches is to bin the data into histograms and compute discrete distribution distances from these histograms. The common Kullback-Leibler distance is used with three different bin sizes. After finding the histograms with  $N$  bins for  $1 \leq i \leq N$ ,  $P_N(i)$  and  $Q_N(i)$ , the discrete KL divergence is computed as