End-To-End Latent Variational Diffusion Models for Inverse Problems in High Energy Physics
Authors: Alexander Shmakov
Kevin Greif
Michael James Fenton
Aishik Ghosh
Pierre Baldi
Daniel Whiteson
Venue: NeurIPS 2023 poster (2023)
Keywords: Diffusion
Variational
VAE
LDM
Physics
Unfolding
TL;DR: We introduce a novel end-to-end Variational Latent Diffusion (VLD) model, unifying variational and latent diffusion models, for a high-dimensional generative inverse problem in high-energy physics.

Abstract: 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.

End-To-End Latent Variational Diffusion Models for Inverse Problems in High Energy Physics
Alexander Shmakov
Kevin Greif
Department of Computer Science University of California Irvine
Department of Physics and Astronomy University of California Irvine
Irvine, CA 92697
Irvine, CA 92697
ashmakovQuci.ed
kgreif@uci... edu
Michael Fenton
Aishik Ghosh
Department of Physics and Astronomy University California Irvine
Department of Physics and Astronomy University of California Irvine
Irvine, CA 92697 mjfentonnuuii. edu
Irvine, CA 92697
Physics Division
I Lawrence Berkeley National Laboratory Berkeley, CA 94720
Pierre Baldi
Daniel Whiteson
Department of Computer Science University of California Irvine Irvine, CA 92697 pfbaldi@uci (elidi@uci..m edu
Department of Physics and Astronomy University of California Irvine Irvine, CA 92697
Abstract
High-energy collisions at the Large Hadron Collider (LHC) provide valuable in- sights 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 mapping detector observations theoretical quantities of the underlying collision are essential parts of many physics analyses at the LHC. We investigate and com- pare various generative deep learning methods approximate this inverse mapping. We introduce novel unified architecture, termed latent variational diffusion mod- els, 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 distribution-fie distance the truth of over 20 times smaller than non-latent state-of of aef-the-art art baseline and 3 times smaller than traditional latent diffusion models.
Introduction
Particle physics experiments at the Large Hadron Collider study the interactions of particles 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
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

of particles with high fidelity, no detector has perfect efficiency and resolution. A common strategy to account for these effects simulation-base inference [. in which the detector resolution and inefficiency are modeled by 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 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 result, there is no true inverse function which can map a given detector observation 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.
single point in the truth-level
Unfolding methods such as Iterative D' Agosttni (2). Singular Value Decomposition [3] and TUn- fold [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 un-binned unfolding with the capacity to handle dishe-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 [1O [1] and variational auto-encoders 12. An alternative approach uses classification models as density estimators which learn to correct imprecise truth-level distributions with re- re-weighting [3 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-evnn 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 [6 to map detector to truth-level distributions. Diffusion models are a class of generative models which learn to approximate 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 are-trained on secondary objectives, such VAE reconstruction tasks 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 develop end-to-end variational latent diffusion model (VLD) achieving state-of-the-ate performance in complex HEP generative tasks.
2 Background
2.1 Unfolding
Let Sam(y) be the distribution which governs an observed detector-level data set y {v: ). An unfolding method aims to sample from a pre-detector distribution Spanto (x), where parton refers to unobserved state of interest physicists. f fpatoo (x) is related to fāct via convolution with "response" function p(y|æ) over the possible true values æ. The response function describes the decay of the initial, unstable particles into stable particles and their interaction with the detector.
[v)(1)) = = f dx ((y))x)) (x)
(I)
In other fields, this kind of problem is often referred to as "deconvolution".

---

W
9
g rebilee
u
W
d
PATLAS
(a) Feynman diagram of top quark pair (tž) produc- tion, Each top decays to W boson and a bottom (anti-)quar) (b). In this example, one W decays leptonically an electron (e and anti-ncutrino (v), the second decays hadronically and down (e)) anti-quark. We store the momentum of all final and intermediate particles.
(b) A example display the high-dimensiona. de- tector observations for an L.HC collision identified as likely to have contained a top quark pair. We show 3D representation of the ATLAS detector and the hadronic jets which produced by the decays. detector-leve data consists of the mo- mentum measurements of these jets.
up (u)
Visual representations of the different perspectives captured by the parton and detector level data. The parton-level data represents 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.
Figure 1:
No closed form expression exists for p(y(z), but Monte-Carlo-bad simulation 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 pseudo-inversion of the response function p(xly), also known as the posterior.
partoon (x) - dy p(x
Generative unfolding methods build the posterior generative model, which can be used sample from p(zly). The desired parton distribution is then obtained by Equation Simulated pairs of parton-detector data, (2r, y). may be used to train the generative model.
important issue when choosing to directly model the posterior is that this quantity itself dependent on the desired distribution fpatton (z),, the prior in Bayes' theorem:
p(zly) p(u)==))
fae(3))
Producing the data used train the generative model requires choosing specific fparton (z). 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 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 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 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-leptone decay mode, Fig. la The 4-momenta (three momentum components, mass) of
3

---

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 trace. The resulting observed detector signature which defines the detector-level space is then quite complex, see Fig. lb
The semi- tł process has been studied by the ATLAS and CMS collaborations measure various properties of the top quark and to search for new particles and interactions (30H35). Many of these 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.
existing unfolding techniques, which limit the unfolded measurements
2.3 Variational Autoencoders
Variational Autoencoders (VAEs) are class of generative models combining an autoencoder architec- ture with probabilistic modeling [36 377. 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 2 conditioned on the input data: (=][r). The decoder network likewise models the generative distribution conditioned on the latent variables p(z[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 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).
VAA = Ez=_(=/>) |- log p(x|=) + DKL(a(z/z)) || p(:))
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.
Conditional VAEs (CVAEs)
LeVAE Ezoq(=|z.n) log v) (2z(a(2) ((+l)))
(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 [0, 1] over the latent space 2, E Z and conditioned on as:
q(2, Jr, v) N(axI, O.I)
(6)
The flow parameters, and aa are defined by a noise schedule. We use the continuous Variance Preserving (VP) framework throughout this work and derive these flow parameters based learned signal-to-noise ratio, e eo((), where:
0 sigmoid (t)) and or sigmoid((nvo())
Assuming it is possible to sample from the terminal distribution p(z1), we may produce samples from the data distribution by inverting the flow and sampling previous latent representations con- ditioned on future latent vectors. The inverse flow is modeled q(=,|z,, @o( I, w)) where áe an approximate denoising of the original data at the current time-step. In practice, the data de- noising is implemented using variance- inniunneeeddeeedeeettt noise prediction network, ég. by the equation (0 v, (=.=0))al).t..)) The noise prediction network, ég. is parameterized using deep neural network. The learnable noise schedule (o) is also parameterized using positive valued, monotonic neural network with learnable end- and-points 7min (0) and "maz 7(1) [4]] Following

---

