Variational Pseudo Marginal Methods for Jet Reconstruction in Particle Physics

TL;DR

This paper tackles jet reconstruction by treating jet histories as latent binary trees and addressing the intractable combinatorial space of possible topologies. It introduces a combinatorial sequential Monte Carlo framework (CSMC) tailored to jets, and builds two variational approaches (VCSMC and VNCSMC) to learn both local tree structures and a global decay parameter $\lambda$, further unifying inference with a variational pseudo-marginal framework. The methodology yields unbiased marginal-likelihood estimators and scalable variational objectives, achieving superior speed and accuracy compared with baseline clustering methods on data generated by the Ginkgo jet generator. The work demonstrates a principled, fully Bayesian treatment of hierarchical jet reconstruction with potential for broad impact in collider data analyses and probabilistic modeling of combinatorial structures in physics. Overall, the approach offers a scalable, uncertainty-aware pathway to learn jet topologies and parameters, enabling more precise physics inferences and model calibration.

Abstract

Reconstructing jets, which provide vital insights into the properties and histories of subatomic particles produced in high-energy collisions, is a main problem in data analyses in collider physics. This intricate task deals with estimating the latent structure of a jet (binary tree) and involves parameters such as particle energy, momentum, and types. While Bayesian methods offer a natural approach for handling uncertainty and leveraging prior knowledge, they face significant challenges due to the super-exponential growth of potential jet topologies as the number of observed particles increases. To address this, we introduce a Combinatorial Sequential Monte Carlo approach for inferring jet latent structures. As a second contribution, we leverage the resulting estimator to develop a variational inference algorithm for parameter learning. Building on this, we introduce a variational family using a pseudo-marginal framework for a fully Bayesian treatment of all variables, unifying the generative model with the inference process. We illustrate our method's effectiveness through experiments using data generated with a collider physics generative model, highlighting superior speed and accuracy across a range of tasks.

Figures (9)

• Figure 1: Jets as binary trees. Left: Schematic representation of the production of a jet at CERN's LHC. Incoming protons collide, producing two new particles (light blue). Each new particle undergoes a sequence of binary splittings until stable particles (solid blue) are produced and measured by a detector. Right: In jet reconstruction, only the leaf nodes are observed and the tree topology is inferred. Each latent tree topology represents a different possible splitting history.
• Figure 2: Illustration of the Ginkgo generative and reconstruction processes. (a) Ginkgo starts with a parent node characterized by a 4-vector $z_p = (E,\vec{p})$ and invariant mass squared $t_P = E^2 - |\vec{p}|^2$. If $t_P$ is greater than the cut off value $t_{cut}$, then the parent node splits (we have a particle decay). The left and right nodes invariant mass squared ($t_L$, $t_R$) are sampled from a truncated exponential distribution defined in Eq. \ref{['exponential']}. (b) The splitting likelihood reconstruction process of a node, defined in Eq. \ref{['splitting_likelihood']} begins with two child nodes $L$ and $R$ along with their respective 4-vectors $z_L$ and $z_R$. The 4-vector for the parent node $P$ is calculated as $z_P = z_L + z_R$ and then $t_P = t(z_P)$. Next, we obtain $t_L = t(z_L)$, $t_R = t(z_R)$ and define $t_P^L = t_P$, and $t_P^R = (\sqrt{t_P} - \sqrt{t_L}) ^ 2$. Finally, the left splitting term likelihood $\ell(t_L,\lambda, t_{cut}, t_P^L, t_P)$ and the right one $\ell(t_R,\lambda, t_{cut}, t_P^R, t_P)$ are evaluated.
• Figure 3: Summary of the Csmc framework: A total of $K$ partial states $\{s_{r}^k\}_{k=1}^{K}$ are retained as collections of tree structures encompassing the data set. A partial state is defined as a collection of trees, which start out as singleton particles $A$, $B$, $C$ and $D$. Each iteration within Algorithm \ref{['alg:csmc']} comprises three key stages: (1) resampling partial states based on their importance weights $\{w_{r}^k\}_{k=1}^{K}$, (2) proposing an expansion of each partial state to form a new one by linking two trees within the forest, and (3) determining the new weights for these new partial states. The illustration above depicts three samples across a jet consisting of observed four particles, denoted as $A,B,C,$ and $D.$
• Figure 4: The likelihood $\mathcal{F}_A$ for a sub-tree defined on leaf nodes $D$, $E$, $F$ and $G$ is defined as the recursive product of splitting likelihoods $\mathcal{F}_B$ and $\mathcal{F}_C$. The intermediate target $\pi(s_3)$ for the partial state $s_3$ also includes the probability of singletons $H$ and $I$ denoted $\mathcal{F}_H$ and $\mathcal{F}_I$.
• Figure 5: Left: Scatterplot comparing log-conditional likelihood of Vncsmc with $K,M = (256,1)$ vs Greedy Search and Center: Scatterplot comparing log-conditional likelihood of Vncsmc with $K,M = (256,1)$ vs Beam Search. Across 100 simulated jets, Vncsmc returns higher likelihood on all 100 cases against Greedy Search and 99 cases against Beam Search. Right: Log-conditional likelihood values for Vcsmc (blue) and Vncsmc (red) with $K = \{256\}$ (and $M = 1$) samples averaged across 5 random seeds. Vncsmc achieves convergence in fewer epochs than Vcsmc and yields higher values, all while maintaining lower stochastic gradient noise. Additional experiments demonstrating the effect of $K$ appear in Fig. \ref{['fig:elbo_conv']} of the Appendix.

Theorems & Definitions (2)

• Definition 1: Partial State
• Definition 2: Natural Forest Extension