SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
D _ {K L, N} = \sum_ {i = 1} ^ {N} P _ {N} (i) \log \left(\frac {P _ {N} (i)}{Q _ {N} (i)}\right)
$$
D Particle Distance Tables
Tables 3 to 6 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 7 to 10 present the distance metrics for each component and model. The mass component seems to vary the most for many of the distance functions, indicating that many models struggle reconstructing the peaked mass distributions. However, the overall results remain consistent with Table ??. We again see the clear benefit of both latent diffusion and end-to-end training.
F Global Distribution Plots
Figures 6 through 16 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 is slightly smoothed but peaks in the correct location. Baseline models struggle with capturing the peaks and shapes of the distributions.
G Posterior Distribution Plots
Figures 17 and 18 present a collection of 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 in Equation 8, adhering to the generative model displayed in Figure 2. Here, data is generated from a latent representation,  $p(x|z_x)$ ; the latent
data from the diffusion end-point,  $p(z_{x}|z_{0})$ ; a diffusion process to generate the sampled configuration  $p(z_{0}|z_{1})$ ; and a simple standard normal distribution prior at the base of the hierarchy  $p(z_{1})$ . All these functions are further conditioned on the (embedded) detector observations  $z_{y}$ . We start with the full conditional hierarchical VAE loss function over  $L$  variational layers as per [64].
$$
\mathcal {L} _ {\mathrm {E L B O}} = \mathbb {E} _ {q \left(z _ {x} \mid x, z _ {y}\right)} \left[ - \log p \left(x \mid z _ {x}, z _ {y}\right) \right] + \sum_ {l = 1} ^ {L} D _ {K L} \left(q \left(z _ {l} \mid x, z _ {y}, z _ {i \neq l}\right) \| p \left(z _ {l} \mid z _ {<   l}\right)\right) \tag {11}
$$
Next, we substitute the layers we defined for the sum to expand the expression for our generative model. For this step, we employ a three-stage hierarchy with the following substitutions:  $z_{1} \gets z_{1}$ ,  $z_{2} \gets z_{0}$ , and  $z_{3} \gets z_{x}$ . We derive each of these components in the following sections.
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {E L B O}} = \mathbb {E} _ {q (z _ {x} | x, z _ {y})} \left[ - \log p (x | z _ {x}, z _ {y}) \right] (12) \\ + D _ {K L} \left(q \left(z _ {1} \mid x, z _ {y}, z _ {x}\right) \| p \left(z _ {1}\right)\right) (13) \\ + D _ {K L} \left(q \left(z _ {0} | x, z _ {y}, z _ {1}, z _ {x}\right) \| p \left(z _ {0} | z _ {1}\right)\right) (14) \\ + D _ {K L} \left(q \left(z _ {x} \mid x, z _ {y}, z _ {1}, z _ {0}\right) \| p \left(z _ {x} \mid z _ {1}, z _ {0}\right)\right) (15) \\ \end{array}
$$
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(z_{1}) \sim \mathcal{N}(\mathbf{0},\mathbb{I})$ . We learn the noise schedule via log  $SNR(t) = -\gamma_{\phi}(t)$ , as defined in VDM [41]. As such, we 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:  $q(z_{1}|x,z_{y},z_{x}) = \mathcal{N}(\alpha_{1}z_{x},\sigma_{1}\mathbb{I})$ , where  $\sigma_t = \sqrt{\sigma(amma_\phi(1))}$ , and  $\alpha_{1} = \sqrt{\sigma(-\gamma_{\phi}(1))}$  as obtained from Section 2.4. We estimate this KL divergence through Monte-Carlo sampling:
$$
D _ {K L} \left(q \left(z _ {1} \mid x, z _ {y}\right) \| p \left(z _ {1}\right)\right) = \mathbb {E} _ {z _ {x} \sim q \left(z _ {x} \mid x, z _ {y}\right)} \left[ \left(\alpha_ {1} z _ {x}\right) ^ {2} + \sigma_ {1} ^ {2} - \log \left(\sigma_ {1} ^ {2}\right) - 1 \right] \tag {16}
$$
H.2 VAE Loss
Equation 15 delineates our contribution to this unified variational model. We derive the base distributions for the right-hand distribution in Equation 9:  $p(z_{x}|z_{0}) \sim \mathcal{N}\left(\frac{1}{\alpha_{0}} z_{0}, \frac{\sigma_{0}}{\alpha_{0}}\mathbb{I}\right)$ . We further describe that the posterior distribution in this KL term is merely the regular VAE posterior  $q(z_{x}|x, z_{y}) \sim \mathcal{N}(\mu_{\theta}(x, z_{y}), \sigma_{\theta}(x, z_{y}))$ . We define the posterior over  $z_{0}$  given  $z_{x}$  by adhering to the definition of the VP diffusion process  $q(z_{0}|z_{x}) \sim \mathcal{N}(z_{x}, \sigma_{0}\mathbb{I})$ . As all these distributions are normal, we can provide an explicit form for this loss:
$$
\begin{array}{l} D _ {K L} \left(q \left(z _ {x} \mid x, z _ {y}\right) \| p \left(z _ {x} \mid z _ {1}\right)\right) = \\ \mathbb {E} _ {z _ {x} \sim q (z _ {x} | x, z _ {y}), z _ {0} \sim q (z _ {0} | z _ {x})} \left[ D _ {K L} (\mathcal {N} (\mu_ {\theta} (x, z _ {y}), \sigma_ {\theta} (x, z _ {y}) \| \mathcal {N} \left(\frac {1}{\alpha_ {0}} z _ {0}, \frac {\sigma_ {0}}{\alpha_ {0}} \mathbb {I}\right)) \right] \\ \end{array}
$$
where the KL term is the regular normal distribution KL.
$$
D _ {K L} (\mathcal {N} (\mu_ {0}, \sigma_ {0}) \parallel \mathcal {N} (\mu_ {1}, \sigma_ {1})) = \frac {1}{2} \left[ \frac {\sigma_ {0} ^ {2}}{\sigma_ {1} ^ {2}} + \frac {(\mu_ {0} - \mu_ {1}) ^ {2}}{\sigma_ {1} ^ {2}} + \log \sigma_ {1} ^ {2} - \log \sigma_ {0} ^ {2} - 1 \right]
$$
H.3 Diffusion Loss
Equation 14 defines the final diffusion loss term for the denoising network. We follow the derivation from Kingma et. al. [41] for a 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(z_{t}|x,y) \sim \mathcal{N}(\alpha_{t}x,\sigma_{t}\mathbb{I})$ . Following [41], we derive a Monte-Carlo estimate
of the integral loss for a noise prediction network:
$$
\begin{array}{l} D _ {K L} \left(q \left(z _ {0} | x, z _ {y}, z _ {1}, z _ {x}\right) \| p \left(z _ {0} | z _ {1}\right)\right) = - \frac {1}{2} \mathbb {E} _ {\epsilon \sim \mathcal {N} (\mathbf {0}, \mathbb {I})} \left[ \int_ {0} ^ {1} S N R ^ {\prime} (t) \| z _ {x} - \hat {z _ {x}} \left(z _ {t}, z _ {y}, t\right) \| _ {2} ^ {2} d t \right] \\ = - \frac {1}{2} \mathbb {E} _ {t \sim \mathcal {U} (0, 1), \epsilon \sim \mathcal {N} (\mathbf {0}, \mathbb {I})} \left[ S N R ^ {\prime} (t) \| z _ {x} - \hat {z _ {x}} (z _ {t}, z _ {y}, t) \| _ {2} ^ {2} \right] \\ = \frac {1}{2} \mathbb {E} _ {t \sim \mathcal {U} (0, 1), \epsilon \sim \mathcal {N} (\mathbf {0}, \mathbb {I})} \left[ \gamma_ {\phi} ^ {\prime} (t) \| \epsilon - \hat {\epsilon_ {\theta}} (z _ {t}, z _ {y}, t) \| _ {2} ^ {2} \right] \\ \end{array}
$$
H.4 Reconstruction Loss
Equation 12 is the final and simplest aspect of this model. Since we are regressing the parton components, we use a Normal decoder on the VAE and use a simple Mean Squared Error loss as the primary reconstruction loss.
$$
\mathcal {L} _ {M S E} = \left\| \operatorname {D E C O D E R} \left(z _ {x}, z _ {y}\right) - x \right\| _ {2} ^ {2}
$$
As we explain in the text, this MSE loss works well for most components but fails to 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.
$$
\mathcal {L} _ {C} = \lambda_ {C} \left| \hat {M} ^ {2} - \left(\hat {E} ^ {2} - \| \hat {p} \| ^ {2}\right) \right|
$$
The total reconstruction loss is simply the sum of these two components.
$$
\mathbb {E} _ {q \left(z _ {x} \mid x, z _ {y}\right)} \left[ - \log p \left(x \mid z _ {x}, z _ {y}\right) \right] = \mathbb {E} _ {q \left(z _ {x} \mid x, z _ {y}\right)} \left[ \mathcal {L} _ {M S E} + \mathcal {L} _ {C} \right]
$$
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νe/νμ	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 b	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 W	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 t	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
tbar System	16.15	9.52	85.58	70.04	384.93	68.85	761.74
Table 3:Particle Distance: Wasserstein Distances
Figure 6: Global Distribution: Leptonic  $b$  Quark
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νe/νμ	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 q1	0.90	0.77	2.71	2.75	16.30	4.82	14.73
Hadronic q2	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 t	0.86	0.44	2.03	2.43	22.16	3.32	24.81
tt System	0.80	0.78	4.36	3.72	18.63	3.47	34.98
Table 4: Particle 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νe/νμ	0.97	1.06	0.73	0.88	1.76	0.77	1.71
Leptonic e/μ	0.48	0.33	0.52	0.35	1.47	0.43	1.02
Hadronic b	0.05	0.05	0.10	1.06	1.95	0.90	1.30
Hadronic q1	0.40	0.40	0.60	0.63	1.63	0.81	1.50
Hadronic q2	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 t	0.10	0.09	0.32	0.28	1.60	0.37	1.48
tt System	0.06	0.09	0.26	0.24	0.82	0.21	1.68
Table 5: Particle Distance: Kolmogorov-Smirnov 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νe/νμ	0.14	3.45	0.87	1.43	0.76	0.88	2.48
Leptonic e/μ	0.14	1.35	0.44	0.93	3.66	1.04	3.43
Hadronic b	0.00	0.01	0.01	3.57	5.99	0.29	1.09
Hadronic q1	0.34	2.07	0.46	1.19	3.80	1.72	6.15
Hadronic q2	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 t	0.02	0.03	1.22	1.47	2.50	0.98	2.24
Hadronic t	0.03	0.04	0.99	1.14	2.39	1.05	2.45
tt System	0.01	0.01	0.06	0.04	0.48	0.04	3.59
Table 6: Particle Distance: KL Divergence with 128 bins.
VLD	UC-VLD	C-VLD	LDM	VDM	CVAE	CINN
mass	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
eta	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	CVAE	CINN
mass	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
eta	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
pz	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
eta	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
px	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-Smirnov 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: KL Divergence with 128 bins.
Figure 7: Global Distribution: Leptonic Neutrino  $\nu_{e} / \nu_{\mu}$
Figure 8: Global Distribution: Leptonic Lepton  $e / \mu$
Figure 9: Global Distribution: Hadronic  $b$  Quark
Figure 10: Global Distribution: Hadronic Light Quark
Figure 11: Global Distribution: Hadronic Light Quark
Figure 12: Global Distribution: Leptonic  $W$  Boson
Figure 13: Global Distribution: Hadronic  $W$  Boson
Figure 14: Global Distribution: Leptonic Top Quark
Figure 15: Global Distribution: Hadronic Top Quark
Figure 16: Global Distribution: Complete  $t\bar{t}$  System
(a) Example Event 1
Figure 17: Posterior distributions for example events. Included is an empirical posterior distribution calculated from the training dataset.
(b) Example Event 2
(a) Example Event 3
Figure 18: Posterior distributions for example events. Included is an empirical posterior distribution calculated from the training dataset.
(b) Example Event 4