Detector Encoder
y
Paston Encoder
x
20
23
Denolsing Model
2 ~
NO0,D)
X
Parton Decoder
C-VLD Only
Figure A block diagram of the end-to-end VLD models with trainable components. The conditional paths are drawn in different colors depending which model variations employ them. We use the continuous, variance preserving SDE diffusion formulation introduced [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(21) (O I). Both the noise prediction network and the noise schedule network are trained using the modified ELBO for continuous-tine diffusion models [41]
LVDM DKI |I p(z1)) + Bg(zol=) 1oggooppr//000o1 v)]
BB-N B. (2((() ||e Eo (.)))
2.5
Latent Diffusion
Latent diffusion models (66 are a deep generative framework that operate the diffusion process in an abstract latent space learned by VAE sample high-dimensional data (D y), possibly conditioned on a secondary dataset (v)). This approach has proven dramatically suc- cessful 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 low dimensional latent representation using a traditional VAE approach, ((>z/2)) and p(x[22), regularizing the latent space towards a standard normal p( (2z) (O, I). A secondary encoder may be trained on the conditioning data p(zyly) kyly) along-side VAE, typically using a CLIP objective (2 map the two datasets into a common latent space. The diffusion process is then trained to reconstruct the latents from the flow latents p( z,). 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 unified conditional generation objective. While several methods allow for training VAE along-side diffusion 42 these approaches either cannot train diffusion in the latent space cannot account for a conditional, fully variational model. We construct a unified variational framework allow for 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 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 conditional data encoder or decoder, and an auxiliary physics-iffrrmed consistency loss which may be enforced throughout the network. We refer this combined method 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 model.
(1yy)l)), Conditioning (Detector) Encoder In traditional LDMs, the conditioning encoder, p(=,, (v), is pre- trained through an auxiliary loss term, such as CLIP 277, which aims to unify the latent space the conditioning and data. While this approach is efficient, may not be optimal: the encoder is trained
5

---

on one objective, and then repurposed to act as 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 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, 2, fo(v). This decision is based on prior work such as the CVAR 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 unified conditional generative model in an end-to-end fashion, we have the option extend the encoder and decoder with conditional probabilistic models: QC-VLD(Z2 [r, 2y) and PC-vLD Raz Zy). 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, QUC 2z [2, zy), but employs an unconditional decoder during generation (zz.)1)(22
VLD ELBO We interpret the continuous VDM as an infinitely deep hierarchical VAE as presented by Kingma el al. [41]. This interpretation allows to seamlessly integrate the VAE into unified diffusion framework by incorporating the VAE as an additional component in the hierarchy. Conse- quently, the hierarchical variational BLBO incorporates
extra KI. divergence term, which serves to regularize the encoder posterior distribution [4]. We combine this hierarchical objective with the denoising loss term derived in [41] define a combined ELBO for the entire generative model.
LVLD = DKI(1(=122 |z. my) || p(=1)) + (e (m.ay) [- log p(r) Rx, Fy)]
DDLL((<z x z,) || p( (20)) + (0.) (1((t) ||e E0(2,, 2y)1 (8)
The additional KL term may be derived explicitly if we assume Gaussian VAE and a Gaussian diffusion process. The posterior parameterized using learned Gaussian, as in a standard VAE: ((zz [x, zy) ~ (Nu "y), op 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 20 in terms of 2g as aguo + doe, where ~ (O I). Solving this equation for 2g yields another reparameterized Gaussian, which allows us define the prior over 2g as;
p(zz|20) ~ N/ 20,
ag
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 p of a particle are related by 2²M = E² - (c !pll) We operate in natural units with unit speed of light = Forcing the predicted mass, energy, and momentum satisfy this equality improves stability and accuracy by capturing this underlying physical relationship between these quantities. We introduce a consistency loss, Lc, in addition to the regular reconstruction loss. weighted by hyper-parameter Àc Similar physics-informed constraints have previously been used for generative models HEP 45) 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.
Lc do Air
( -1=1)))))
(10)
4 Unfolding Semi-Leptonic tt Events
Generative models can be trained estimate conditional density given any set of paired data. In the unfolding context, Monte Carlo simulation can be used to generate pairs of events at detector and parton level. The density of parton level events fpurton (z) can be taken as the data distribution, and the density of detector level events faee (v) can be taken as the conditioning distribution. A generative model can then be used to unfold 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 ~ (z).). This can be done using matrix element solver such as MADGRAPH 49
2. Sample possible detector observation y ~ pc(y)) x)) using the tools PYTHIA8 [50 and DELPHES [S]], which simulate the interactions of particles in flight and the subsequent interactions with a detector.
3. Train generative model to approximate the inverse distribution (o(x)y)).
4. Produce new posterior samples for inference data with unknown parton configurations.
4.1 Generative Models
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 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 12 which represents the latest deep learning approach that has demonstrated success in unfolding tasks. This model utilizes conditional normalizing flow to train mapping from 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 conditional affine layer with global affine and permutation transforms to create powerful invertible block. In this work, the MMD objective defined in [52) is replaced with MSE reconstruction objective and the physics-informed consistency loss, for fair comparison with other models.
VDM A Variational Diffusion Model (VDM) [4] 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 Mean Squared Error loss against the generated noise.
LDM A Latent Diffusion Model (LDM) with a pre-trained VAE, popularized by recent achieve- ments in text-to-image generative models (26). The VAE is pre-trained using a Gaussian likelihood and minimal prior weight (10 (10-+). This baseline is meant to highlight the importance of the unified end-to-end architecture all other aspects of the network are identical to the proposed method.
VLD, C-VLD, UC-VLD These models are variations the proposed unified Variational Latent Diffusion (VLD) architecture. They correspond to an unconditional VAE (VLD), a conditional encoder and decoder (C-VLD), conditional encoder with 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, described in Section Additionally, the missing transverse momentum (MET) from the neutrino is included as fixed-size global variable. As there is no inherent ordering these jets, it is crucial to use a permutation-innaaian network architecture for the encoder. We use the jet transformer encoder from the SPANet (v2.1, 1, BSD-3) (54) jet-parton reconstruction network embed detector variables. This architecture leverages the permutation invariance of attention to contextually embed 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 (vyy) = Zy E D..
4.3 Parton Encoder-Decoder
For 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
7

---

Wasserstein
Energy
K-S
KT64
K
L256
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 Total distance measures across all 55 components for every model and metric. The indepen- dent sum of I-dimensional distances for each component are summed across all the components to compute the total metrics.
which consists of single 55-dimensional vector for each event. The encoder and decoder network employ a Con-NeXt-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, no (ar, zy). and ltg-standard deviation, og (2, 2, of the encoded vector, possibly conditioned on the detector observations Σy. The decoder similarly accepts a latent parton representation, possible conditioned on the detector, and produces a deterministic estimate of the original parton configuration & = D(2z; my).
Experiments
Dataset Each of the generative approaches is trained to unfold simulated semi-leptonic tÊ pro- duction data set. Matrix elements are evaluated a center-of ofasss mass energy of /s 13 TeV using MADGRAPH AMC MEDGRAAH_AMCO NLO [49 (v2.7.2, NCSA license) with a top mass of m, = 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 pair of light (u, d, s, c) quarks gg' lepton-neutring pair lv (6 = e). 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 PT > 25 GeV and absolute value of pscudorapidity Iy/ < 2.5. The b and light quarks are reconstructed with the anti-kr jet algorithm [57] using a radius parameter R 0.5 and the same Pr and by) requirements the leptons. Jets originating from b-quarks are identified with a "b-tagging" algorithm that incorporates a Pr and angular (1) Φ) dependent identification efficiency and mis mis-tagging rate. Selected events are then required to contain exactly one lepton and at least jets, of which at least two must be b-tagged. Events are separated into training and testing data sets, consisting of 9,865,4000 and 1,332,514 14 events respectively.
Parton Data The kinematics for the six final state partons are used as unfolding targets (6, 41 82. b, 1), along with the kinematics of the intermediate resonance particles (Wipp: Whas, t, f). and the entire lī system. The parton-level data consists of 11 momentum vectors, each represented by the five quantities (M, , log E,Px, Pr,P,, Pz); where M the invariant mass of the particle; the energy; and Pa. Py: and y> 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 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.
Py));
Detector Observations The detector-level jets and leptons are used as the conditioning data. The jets are stored variable-length sets of momentum vectors with a maximum of 20 jets in each
8

---

Ther UCALD LDN CINN
Histogram
1.2
0.5
9.5
170
160
1.00
1,85
200
Ratios
-.0
20 Leptonic Mass (GeV)
10
50
50 50
250 Leptanic Moss (GeV)
163
300
Lec Hadronic Mass (GoV)
100
200
100
200
Hadranic Mass (C4V)
Leptonic (GAVV)
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-trainn pre-trained and non-latent models.
event. This study is limited to semi-leptonic tÝ events, each event is guaranteed have single lepton. The missing transverse momentum each event (MET) is also computed and included in the conditioning. The jets and leptons represented using both polar, (M, Pr, Q, 17). and Cartesian, (E, Pa,Py. Ps), representations. We also include a one-hot particle identity, encoding either or for the lepton, or b non-b for the jets 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 À = 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 more traditional discrete noise schedule, and this approach 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 presented in Appendix |F] Several highlighted reconstruction distributions are presented in Figure Additionally, each model was assessed using several distribution -ree measures of distance. The bin bin-independent Wasserstein and Energy distances, the non parametric Kolmoogoov-Smmima test, as well as three different empirical KI. 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 Appendices and
(K-S)
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 interesting to note that UC-VLD has lower distribution-fiee distance wheras VLD has 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 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, (2z. 2c), whereas the inference procedure estimates the latent parton vector while using the true encoded detector variables for conditioning, (za, zz). 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 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).

---

Event 0
Event 0
Event 1
Event 2
Event 3
Leptonic Ms
Hadronic
Leptonic
Hadronica
System
(Gev)
Mass (GeV)
Energy (GeV)
Energy (GeV)
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 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 distribution derived from the training data. This posterior is defined by a re- re-weighting of the parton level training data, where the weights given by the inverse exponential of the Lg distance between the testing event's detector configuration, yr, and every training event's ' detector configuration, ys: w; = yr 177---11 a-ll, Selected posterior distributions are presented in Figure 4 and complete posterior examples for individual events 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 is difficult to recover the true posterior. We also note that the VLD model was also able to reproduce bimodal neutrino 7) posterior. Neutrinos are not directly measurable at the detector, and their kinematics must be inferred from the event's missing energy. For events with single neutrino, the missing neutrino energy defines 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 values for the neutrino.
brute-force posterior
6 Conclusions
This paper introduced novel extension to variational diffusion models, incorporating elements from latent diffusion models to construct powerful end-to-end latent variational generative model. array of generative models were used to unfold semi-leptonic 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 method consistently improved bascline
particle-level. Despite being tested single topology, results, underscoring the importance of latent methods for such high- 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 's applicability to different event topologies, unfolding 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.
7 Acknowledgements
We would like to thank Ta-We Wei Ho and Hideki Okawa for assistance in generating the tł sample used in this study. DW, KG, AG, and MF are supported by DOE grant DE- DE-SC0009920, and AG is also supported under contract DE-ACO2-0SCHI122 grant 76649-CS to PB. We thank Vinicius Mikuni for fruitful discussions on unfolding methods.
The work AS and PB in part supported by ARO
10

---

References
[1] Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The Frontier of Simulation- Based Inference.
Proc. Nat. Acad. Sci., 117748):30055-30063 2020. doi: 10.1073/pnas. 19 19178891
[2] G. D'Agostini. A Multidimensional Unfolding Method based on Bayes' Theorem. Nucl. Instrum, Meth., A362:487-498, 1995. doi:
[3] Andreas Hocker and Vakhtang Kartvelishvili. SVD Approach to Data Unfolding. Nucl. Instrum. Meth., A372:469-44 481, 1996. doi: 10.1016/0168 10.101660066-9900((5))117733
[4] Stefan Schmitt. TUnfold: an Algorithm for Correcting Migration Effects in High Energy Physics. JINST, 7:T10003, 2012. doi: 00.1088/778--02211/0000100
[5] The ATLAS Collaboration. The ATLAS experiment the CERN large hadron collider. Journal of Instrumentation, 3(08):508003--08000, aug 2008. doi: 10.
[6] The CMS Collaboration. The CMS Experiment at the CERN LHC. JINST, 3:S08004, 2008. doi:
[7] P. Baldi. Deep Learning in Science. Cambridge University Press, Cambridge, UK, 2021.
[8] Kaustuv Datta, Deepak Kar, and Debarati Roy. Unfolding with Generative Adversarial Networks. 2018.
[9] Marco Bellagente, Anja Butter, Gregor Kasieczka, Tilman Plchn, and Ramon Winterhalder. How to GAN away Detector Effects. SciPost Phys., 8(4):070, 2020. doi: 10.21468/SiPostPhyy. 8.4.070.
[10] Marco Bellagente, Anja Butter, Gregor Kasieczka, Tilman Plehn, Armand Rousselot, Ramon Winterhalder, Lynton Ardizzone, and Ullrich Köthe. Invertible Networks or Partons to Detector and Back Again. SciPost Phys., 9:074, 2020. doi: 09.21666/iiiossstyy.55......
[11] Mathias Backes, Anja Butter, Monica Dunford, and Bogdan Malaescu. An unfolding method based on conditional Invertible Neural Networks (cINN) using iterative training. 12 2022.
[12] Jessica N. Howard, Stephan Mandt, Daniel Whiteson, and Yibo Yang. Learning to Simulate High Energy Particle Collisions from Unlabeled Data. Sci. Rep., 12:7567, 2022. doi: 10.1038/ s41598-022- 10966-7.
[13] Anders Andreassen, Patrick T. Komiske, Eric M. Metodiev, Benjamin Nachman, and Jesse Thaler. OmniFold: Method Simultancously Unfold All Observables. Phys. Rev. Lett., 124 1180:182001, 2020. doi: 10.1 10.1100000011 110//hysRewL.et..2222..
[14] Patrick Komiske, W. Patrick McCormack, and Benjamin Nachman. Preserving New Physics while Simultancously Unfolding all Observables. Phys. Rev. D, 104(7):076027, 2021. doi: 10.1 1e..10//yyyRee..00..766
[15] Jay Chan and Benjamin Nachman. Unbinned Profiled Unfolding. 2023.
[16] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[17] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. 2021.
[18] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[19] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[20] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents, ArXiv, abs/2204.06125 2022.

---

[21] Ilia Igashov, Hannes Stärk, Clement Vignac, Victor Garcia Satorras, Pascal Frossard, Max Welling, Michael M. Bronstein, and Bruno Correia. Equivariant 3d-conditional diffusion models for molecular linker design. In NeurlPS 2022 Al for Science: Progress and Promises, 2022.
[22] Vinicius Mikuni and Benjamin Nachman. Score-based generative models for calorimeter shower simulation. Phys. Rev. D, 106(9):092009, 2022. doi: 10. 10.100//yyyee..... DD0SPPhyyRRVD.10.69920000
[23] Matthew Leigh, Debajyoti Sengupta, Guillaume Quétant, John Andrew Raine, Knut Zoch, and Tobias Golling. PC-JeDi: Diffusion for Particle Cloud Generation in High Energy Physics. 2023.
[24] Vinicius Mikuni, Benjamin Nachman, and Mariel Pettee. Fast Point Cloud Generation with Diffusion Models in High Energy Physics. 4 2023.
[25] Erik Buhmann, Sascha Diefenbacher, Engin Eren, Frank Gaede, Gregor Kasieczka, Anatolii Korol, William Korcari, Katja Krüger, and Peter McKeown. CaloClouds: Fast Geometry- Independent Highly-Grandne Granular Calorimeter Simulation. 5 2023.
[26] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models, 2021.
[27] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning,
pages 8748-8788 8763.87633777 PMLR, 2021.
[28] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 14:21696-21701 2021.
[29] Anja Butter, Theo Heimel, Sander Hummerich, Tobias Krebs, Tilman Plehn, Armand Rousselot, and Sophia Vent. Generative Networks for Precision Enthusiasts. 10 2021.
[30] ATLAS Collaboration. Measurement of the top quark mass in the -> lepton+jett channel from / 8 TeV ATLAS data and combination with previous results. Eur. Phys. J. C, 79(4): 290, 2019, doi: 10.
[31] ATLAS Collaboration. Measurements of top-quark pair differential and double-differential cross-sections in the É+jets channel with pp collisions at vs = 13 TeV using the ATLAS detector. Eur. Phys. J. C, 19(2):1028, 2019, doi: 10. 1001 Eur.Phys.J.C 80, 1092 (2020)).
[Erratum:
[32] ATLAS Collaboration. Search for heavy particles decaying into top-quark pairs using lepton- plus-jets events in proton-proton collisions / 13 TeV with the ATLAS detector. Eur. Phys. J. C, 78())5665, 2018. doi: 10.1140/epjc/s 01052-018-599---1
[33] CMS Collaboration. Measurement of differential tī production cross sections using top quarks at large transverse momenta in pp collisions at /a = 13 TeV. Phys. Rev. D. 1035500050008, 2021. doi: 10. 10.1103/PhyyReeD.113..2222
[34] CMS Collaboration. Measurement of differential tł production cross sections in the full kinematic range using lepton+jets events from proton-proton collisions at v = 13 TeV. Phys. doi: 10. 10.110//Physeee...44.... 104.99013.
Rev. D, 104())092003,, 2021.
[35] CMS Collaboration. Search for dark matter produced in association with single top quark or a top quark pair in proton-proton collisions v 13 TeV. JHEP. 03:141, 2019. doi: 10. (0.100///HEE033(011914
[36] Diederik P Kingma and Max Welling. Auto-encoding variational bayes, 2022.
[37] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Eric P. Xing and Tony Jebara, editors, Proceedings of the 31st International Conference on Machine Learning, pages 1278 1286. Bejing, China, 22-24 Jun 2014. PMLR.
12

---

[38] Xun Huang, Ming- Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised image- to-image translation. In Proceedings of the European conference on computer vision (ECCV), pages 172-189, 2018.
[39] Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz. and Samy Bengio. Generating sentences from a continuous space. arXiv preprint arXiv: 1e1.1/5111..999 2015.
[40] Kibyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28, 2015.
[41] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. On density estimation with diffusion models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances Neural Information Processing Systems, 2021.
[42] Kushagra Pandey, Avideep Mukherjee, Piyush Rai, and Abhishek Kumar. Diffuse Efficient, controllable and high-fidelity generation from low-dimensional latents. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
[43] Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[44] Casper Kaae SI nderby, Tapani Raiko, Lars Maalø Søren Kaae So nderby, and Ole Winther. Ladder variational autoencoders. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, 2016.
[45] Michela Paganini, Luke de Oliveira, and Benjamin Nachman. CaloGAN Simulating 3D High Energy Particle Showers in Multilayer Electromagnetic Calorimeters with Generative Adversarial Networks. Phys. Rev. D, 97(1):01402,, 2018. doi: 10. 10w.101yseee...7.000
[46] ATLAS Collaboration. Deep Generative Models for Fast Photon Shower Simulation in ATLAS. 10 2022.
[47] Simon Badger et al. Machine Learning and LHC Event Generation. SciPost Phys., 14:079, 2023. doi: 10.21 1wt////pPosshhyy..4...0.99999999
[48] Yadong Lu, Julian Collado, Daniel Whiteson, and Pierre Baldi. Sparse autoregressive models for scalable generation of sparse images in particle physies. Physical Review D, 103(3):03002,, 2021.
[49] Alwall, R. Frederix, S. Frixione, V. Hirschi, F. Maltoni, O. Mattelaer, H. Shao, T. Stelzer, P. Torrielli, and M. Zaro. The Automated Computation of Tree-level and Next-to-leading Order Differential Cross Sections, and Their Matching to Parton Shower Simulations. JHEP, 07:079, 2014. doi: 10. 10.10077/HEE077(2140077
[50] Torbjörn Sjöstrand, Stefan Ask, Jesper R. Christiansen, Richard Corke, Nishita Desai, Philip Ilten, Stephen Mrenna, Stefan Prestel, Christine O. Rasmussen, and Peter Z. Skands. An introduction to PYTHIA 8.2. Comput. Phys. Commun., 191:159-177, 2015. doi: 10.1016/j..pe. 2015.01.024.
[51] Favereau, C. Delaere, P. Demin, A. Giammanco, V. Lemaître, A. Mertens, and M. Selvaggi. DELPHES 3, A Modular Framework for Fast Simulation of Generic Collider Experiment. JHEP, 02:057, 2014. doi: 10.1007//HEP0222014007...
[52] Marco Bellagente, Anja Butter, Gregor Kasieczka, Tilman Plehn, Armand Rousselot, Ramon Winterhalder, Lynton Ardizzone, and Ullrich Köthe. Invertible networks or partons to detector and back again. SciPost Phys., 9:074, 2020. doi: 10.21468/saPosPPyy...555
[53] Lynton Ardizzone, Till Bungert, Felix Draxler, Ullrich Köthe, Jakob Kruse, Robert Schmier, and Peter Sorrenson. Framework for Easily Invertible Architectures (FrEIA), 2018-2022.
13

---

[54] Alexander Shmakov, Michael James Fenton, Ta- 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.2 00.21468/sc/oosthys. 12.5.178.
[55] Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. CoRR, 0ass/2001.0354, 2022.
[56] J. de Favereau, C. Delaere, P. Demin, A. Giammanco, V. Lemaître, A. Mertens, and M. Selvaggi. DELPHES 3, A modular framework for fast simulation of a generic collider experiment. JHEEP 02:057, 2014. doi: 10. 1007//HHEP02(221400557
[57] Matteo Cacciari, Gavin P. Salam, and Gregory Soyez. The Anti-k; Jet Clustering Algorithm. JHEP, 04:063, 2008. doi: 10.1088/1 126 120 670660000700000//06633
[58) Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. International Conference on Learning Representations, 2022.
[59] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan 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//100..08.... 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, Çaglar Gülçehree 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/ 2019.
[63] Michael James Fenton, Alexander Shmakov, Ta-Wei Ho, Shih-Chich Hsu, Daniel Whiteson, and Pierre Baldi. Permutationless Many-jet Event Reconstruction with Symmetry Preserving Attention Networks. Phys. Rev. D, 105(11):11 1008,, 2022. doi: 10. 1..100/yyseeDD.005..2200.
[64] Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in neural information processing systems, 13:19667-19679, 2020.
14

---

ED-Upscale Linear
D Dimension
Layer Norm
D-Linear
Gelu
Linear
GRU
Skip Connection
Figure 5: Block diagram from the gated invereebottleenck 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], the foundation the networks this study. This block features an inverted bottleneck formed by three linear layers, GELU activation [60). layer normalization [61] for regularization, and gated residual connection inspired by GTRXL [62]. The block uses GRU recurrent layer instead of 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
S
A.2 Detector Encoder
The detector encoder network, as defined in SPANet observations into a fixed-size D-dimensional vector. It uses the gated transformer architecture from SPANet version 2, with Ns 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.
63 54
processes variable-length sets of detector
A.3 Parton Encoder
The encoder starts with an embedding layer, transforming the fixed-size 55 dimensional parton representation into D-dimensional vector via linear layer. The encoder's body comprises Ng feed-forward blocks arranged in series. The input can be the embedded D-dimensional parton vector, its concatenation with the D-ddmensional encoded detector data. Two independent networks, each accepting identical input and sharing the same block structure, predict the mean, μ4 (z; 2y), and the log standard deviation, log (2; 2y), respectively. Normalizing and then scaling the mean also /D helped prevent the encoder from learning very small-valued components: [| )|] (40 (2;; zy)
A.4 Parton Decoder
The decoder retains the encoder's linear block structure. a deterministic decoder, comprises single stack of Np feed-forward blocks and concluding linear layer mapping D dimensions back 55. Conditional decoders also append the D-dimensional detestor vector to the input before processing.
15

---

A.5 Denoising Network
The denoising network, employing the same linear block structure, consists N, feed-forward blocks. It maps the D-dimensional latent sample 2 to the approximate noise that produced it, eø t, zy). The time, E [0, 1], encoded with 32-dimensional sinusoidal position encoding as detailed in 59 The latent vector, position encoding, and conditioning are concatenated and fed through the foed-forward network produce the noise estimate.
B Hyperparameters
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 contain similar number of parameters 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 (Ng)
6
Parton Decoder Blocks (ND)
6
Denoising Network Blocks (N₂)
10
Primary Learning Rate
5 10
Fine-tuning Learning Rate
1 10-1
T.2 Gradient Clipping Limit
Consistency Loss Scale (Ac)
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 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 only defined for I-dimensional distributions. As there is no commonly accepted way of measuring distance for N-dimensional distributions, the I-dimensional distances are simply summed across the components. Although not idcal, it is enough to compare different models and rank them based on performance.
C.1 Wasserstein Distance
The Wasserstein distance, often referred to 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 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 (xx) and (ax).
D DWasserstin (u, (U(+) V(x)ddr
C.2 Energy Distance
Energy distance another statistical measure used to quantify the difference between two probability distributions based empericial CDFs. It compares the expected distance between random variables
16

---

drawn from the same distribution (intra-distriine distribution) with the expected distance between random variables drawn from different distributions .istribuuii......... The Energy distance may be defined as the squared variant of the Wasserstein distance.
/
/2 { (U(x) V(x))² dr
D Energy (u, u)
C.3 Kolmogorov-Smirnox Test
The two-sample Catmogorov-Smimeee (K-S) test is non- con-parametric statistical hypothesis used to compare the underlying probability distributions of two independent samples. It is particularly useful in machine learning applications where the goal is assess whether two datasets come from the same distribution or if they differ significantly, without making any assumptions about the underlying distribution shape. 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 s < N, (N) and Q (i), the discrete KL divergence computed as
N DKL,N = Σ PN(()
PN(2) (i)
D Particle Distance Tables
Tables 3 to present the distance metrics for each parton and model. The general trends in Table ?? remain generally consistent across partons. The neutrino reconstruction can prove difficult for Latent diffusion models, likely due to its very peaked components.
E
Component Distance Tables
Tables present the distance metrics for each component and model. The mass component seems to vary the most for for many of the distance functions, indicating that many models struggle reconstructing the peaked mass distributions. However, the overall results remain consistent with We again see the clear benefit of both latent diffusion and end-to-end training.
Table ??
F
Global Distribution Plots
present a collection of global distributions for the three primary classes of generative models for every particle and component. The proposed method (VLD) closely matches the truth distributions across all components, including the mass which slightly smoothed but peaks in the correct location. Baseline models struggle with capturing the peaks and shapes of the distributions.
Figures 6 through 16
G Posterior Distribution Plots
Figures and 18 present a collection posterior distributions for four testing events, along with several models and the brute-force empirical approach.
H Loss Function Derivation
We provide a derivation for the loss function presented Equation g adhering to the generative model displayed in Figure 2 Here, data is generated from latent representation, the latent
17

---

data from the diffusion end end-point, p0)=|z0); diffusion process to generate the sampled configuration p(zo|21) 21); ); and a simple standard normal distribution prior the base of the hierarchy p(21) ). All these
functions are further conditioned on the (embedded) detector observations Ey- We start with the full conditional hierarchical VAE. loss function over variational layers as per
L Σ
CLBOO = (g(==)1,,))) I- log p(x) &es 2y)1 DKLla(zikz By. B(41) ll p(z((ze))
(11)
Next, substitute the layers defined for the expand the expression for our generative model. For this step, we employ a three-stage hierarchy with the following substitutions: 21 < 21, 22 +- 20, and 23 &- Zz. We derive each of these components in the following sections.
LEBO = Bu( |- log p(z) zy)]
(12)
+ (x 2n) p(21))
(13)
+ ( 2, 21, ((20) 21)
(14)
+ Dx 21, 21, p( 21, 20))
(15)
H.1 Prior Loss
Equation 13 establishes the prior loss and the base layer in the hierarchy. In accordance with the VP framework, the correct prior distribution for the final latent representation is the standard normal p(21) ~ (0, I).. We learn the noise schedule via log SN R(t) = -%o(t) (t), as defined in VDM [41] As such, must ensure the terminal state in the forward diffusion process aligns with the prior distribution. Substituting the VP noise schedule yields the following distribution for the posterior:
ạ<2 jr., Zy: 2 [(1 a I), where a₂ = / (ammmas()) and on = vo (-(>(>(0)) (1)) as obtained
from Section 2.4
We estimate this KI. divergence through Monte- Carlo sampling:
Dxz(o(z)l, p(=)) = E2 (( [(012=) + of log(o?) 1]
(16)
H.2 VAE
15 delineates our contribution to this unified variational model. We derive the base distributions for the right-hand distribution in Equation 9 p(zrlzo) ~ N ag 20, do II). We fur- ther describe that the posterior distribution in this KL term is merely the regular VAE posterior (22 (r, Zn) ~ (Nuu (z, Zy), (2 2y)). We define the posterior over 20 given 2n by adhering to the definition of the VP diffusion process g(20 2xx) N(zz, ool). As all these distributions are normal, we can provide an explicit form for this loss:
Equation
as
DRI(I(ZZla) (z, Zy) (((2((21))
E Ex (ts.oglsll..y)..oooooo [.,.). D
DKL(N(uo(ss 11 N (_20, _
where the KL term the regular normal distribution KL.
(No m2) DNL(N(oo on) (V(1. = 2 o of
o
+ log of log o
H.3
Diffusion Loss
Equation 14 defines final diffusion loss term for the denoising network. We follow the derivation from Kingma et. al. [1 for continuous-time diffusion process. The key insight is to interpret the diffusion process as infinitely deep hierarchical variation model. The VP framework defines intermediate steps as: q y) ~a(a, ozIl), Following [1], we derive Monte-Carlo estimate
18

---

of the integral loss for noise prediction network:
[% SNR'(1) (1) ||2a (2t, Fyy 1)12 dll
DK 2gp
|| p(zo|z1)) B Be
2 1
SNN (t) || - 2ne (2e, 2y. t) 1)
2 B (((()) (t) ||e Eg(ze, 2zy 1)]2
H.4 Reconstruction Loss
Equation 12 is the final and simplest aspect of this model. Since regressing the parton components, we use a Normal decoder on the VAE and use simple Mean Squared Error loss as the primary reconstruction loss.
MSE [DECODER(Za; 2y) x
As explain in the text, this MSE loss works well for most components but fails capture the peaked nature of the mass term. Therefore, we add the physics consistency loss described by Equation 10 in order to assist with this mass reconstruction.
Lc Ac M (E² (E-- (1=-11=
The total reconstruction loss is simply the sum of these two components.
(x(zal |- log p(x]2222 zy)1 = Eg (m,,,) [AMSE + Lcl
VLD
UC-VLD
C-VLD
LDM
VDM
CVAE
CINN
Leptonic b
11.96
5.23
9.08
16.74
212.29
33.85
143.49
Leptonic Ve M4
10.33
7.40
36.69
29.63
117.37
36.95
149.22
Leptonic e μ
2.95
2.76
3.07
11.13
155.87
8.72
96.89
Hadronic
9.52
4.72
11.38
15.66
226.58
39.02
132.69
Hadronic Q1
8.88
7.19
24.26
35.33
187.24
58.33
180.42
Hadronic Q2
8.56
5.93
53.15
43.56
99.16
46.42
123.12
Leptonic
8.18
8.19
43.63
37.87
213.51
32.96
260.12
Hadronic W
7.55
6.94
39.97
50.28
215.10
53.85
308.33
Leptonic
9.57
9.12
46.56
43.27
323.00
46.51
408.39
Hadronic t
15.12
6.57
36.25
48.82
343.31
59.08
444.66
tt System
16.15
9.52
85.58
70.04
384.93
68.85
761.74
Table 3: Particle Distance: Wasserstein Distances
1.00
Leptonic
0.75
Truft
UEALD
0.5
LON
0.25
0.5
CRNN
0.00
0.0
0.0
Ratios
2.8
1.0
0.5
Mass (GeV)
Energy (GeV)
Pc (GeV)
Figure 6: Global Distribution: Leptonic b Quark
19

---

VLD
UC-VLD
C-VLD
LDM
VDM
CVAE
CINN
Leptonic b
0.93
0.59
0.62
1.01
17.04
2.53
10.33
Leptonic Ve M
0.79
0.73
2.31
1.76
11.02
2.72
12.15
Leptonic e / 0.24
0.29
0.26
0.70
13.99
0.59
7.56
Hadronic b
0.61
0.43
0.94
1.00
17.83
2.75
9.06
Hadronic
0.90
0.77
2.71
2.75
16.30
4.82
14.73
Hadronic 42
0.79
0.57
3,90
2.93
10.00
4.09
9.85
Leptonic W
0.54
0.66
2.94
2.37
16.48
2.09
17.58
Hadronic W
0.49
0.45
2.70
2.99
16.18
3.38
20.81
Leptonic t
0.64
0.64
2.61
2.42
21.72
2.52
23.27
Hadronic
0.86
0.44
2.03
2.43
22.16
3.32
24.81
tī System
0.80
0.78
4.36
3.72
18.63
3.47
34.98
Table 4: Partiele Distance: Energy Distances
VLD
UC-VLD
C-VLD
LDM
VDM
CVAE
CINN
Leptonic b
0.85
0.07
0.06
0.59
1.91
1.01
1.30
Leptonic Ve /
0.97
1.06
0.73
0.88
1.76
0.77
1.71
Leptonic /
0.48
0.33
0.52
0.35
1.47
0.43
1.02
Hadronic
0.05
0.05
0.10
1.06
1.95
0.90
1.30
Hadronic Q
0.40
0.40
0.60
0.63
1.63
0.81
1.50
Hadronic
0.90
1.00
0.83
0.76
1.64
0.95
1.59
Leptonic W
0.10
0.11
0.46
0.37
1.39
0.30
1.27
Hadronic W
0.07
0.09
0.42
0.42
1.36
0.32
1.46
Leptonic t
0.09
0.10
0.35
0.32
1.61
0.30
1.43
Hadronic
0.10
0.09
0.32
0.28
1.60
0.37
1.48
System
0.06
0.09
0.26
0.24
0.82
0.21
1.68
Table 5: Particle Distance: Kolmogorov- Test Statistics
VLD
UC-VLD
C-VLD
LDM
VDM
CVAE
CINN
Leptonic b
1.50
0.01
0.01
0.49
5.91
0.27
1.44
Leptonic Ve F
0.14
3,45
0.87
1.43
0.76
0.88
2.48
Leptonic
0.14
1.35
0.44
0.93
3.66
1.04
3.43
Hadronic
0.00
0.01
0.01
3.57
5.99
0.29
1.09
Hadronie 4
0.34
2.07
0.46
1.19
3.80
1.72
6.15
Hadronic 42
1.54
0.08
0.67
0.96
2.31
0.98
2.67
Leptonic W
0.02
0.03
2.94
2.59
2.39
0.69
2.13
Hadronic W
0.01
0.02
2.43
2.54
2.11
1.25
2.53
Leptonic L
0.02
0.03
1.22
1.47
2.50
0.98
2.24
Hadronic
0.03
0.04
0.99
1.14
2.39
1.05
2.45
Iī System
0.01
0.01
0.06
0.04
0.48
0.04
3.59
Table 6: Particle Distance: KL Divergence with 128 bins.
20

---

VLD
UC-VLD
C-VLD
LDM
VDM
CVAE
CINN
2.06
2.61
17.76
19.12
87.24
20.91
113.29
pt
7.18
7.65
39.39
33.44
429.70
74.56
257.20
0.37
0.26
0.47
0.40
0.44
0.51
4.35
phi
0.22
0.19
0.20
0.38
0.17
0.30
1.29
energy
27.33
16.26
140.80
152.24
772.94
155.16
1098.99
px
13,89
10,48
25.86
25.37
272.43
50.46
166.93
py
8.90
7.95
26.48
22.95
270.28
47.97
168.37
pz
48.80
28.16
138.66
148.42
645.15
134.69
1198.66
Table 7: Component Distance: Wasserstein Distances
VLD
UC-VLD
C-VLD
LDM
VDM
CVAB
CINN
0.51
0.76
3.77
3.50
15.30
3.00
9.70
pt
0.77
1.04
4.27
3.42
44.65
7.71
25.36
0.26
0.17
0.28
0.25
0.25
0.33
2.76
phi
0.15
0.13
0.14
0.25
0.11
0.21
0.86
energy
1.40
0.96
7.50
7.36
53.11
8.32
70.01
px
1.25
1.05
2.23
2.20
20.68
4.13
12.49
py
0.87
0.81
2.23
1.88
20.54
3.90
13.15
2.38
1.43
4.97
5.23
26.71
4.70
50.79
Table 8: Component Distance: Energy Distances
VLD
UC-VLD
C-VLD
LDM
VDM
CVAE
CINN
mass
3.39
2.78
3,20
4.54
7.55
4.21
4.66
pt
0.08
0.14
0.36
0.31
3.30
0.67
1.92
0.14
0.09
0.13
0.12
0.11
0.18
1.34
phi
0.07
0.06
0.07
0.12
0.05
0.10
0.39
energy
0.07
0.08
0.36
0.30
2.94
0.45
3.81
0.10
0.10
0.18
0.18
1.16
0.31
0.78
py 0.08
0.08
0.17
0.16
1.14
0.29
0.92
pz
0.15
0.09
0.18
0.17
0.88
0.17
1.91
Table 9: Component Distance: Kolmogorov-Smimon Test Statistics
VLD
UC-VLD
C-VLD
LDM
VDM
CVAE
CINN
mass
3.69
7.04
9.78
15.80
22.78
8.53
12.34
pt
0.01
0.02
0.07
0.07
3.08
0.22
1.66
eta
0.01
0.01
0.03
0.04
0.04
0.03
1.87
phi
0.00
0.00
0.00
0.01
0.00
0.01
0.07
energy
0.01
0.01
0.08
0.24
2.99
0.15
6.53
px
0.01
0.01
0.03
0.03
1.28
0.10
0.92
py
0.00
0.01
0.04
0.03
1.27
0.09
0.92
pz
0.01
0.00
0.05
0.13
0.86
0.05
5.89
Table 10: Component Distance: KI. Divergence with 128 bins.
21

---

Leptonic Ve Vv
10.2
2.0
2.0
1.5
1.0
100
0.5
0.5
CINN
0.5
0.0
0.
0.0
200
400
0.0
100
200
2
0.0
2
1.2
Ratios
0.5
1.0
2.0
0.0
400
Mass (GeV)
Energy (GeV)
Pt (GeV)
Figure 7: Global Distribution: Leptonic Neutrino V, U/
Leptonic e 1
2
1.0
1.0
LDH
0.5
CINN
0.00
0.0
200
400
0.0
0.0
1.2
Ratios
1.0
2.5
0.5
0.10
200
400
Mass (GeV)
Energy (GeV)
Pr (GeV)
Figure 8: Global Distribution: Leptonic Lepton e 1
Hadronic b
3.0
15
4
1.0
0.0
0.0
4.5
5.0
200
400
600
100
200
Ratios
1.1
1.0
0.5
0.0
200
400
800
100
200
Mass (GeV)
Energy (GeV)
P. (GeV)
Figure 9: Global Distribution: Hadronic b Quark
Hadronic dt.
2.0
1.5
2.0
1.5
1.0
3.0
1.0
LDN
0.5
0.5
CNN
3.0
1.0
1.5
0.0
200
400
0.0
200
0.0
Ratios
0.5
3.0
200
460
200
Mass (GeV)
Energy (GeV)
Pú (GeV)
Figure 10: Global Distribution: Hadronic Light Quark
22

---

A0-2
1.5
Hadronic
1.0
1.0
Weot CNN
0.5
0.5
0.5
-0.
0.0
200
400
0.3
0.0
Ratios
1.0
0.5
Mass (GeV)
Energy (GeV)
Pt (GeV)
Figure 11: Global Distribution: Hadronic Light Quark
10.
1c
10
1
10
1.5
Leptonic
Wi0 CNN
0.50
0.5
Ratios
1.1
1.0
1091
Mass (GeV)
Energy (GeV)
(GeV)
Figure 12: Global Distribution: Leptonic Boson
1
1
10:1
1.00
2.0
1.5
Hadronic
0.73
0.3
0.00
0.0
0.0
0.0
80
500
1.25
Ratios
1.5
1.0
1,00
0.5
0.75
Mass (GeV)
Energy (GeV)
Pt (GeV)
0
Figure 13: Global Distribution: Hadronic W Boson
Leptonic t
Truth UCALD
1.0
1.5
1.0
0.5
1.60
1770
180
0.0
500
1000
200
400
2
0.0
-2
1.1
Ratios
1.0
5.0
1000
600
Mass (GeV)
Energy (GeV)
P; (GeV)
Figure 14: Global Distribution: Leptonic Top Quark
23

---

10
Ic
Hadronic
LDN
1.0
CINN
1.0
0.5
160
170
180
500
1000
200
4.0
-2
Ratios
1.0
1.0
160
170
180
500
1000
200
400
Mass (GeV)
Energy (GeV)
(GeV)
Figure 15: Global Distribution: Hadronic Top Quark
t System
1.5
CINN
0.5
24
0.5
Ratios
1.00
500
1000
500
3000
1500
2000
200
Mass (GeV)
Energy (GeV)
(GeV)
Figure 16: Global Distribution: Complete System
24

---

(a) Example Event
(b) Example Event
Figure 17: Posterior distributions for example events. Included is an empirical posterior distribution
calculated from the training dataset.
25

---

(a) Example Event
(b) Example Event
Figure 18: Posterior distributions for example events. Included is an empirical posterior distribution
calculated from the training dataset.
26


Reviews:
[{"review_id": "aivaHrKbRy", "reviewer": "Reviewer_aaox", "summary": "Thanks for the author's rebuttal. I have read all of the rebuttals and reviews and decided to keep the current rating.\n\nCurrent work proposed an extension of the variational diffusion model, called the variational latent diffusion model. It is applied to inverse problems in the high energy physics field and tested on a single topology. The result shows that the distance is three times less than current latent diffusion models.", "questions": "1. There is no general caption for Figure 1. Moreover, the sub-caption is not consistent with the sub-figure. For example, ($e$) in caption while ($e^{-}$) in the subfigure. ($\\nu$) in caption while $\\bar{v}$ in subfigure. $d$ in caption while $\\bar{d}$ in subfigure. Moreover, subfigure (b) is too small to see the detail; the reviewer did not get the point of how it is related to your problem or the model. The font in this subfigure is too small to see. The review suggests presenting the subfigure (b) in a better way to give more insights into your problem or model.\n\n2. There are several typos in the paper. For example, line 8, ``latent variation diffusion model``, is supposed to be ``latent variational diffusion models`` Line 192, the $\\hat{\\cdot}$ is on $z_t$ or $t$?\n\n3. Some definitions are not clear. For example, how did you calculate $\\hat{E}$ and $\\hat{p}$ from the predicted value in Equation 10? What is the norm of the $\\lambda_c|\\cdot|$ in Equation 10? Is it $L_1$ norm or $L_2$ norm?\n\n4. The author uses a deterministic encoder and claims it is better than the variational encoder. According to the author, the latter only has limited benefits while increasing training variance and complexity. However, there is no experiment result to support the claim. More ablation study is needed to justify the claim. Still, it sounds strange to the reviewer that the model is called a variational latent diffusion model but only uses a deterministic encoder.\n\n5. The author proposed three variants of the model. Conditional, unconditional, and the last one, conditional encoder and unconditional decoder. The way of presenting the results confuses the reviewer. Firstly, which one do you promote to use in the conclusion? It seems in Table 1. VLD is better in the latter three metrics, while the UC-VLD is slightly better in the first three metrics. However, Figure 2 shows the framework of VLD, while all the Figures in the result part and the appendix part only show the results for the UC-VLD. If the authors want to promote VLD, consider adding relevant plot results for VLD. If the UC-VLD is better, consider changing Figure 2 to reflect this point. \n\n6. A physics-informed consistency loss is proposed with a hyperparameter $\\lambda_{c}$; how did you adjust the value of this hyperparameter? Moreover, what would be the effect of including and not including this loss?\n\n7. The author claims the model is aimed at high-dimensional inverse problems. However, the training cost for $55$ dimensional variables is expensive, $24$ hours. Moreover, the designed latent space has a higher dimension than the input instead of the commonly used lower latent space dimension for compression. This design could introduce additional costs. And the reviewer is concerned with the scalability of the current model to higher dimensional problems. \n\n8. What is the limitation of the current work?", "limitations": "The author did not mention the limitations explicitly. I would be curious what would be the limitation of the proposed model.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Originality: The model is a combination of two existing models, the latent diffusion model and the variational diffusion model. Moreover, the author defined an appropriate loss function to train the model.\n\nQuality: The result demonstrated that the proposed model outperforms other baseline models.\n\nSignificance: The experiment provides unique data in the high energy physics field.", "weaknesses": "The paper has several typos and unclear points. Moreover, some claims are not well supported by the experiment results. The reviewer is confused to some presented points. Details can be found in the questions part.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful questions and comments.\n\n1. We agree with the reviewer that Figure 1 could benefit from additional context. The intent of this figure was to offer a visual illustration of parton configurations and detector observations, as subsequent sections describe this data solely in terms of kinematic vectors. We have augmented the captions to clarify this intent. \n\n2. We thank the reviewer for noticing these  detailed errors. We have performed a pass through the paper and fixed these errors along with other small wording changes. We have also added the missing Figure number for Figure 1 which did not render in our original submission.\n\n3. As we train our model on the task of reconstruction, the inputs and outputs of our network consist of the parton momentum vectors, defined by the components $(M, \\log E, p_x, p_y, p_z)$. These vectors represent specific physical quantities, such as mass and energy, and the network provides approximate values for these kinematic parameters. In response to this question and comments from other reviewers who requested clarification on the letter symbols and their associated physical quantities, we have added further details to this description. In relation to Equation 10, as all values are scalars, the bars represent a simple absolute value rather than a vector norm. We have clarified this with an explicit formulation in terms of mean absolute error.\n\n4. We use a deterministic encoder only for the detector (conditioning) inputs, not the parton encoder. We have clarified this distinction in the text. We note in the text that we have the option to make the detector encoder variational, but we did not expect any benefit from this approach. Our decision aligns with previous work, including methodologies like CVAE and LDM, which similarly utilize deterministic encoders for their conditioning. This choice was also informed by preliminary experiments, which were insufficiently conclusive to include in our study. We have therefore also revised this section to more precisely articulate that we opted not to explore a variational encoder for the detector observations, owing to both a lack of compelling motivation and following existing precedent.\n\n5. We thank the reviewer for identifying the discrepancy and the absence of detail in Figure 2, as well as for their insightful comment. We have revised Figure 2 to include annotations that highlight the differences between the three variations, thereby providing additional visual intuition to their definitions. We have included this updated figure in this response. We have also clarified that we view the UC-VLD as the best model for this dataset  due to its performance on the bin-independent metrics and present its results in the graphical comparisons.\n\n6. The consistency loss scale functions as a hyper-parameter, comparable to the weight regularization scale, and requires empirical tuning. Our computational limitations precluded a comprehensive hyper-parameter sweep for this parameter, so we selected a conservative value of 0.1, guided by the loss magnitudes observed during training. We observed that omitting this loss leads to a broader mass reconstruction, failing to capture the sharp peak in the mass features, and this trend was consistent across all models. To ensure a fair comparison, we applied the same loss and loss scale for all models.\n\n7. We are indeed looking into higher dimensional problems and problems where the data could have a variable number of dimensions. Diffusion models are infamously very computationally expensive to train, even within existing application domains. We note that our problem has 55 dense features, where every feature is approximately independent and meaningful; this is in contrast to images, which have much higher dimensionality, but each individual feature is less informative. Due to this difference, we opt to make our latent space with higher dimensionality than our data to allow the VAE to learn a fine-tuned latent space for the diffusion objective.\n\n8. We address this question in the global rebuttal."}, {"review_id": "bVZCMvstVc", "reviewer": "Reviewer_Btsg", "summary": "The paper benchmarks several network architectures on a real-life problem in particle physics, i.e. the problem of inverting the effect of limited detector resolution and guess the features of a given collision from what is actually observed in the detector (unfolding). Considering several metrics to assess the accuracy of a given unfolding, the authors show that a novel architecture for diffusion models provides the best performance. ", "questions": "none", "limitations": "none", "rating": 9, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Very solid analysis, with clear explanation of the various steps.\nShows potential progress in applications, thanks to novel architecture", "weaknesses": "none", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their kind comments and support of our paper."}, {"review_id": "dQG7oeKFdZ", "reviewer": "Reviewer_sYuc", "summary": "This paper proposes a unified framework, which combines the latent diffusion model and variational diffusion model, and applies the method to the inverse problem in the field of High Energy Physics. The loss terms for VAE and variational diffusion model are combined to achieve end-to-end training. The proposed VLD and corresponding variants achieve state-of-the-art performance as shown in experiments.", "questions": "* How to define that the dataset at L273 includes enough variations for the high energy physics? For example, is there any existing similar settings for reference?", "limitations": "Technically, this paper provides a unified way, which seems like a straight-forward combination of the loss terms, to train a latent variational diffusion model in an end-to-end manner. However, experiment results are insufficient to support the importance of the extra loss term, or the benefits brought by the unified training process. ", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* A unified framework to support end-to-end training for latent variational diffusion model.\n* The proposed VLD and variants achieve state-of-art performance existing generative models.\n", "weaknesses": "Overall, the main contribution of this paper is the unified framework to combine latent diffusion model and variational diffusion model, where the core lies in the extra VAE loss (3rd term in Equation 8). \n1. Technical novelty seems limited, since the only contribution could be summarized as an extra loss term in diffusion model.\n2. The experiments to emphasis the importance of the extra loss term are limited. Among all the baselines, the LDM at L243-245 is the most similar method to the proposed one, where the only difference is the extra loss term that leads to an end-to-end training process. The small gaps between LDM and C-VLD in Table 1 also suggest the similarities. The main question is, is the comparison between LDM and VLD fair enough? The details about training LDM are missing, such as whether they adopt the same feedforward block as VLD, especially how to pretrain the VAE and what is the performance of the pretained-VAE. Existing experiments seem insufficient to support that the benefits of VLD are from the unified training process.\n3. Since there are results for C-VLD and UC-VLD, which mainly differ in the conditional signal, is it possible to compare LDM with similar settings?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful questions.\n\n1. We respectfully disagree with this characterization. Many important contributions in machine learning, such as weight regularization and the original VAE, may be characterized in this reductive manner if one overlooks the reasoning behind the  respective loss modification. The primary contribution of this paper is a formal derivation of a unified variational architecture, together with its application to an important inverse problem in physics and an assessment of each component of this architecture. End-to-end training of this architecture in a theoretically justified fashion necessitates the addition of a new loss term, derived from the more foundational variational objective. Our experiments demonstrate that the main performance advantage originates from this unified end-to-end training of all the components, allowing the latent space to be fine-tuned to the diffusion task, in contrast to VAEs trained without diffusion or methods lacking a latent space.\n\n2. We respectfully disagree with the reviewer's interpretation of the results. The VLD and LDM share identical network architectures for the denoising network, detector encoder, and VAE. The primary distinction between these two methods lies in pre-training the VAE (LDM) as opposed to training all components end-to-end (VLD). We have highlighted this distinction in the revised text to avoid confusion. Our findings, presented in Table 1, reveal that the regular VLD significantly outperforms the LDM, with the sole discrepancy being the unified training. We did observe that including a conditional decoder unexpectedly worsened performance, a phenomenon we discuss in the text in the paragraph starting on line 321. We hypothesize that this reduction may stem from the conditional decoder's sensitivity to minor errors in the diffusion process during inference.\n\n3. We thank the reviewer for their interesting suggestion. While it may be possible to formulate a pre-trained LDM which uses a conditional VAE, doing so is non-trivial. The design would need to decide between using individual conditioning encoders for each phase or sharing the encoder for both phases. This is further complicated by the fact our domain lacks a pre-trained conditioning akin to CLIP for text. We encourage future work to examine this problem, but one of the benefits of our unified approach is that this change is simpler to implement while maintaining a common framework.\n"}, {"review_id": "7XbkFmLE7C", "reviewer": "Reviewer_GAKk", "summary": "The paper introduces a diffusion model based approach to tackle inverse problems. The method is then applied to High Energy Physics in order to reconstruct kinematic quantities.", "questions": "1. Page 4, line 143, typo, should be “in an abstract”\n2. Page 6, line 204, please define E to be energy, m mass, p momentum, and mention that you are using the HEP convention of setting c to be 1.\n3. Please discuss the reason why UC-VLD outperforms the VLD on some metrics and vice-versa.\n4. Figure 4, b quark, I don’t really see the bimodal nature of the distribution\n5. Figure 4, neutrino, how is the neutrino “truth” line showing the bimodal nature?\n6. What would happen with the LDM’s performance if a different prior weight is given?\n7. What would happen in the benchmarks if the CINN uses the MMD loss instead?\n8. To further show the improvement of the described unified training loss, it would be beneficial to show what would happen if the various components would be trained in parallel.", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well written with clear structure, definitions, and figures. The paper includes testing of the proposed methodology, VLD, on an important HEP inverse problem, which is unfolding semi-leptonic tt events. It is benchmarked on this task against other state of the art algorithms such as CINN, LDM, and VDM, as well  as variations of the authors own proposed algorithm. Their VLD algorithm does appear to significantly outperform the other considered algorithms.\n\nAnother strength of this paper is the fact that the authors are deliberate in construction of their network, explaining the advantages and necessity of each of the components.", "weaknesses": "Limited testing of the proposed unifying architecture is a weakness of the paper. Indeed, the authors do acknowledge that further testing on different event topologies would be beneficial, yet including at least one more benchmark would increase the confidence in the algorithm’s performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their king comments and detailed questions. We address the limitations comment in the global rebuttal. We answer the questions here.\n\n1. & 2. We thank the reviewer for noticing these points. We have done a pass through the paper and fixed these presentation issues along with other small wording changes.\n\n&#8291;\n\n3. We thank the reviewer for their insightful question. Unfortunately, we do not yet have an intuitive explanation for this behavior. The complexity is further compounded by the absence of a standard (and computationally feasible) method for measuring the distance between high-dimensional distributions. Each of the metrics we provide captures different aspects of distribution distance. We think the bin-independent metrics are more sensitive to outliers and long tails whereas the binned metrics focus more on the distance between the high-mass regions of the distribution. An examination of the detailed appendix results reveals that the UC-VLD performs more favorably with long-tailed features such as energy, but less so with the peaked mass term.\n\n4. & 5. We regret the confusion surrounding the bimodal posteriors. Only the neutrino $\\eta$ posterior is expected to be bimodal. This arises because 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 $\\eta$ term which often leads to two configurations satisfying both energy and momentum conservation. In Figure 4, we notice that the empirical estimate fails to capture this expectation, whereas the LVD captures this behavior without explicit information about this expectation. We also note that the truth sample will still be a single value sampled from this theoretically bimodal distribution. We have clarified this phenomenon in the revised text to better describe this expectation and our results.\n\n&#8291;\n\n6. This is an interesting question. We use a low prior weight in accordance with the LDM paper, as we replicate their VAE pre-training. We suspect that, since our use case makes the latent dimension larger than the data dimension, the prior loss is dominating the pre-training step. This is because reconstruction loss may trivially be near-perfect without a bottleneck, so the main constraint for the VAE is the Gaussian prior over the latent space.\n\n7.  The MMD (Maximum Mean Discrepancy) loss was used in the CINN paper to more accurately represent the true distributions associated with kinematics, such as the peaked mass term. We could not incorporate this MMD loss into our variational framework in a natural way as we employ a Gaussian VAE which necessitates a mean squared error loss term for the reconstruction. Instead, we chose to incorporate an additional self-consistent physics-informed regularization loss to enhance mass reconstruction, citing analogous strategies previously implemented for other generative models in physics. These two approaches share similar goals, although a detailed comparison of these methods could be fruitful for future work. We use exactly the same loss for all methods with the goal of evaluating only the generative model architecture's effect on performance, \n\n8. We agree with this reviewer and in fact we reported the results of such an experiment in Table 1. In this table, we compare the component-wise training (LDM) to the unified end-to-end training (VLD), and show that the latter outperforms the former by roughly one order of magnitude. Since we use identical network architectures for all methods, the primary difference between these two methods is the unified training. We have further clarified this point in the revised version, "}]

Author rebuttal:
We thank all of the reviewers for their detailed comments and questions. We begin with a discussion of the limitations as requested by two of the reviewers and then answer the remaining questions in individual responses. We also include a document with updated figures.

Limitations
--------------
We note in the text that the experiments performed in this study are limited to parton unfolding of a specific event topology and on simulated data. The method is general and may be applied to arbitrary topologies, with our choice guided by the limits of current baseline methods. We think this exploration is robust as many practically explored topologies are simpler than our semi-leptonic $t\bar{t}$ test-case. We note in the text that the next step would be to perform particle-level unfolding which is not specific to a particular topology, and we believe this technique may be applied to this more general problem as well with little modification. Another common limitation of training on simulated data is the imperfection of the simulation and the possibility of skewing the results compared to real detector data. We note in the text that we must adjust the model’s simulation bias by employing real data, and we cite potential techniques, such the as iterative approach of ICINN, to accomplish this adjustment. We have added further discussion to the text to illustrate these limitations and future solutions.

Meta-review: The reviewers were mixed on this paper and sought clarification on a number of points.  The authors engaged in the rebuttal process to provide answers to many of the questions raised.  While not all reviewers recommended acceptance, after reading the reviews, discussion and paper I recommend acceptance.  I believe the paper considers a challenging and interesting problem.  The proposed methodology appears to be well tailored to the task and demonstrates significant performance improvements over existing techniques.  The authors should revise the paper to consider the issues raised in the review and discussion process, including additional clarifications.
Decision: Accept (poster)