Variational Pseudo Marginal Methods for Jet Reconstruction in Particle Physics
Authors: Hanming Yang
Antonio Khalil Moretti
Sebastian Macaluso
Philippe Chlenski
Christian A. Naesseth
Itsik Pe'er
Venue: Accepted by TMLR ()

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 of 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.

Published in Transactions on Machine Learning Research (12/2024)
Variational Pseudo Marginal Methods for Jet Reconstruction in Particle Physics
Hanming Yang"
hy2781 @columbia. edu calloniomoreltiisspitin edu seb. macaluso @gmaail com pacBes. cobumbia edu c. naesseth @ruvo. nll itsik&os. www.:eessssccoboooooooo columbia edu
Antonio Khald Moretti 1.2,*
Sebastion Macaluso
Philippee Chlenski
Christian Naesseth 4
Itsik Pe'er"
DDeportment of Computer Science, Columbia University
DDpprtment of Computer Science, Spehman College
Telefonica Research University of Amsterdam
*Equal contribution
Reviewed on OpenReview: https: openreview. net/ forum? id PCopRF2uF
Abstract
Reconstructing jets, which provide vital insights into the properties and histories of sub- atomic particles produced in high-energy collisions, is main problem in data analyses of collider physics. This intricate task deals with estimating the latent structure of jet (bi- nary tree) and involves parameters such as particle energy, momentum, and types. While Bayesian methods offer natural approach for handling uncertainly and leveraging prior knowledge, they face significant challenges due to the super-exponential growth of potential jet topologics 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 develop a variational inference algorithu for parameter learning. Building on this, we introduce variational family using ggeudo-marging framework for fully Bayesian treatment of all variables, unifying the gen- erative 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.
Introduction
Reconstructing jets in particle physics deals with estimating a high-quality hierarchical clustering, A com- prehensive approach to this process also involves inference on model parameters, which could provide insights into our understanding of quantum chromodynamics (QCD), i.c. the theory of the strong interaction between quarks mediated by gluons. Hierarchical clustering forms natural data representation of data generated by Markov tree, and has been applied in wide variety of settings such entity resolution for knowledge- bases (Green et al. 2012 Vashishth et al. 2018), personalization (Zhang et al. 2014), and jet physics CCcciari ct al. 2008 Catani et al.] 1993 Dokshitzer et al. 1997 Ellis & Soper 1993). Typically, work has focused
on approximate methods for relatively large datasets Bateni el al. 2017 Monath el. al. 2019 Naumov Monath et al, 2021 However, there are relevant use cases for hicrarchical clustering that require exact or high-quality approximations on small medium-sized datasets (Greenberg et al., 2020 2021). This paper deals with one of these cases: reconstructing the latent hierarchy of jets in particle physics. Within this context.
le al. 2020: Dubey et al., 20144] Hu et al.] 2015 Monath et al.) 2020 Dubey et al. 2020

---

Published in Transactions on Machine Learning Research (12/2024)
Jet generation
Jet reconstruction
Proton
Leaves
Proton
Observed
Particles
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 roconstruction, only the leaf nodes observed and the tree topology is inferred. Each latent tree topology represents different possible splitting history.
Bayesian methods provide a natural approach for handling uncertainty, but the super-exponential scaling of the number of hierarchies with the size of the datasets presents significant difficulties, i.e. the number of topologies grows as (2.V - 3)!!, with N being the number of leaves. This super-exponenti growth in the space of configurations makes brute force and exact methods intractable.
1.1 Jet physics
During high-energy particlc collisions, such as those observed at the Large Hadron Collider (LHC) at CERN, collimated sprays of particles called jels are produced. The jet constituents are the observed final-state particles that hit the detector and are originated by a showering process (described by QCD) where an initial (unstable) particle goes through successive binary splittings. Intermediate (latent) particles can be identified as internal nodes of a hicrarchical clustering and the final-state (observed) particles correspond the leaves in binary tree. Fig. | provides schematic representation of this process. This results in several possible latent topologies corresponding to set of leaves. This representation, first suggested in Louppe ct al. 2019), connects jets physics with natural language processing (NLP) and biology.
1.2 Collider data analysis
A main problem in data analyses of collider physics deals with estimating the latent showering process (hierarchical clustering) of a jet, which is needed for subsequent tasks that aim to identify the creation of different types of sub-atomic particles. The final goal is to perform precision measurements to test the predictions of the Standard Model of Particle Physics and explore potential models of new physics particles, thereby advancing our comprehension of the universe's fundamental constituents. The improved performance of deep lcarning jet classifiers (Butter et al. 2019) (for the initial state particles) over traditional clustering- based physics observables gives evidence of the limitedness of current. clustering algorithms since the right clustering should be optimal for classification tasks. Thus, high-quality approximations in this context would highly beneficial for data analyses in experimental particle physics.
1.3 Jet simulators
Currently, there are high-fidelity simulations for jets, such as Pythia (Soostrand al. 2006). Herwig (Bellm) et al. 2016), and Sherpa (Gleisberg et al. 2009). These simulators are grounded in QCD but make several approximations that introduce parametric modcling choices that are not predicted from the underlying the- ory, so they are commonly luned to match the data. Many lasks in jet physics can be framed in probabilistic terms Cranmer et al. 2021). In particular, we consider the challenges of calculating the maximum likelihood hicrarchy given a set of leaves (jet constituents), the posterior distribution of hicrarchics, as well as estimat- ing the marginal likelihood. Also, these quantities are relevant to tune the parameters of the simulators. While these formulations are helpful conceptually, they are not practical in current high-fidelity simulations
2

---

Published in Transactions on Machine Learning Research (12/2024)
for jets, given that the likelihood is typically intractable (they are implicit models). Thus, we consider Ginkgo (Cranmer, Kyle et al.) 2021 Cranmer et. al. 2019b;: semi-realistic generative model for jets with a tractable joint likelihood and captures essential ingredients of parton shower generators in full physics simulations. In particular, Ginkgo was designed to enable implementations of probabilistic programming. differentiable programming, dynamic programming and variational inference. Within the analogy bet ween jets and NLP, Ginkgo can be considered as ground-truth parse trees with a known language model.
1.4
Related work
1.4.1 Jet clustering
For each jet produced at the LHC, there is an inference task on the latent hierarchy that typically involves 10 to 100 particles (leaves). Though this is a relatively small number of clements, exhaustive solutions are intractable, and current exact methods, e.g., Greenberg et al. (2020; 2021 have limited scalability. The industry standard uses agglomerative clustering techniques, which are greedy and based on heuris- tics (Cacciari et al.) 2008) Catani al.) 1993 Dokshitzer et al. 1997 Ellis & Soper 1993), typically finding low-quality hicrarchical clusterings. Regarding likelihood-based clustering (applied to Ginkgo datascts), previous work Greenberg el al. (2020) introduced a classical data structure and dynamic programming al- gorithm (the cluster trellis) that exactly finds the marginal likelihood over the space of configurations and the maximum likelihood hierarchy. Also, an A search algorithm combined with a trellis data structure that
finds the exact maximum likelihood hierarchy was introduced in Greenberg et al. (2021). ). Finally, Cranmer et al. (2022) pairs Ginkgo with the cluster trellis (Greenberg et al. 2020), to use the marginal likelihood
directly characterize the discrimination power of the optimal classifier 4J. Stuart & Arnold 1994 Cranmer
& Plehn 2007) as well to compute the exact. maximum likelihood estimate for the simulator's parameters.
While these works provide exact algorithms that extend the reach of brute force methods, they have exponential space and time complexity, becoming intractable for datasets with as few as 15 leaves. For this
reason, Greenberg et al. 2200 2021 also provide approximate solutions al the cost of finding lower-quality hierarchies.
1.4.2 Bayesian inference
A recent body of research has melded variational inference (VI) and sequential search. These connections are realized through the development of a variational family for hidden Markov models, employing Sequential
Monte Carlo (SMC) as the marginal likelihood estimator (Maddison et al.) 2017 Naesseth et al. 2018
Le et al., 2018: Moretti et al., [201 2020 2021). Within the field of Bayesian phylogenctics (the study of evolutionary histories), various methods have been proposed for inference on tree structures. Common approaches include local search algorithms like madomm-walk McMc (Ronquist et al.] 2012)) and sequential
scarch algorithms like Combinatorial Sequential Monte Carlo (CSMC) (Bouchard-Côté et al. 2012 Wang
ct al. 2015 MCMC methods also handle model learning. Dinh et al. (2117) proposes PPHMC which extends Hamiltonian Monte Carlo to phylogenies. Evaluating the likelihood term in MCMC acceptance ratios can be challenging. As workaround, particle MCMC (PMCMC) algorithms use SMC to estimate the marginal likelihood and define MCMC proposals for parameter learning (Wang & Wang 2020).
Pseudo-marginal methods are a class of statistical techniques used to approximate diffieult-to-compute prob- abilities, typically by introducing auxiliary random variables to form an unbiased estimate of the target prob- ability (Andricu & Roberts 2009 Beaumont (2003) introduced a method in genetics to sample genealogies in fully Bayesian framework. Tran et al. (2016) utilizes beeudo-magginn
methods to perform variational
Bayesian inference with an intractable likelihood. Our work is a synthesis of Wang et al. (2015) and Moretti et al. (2021) in that we introduce a variational approximation on topologies using SMC and a VI framework to learn parameters.
1.5 Contributions of this Paper:
We expand upon the CSMC technique introduced by Wang et al. (2015) and the NCSMC method from Moretti al. (2021) to introduce a Combinatorial Sequential Monte Carlo framework for inferring hierarchical clusterings for jets. The resulting estimators are unbiased and consistent. To the best
3

---

Published in Transactions on Machine Learning Research (12/2024)
of our knowledge, this marks the first adaptation of SMc methods to jet reconstruction in particle physics.
2. We leverage the resulting SMC estimators to develop two approxinate posteriors jet hierarchies and correspondingly two VT methods for parameter learning. We illustrate the effectiveness of both methods through experiments using data generated with Ginkgo (Cranmer, Kyle et al. 2021 highlighting superior speed and accuracy across various tasks.
3. In order to circumvent parametric modeling assumptions, we propose a unification of the generative model and the inference process. Building upon the point estimators, we define a distinct variational family over global and local parameters for a fully Bayesian treatment of all variables.
We show how partial states and re-sampled indices generated by SMC can be interpreted as auxil- iary random variables within a pseudo-marginal framework, thus establishing connections bet ween variational greudo-marging methods and VSMC (Naesseth al. 2018 Moretti al. 2021).
2
Background
Section 2.1
provides an overview of the Ginkgo generative model for jet physics (Cranmer, Kyle et al. 2021 summarizes the approximate inference techniques this work builds upon.
Section 2.3.
2.1
Ginkgo generative model
In this subsection we provide an overview of the generative process in Ginkgo as well as jet (binary tree) reconstruction during inference. As mentioned in the introduction, Ginkgo is a semi-realistic model designed to simulate a jet. The branching history of jet is depicted as binary tree structure T (V) where denotes topology, V comprises the set of vertices, and È the set of edges. Each node is characterized by 4D momerty-mommntumm) vector = (E' E RR 5 E R$) where E' denotes energy and p (Px, Py, Px) denotes momentum in the respective dimensions. The squared mass t t(z) = E² - 1/ is calculated using the energy-momentum vector 2. The terminal nodes (or leaf nodes), represented as X = (xI, ... ,I,,,, correspond to the observed energy-momentun vectors measured at the detector. The tree topology T and the energy-momentum vectors associated with internal nodes, denoted as {21.-.. - are latent variables in the model.
2.1.1 Generative process
In Ginkgo the generative process begins with the splitting of a parent (root) node with invariant mass squared tp into two children, as shown schematically in Fig. 2 (left). The process is characterized by a cutoff mass squared fext, and a rate parameter À for the exponential distribution governing the decay. During generation,
as long as the invariant mass squared of node exceeds the cutoff value (tr > feut), that node is promoted to be parent and the algorithm recursively splits it. The squared masses of each new left. (L) and right (R) child nodes (fL, tR) are obtained from sampling from the exponential distribution f(t)A, tp) defined in Eq. [ with parameters specific to each child (L, R). Finally, once tr and tn are sampled, the corresponding energy-
momentum vectors (ZL, #R) for the (L, R) nodes are derived from (p, iy and iR following meergy-mmmentum
conservation rules, i.e. applying a 2-body particle decay (see (Cranmer, Kyle et al.] 2021 2019b) for more details). Next, we specify the exponential distribution
Cranmer et al.
1 X -
f(I)) t/p)) - 1- Ep
(1)
where the first term (1 (1-e- (1---A)-1//1 is normalization factor, i E (L, R), 1/. =r, and tB = (ytr vtz)² To satisfy energy-momentum conservation (VTz + VtR < VTr,, l is sampled before Ln- Thus, in Ginkgo the left child mass squared lr. is sampled first with 1/ = (p and then an auxiliary value 1 = (vlr vTr)² is calculated to sample the right child mass squared fR (see Fig. (left)).
There are different types of jets, depending on the type of initial state particle (root of the binary tree). To simulate QCD-like jets, single À parameter is employed for the entire process. However, for heavy
4

---

Published in Transactions on Machine Learning Research (12/2024)
1p = ((p)) >
beut
Lp 1(21. + zR)
E(1z, X, beuts 14. Lp) f(t)),tp)
P
f(In: X, lout, (pp (p)
[(JA, (Vtp -V1))
L
R
L
R
tu ~ f(1),tp) tR ~ f(-1, (Vtp -tt))-vt)' vLz)²)
tr = t(ZL)
iR = ((zR)
(a) Ginkgo generative process.
(b) Node splitting likelihood reconstruction process.
Figure 2: Illustaatinn of the Ginkgo generalive and reconstruction processes. (a) Ginkgo starts with a parent node characterized by a 4-vector 2p = (E. p) and invariant mass squared tp = E² 1p]². If tp is greater than the cut off value feut, then the parent node splits (we have a particle decay). The left and right nodes invariant mass squared (tz, tR) are sampled from a truncated exponential distribution defined in Eq. D (b) The splitting likelihood reconstruction process of node, defined in Eq. El begins with two child nodes L. and R along with their respective 4-vectors zz and ZR. The 4-vector for the parent node P is calculated as 2p = 2L + zR and then tp = t(zp). Next, we obtain tr t(ZL), tR = t(zg) and define tp = tp, and 1/ = (vlp- vTz)². Finally, the left splitting term likelihood f(LL, X, leaa (h.pp) and the right one Eltn, X, Leat, tB, tp) are ovaluated.
resonance particle decay, such a W boson jet, the initial (root node) splitting is governed by model parameter \1, while the subsequent ones are characterized by X2.
2.1.2 Jet reconstruction during inference
In addition to the generative model, also need to be able to assign likelihood value to proposed jet clustering (binary tree) during inference. To do this use the same general form for the jet's likelihood based on product of likelihoods over each splitting. In order to evaluate this we need to first reconstruct each parent from its left and right children. Different tree topologies give rise to different tp values for the inner nodes and thus different likelihoods. Notably, in Ginkgo the likelihood of a tree is expressed in terms of the product (in linear space) of all splitting likelihoods (specified in Eq. and referred as the partial likelihood) of a parent into two children. The likelihood of a splitting connects parent with child nodes (i.c. we have a likelihood of sampling child with squared mass given parent with squared mass tp). Thus, parent and child nodes not independent. However, splitting likelihoods of different parent nodes are independent, given tree. For set of observed energy-momentum vectors X = (æ1, .NN} (leaf nodes), parameters 0, and a tree topology T, the likelihood of a splitting history can be evaluated efficiently.
2.2 Parent node splitting likelihood reconstruction
At inference time, parent node with energy-momentum vector ap is obtained by adding its children values (zL, ZR) as shown schematically in Fig. 2 (right), and tp is calculated deterministically given zp. The likelihood of a parent splitting into a left (right) child is defined as follows:
f(t.)),th)
tp > teut
El1e, N,,,,,out (xx (p)1p)
(2)
F.lteur:tp)
tp I teut
where i e {4, R} (note that to fix a degree of freedom, to = tp, and E = (vtp - ?)z)²) and teut is the cutoff mass squared scale for the binary splitting process stop (if t, S teuts the corresponding node is a leaf of the binary tree). We introduce the cunulative density function F,(lour for a given generative process stop, İ.e. the probability of having sampled value of tp < leut;
1 e -tteut/pp
tp > text
Fa (p)
1 1-e -
(3)
tp S text
5

---

A Published in Transactions on Machine Learning Research (12/2024)
RESAMPLE
PROPOSE
WEICITTING
A B
B
B
C
C
C
C
D
D
D
D
]
A
1
A B C
A
A
B C
B C
D
D
D
D
A C
A
B
B
C
C
C D
D
D
D
Figure 3: Summary of the CSMC framework: A total of K partial states {s*] are retained as collections of tree structures encompassing the data set. A partial state is defined as collection of trees, which start out as singleton particles A, B, C and D. Each iteration within Algorithm comprises three key slages: (1) resampling partial states based on their importance weights {w<1K (2) proposing an expansion of each partial state form now 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 jet consisting of observed four particles, denoted as A, B, C, and D.
Taking this into account, the probability F(tL, Lu, X, teut, tp) of a parent node splitting at inference time can reconstructed as the product the probability of splitting (1 Fletuus (p)) times the likelihood of a parent splitting into left and right children as follows:
F(t,, tli,, tp) - (1 Fs(teut,tp)) Eltz,,, A, teu l(tR, tetua,
(4)
4T
where the factor I comes from the likelihood of sampling uniformly over the two-sphere during the 2- दt body particle decay process. Also, at inference time, given two particles, we assign t - max(tz, ta and Ln + min(Lr,4n)
2.3
Approximate Inference
2.3.1 Bayesian jet reconstruction.
The goal of jet reconstruction to infer the splitting history and properties of subatomic particles pro- duced during high-energy collisions. This involves estimating various paramcters such as particle momenta, positions, and lypes. Bayesian methods provide natural framework for modeling uncertainty and incor- porating prior information into the reconstruction process. Let X {xI,... .IN} denote the matrix of observed monegy-m vectors of the Ginkgo model. The posterior distribution can be expressed follows:
P(r|X. P(X A))((()
P(X
Calculating the denominator requires marginalizing over (2N- 3)!! distinct jet topologies which is intractable. The two interconnected lasks for Bayesian jel reconstruction are: (i) computing the normalization ion constant P(X[A) by marginalizing all possible candidate topologies:
P(X[A) = Σ P(X)1)))((1)),
(6)

---

Published in Transactions on Machine Learning Research (12/2024)
and (ii): learning optimizing the likelihood in Eq. obtained by marginalizing Eq. X
arg max log P(X Variational Inference (VT) offers approach to tackle both tasks for non-trivial poste- λελ rior distributions.
2.3.2 Variational Inference.
VI is a method used to estimate the posterior distribution P(T, X(X) when its direct computation is in- tractable (due to the complexity of marginalizing the latent variables r). To address this challenge, VI introduces a tractable distribution T(),T(X) create lower bound LELBO on the log-likelihood:
P(T,
log P(X) 2 LELBO X)= = B Q[T,
(7)
In the context of trained (Kingma & Welling 2013 Rezende et al. 2014). To approximate the expectation in Eq Monte Carlo samples from Q(T, X are averaged, and these samples are reparameterized using deterministic function of random variable that is independent. of 7.
Variational Bayes (AEVB, both Q(T, and P(A, T, are jointly
Obtaining feasible approximation for jet structures can be complex task, leading us to modify and adapt CSMC.
2.3.3 Combinatorial Sequential Monte Carlo.
CSMC, tailored for phylogenetic tree models, approximates a sequence of increasing probability spaces, ultimately aligning with Eq. 5 (Wang al. 2015p. CSMC employs sequential importance resampling across {r) steps to approximate both the unnormalized larget distribution and its normalizat ion constant, denoted as ||T]], constituting the numerator and denominator in Eq. by K partial states {s Jk=1 E S, to form a distribution (see Wang et al. (2015) or Appendix
k K
7, |
W (6)
Vs
(8)
CSMC, in contrast to standard SMC techniques, manages a combinatorial set representing the realm of tree topologies alongside the continuous branch lengths- both of which are characteristic features of phyloge- nies (Wang et al. 2015). Partial states (Monte Carlo samples) are resampled at each rank ensuring samples remain in high-probability regions, and importance weights are defined as:
(sa)) (3,-T11) (
- w
where q(s d specifics proposal distribution, a E f1..... ......KK K} denote resampled ancestor indices with P(ay_1 i) w/-1/ wi_., and -ー is an overcounting correction defined in Wang et al. (2015)
Resampled states are then extended via proposal distribution simulations (see Fig. This framework allows for the construction of an unbiased estimate of the marginal likelihood, converging in L2 norm:
R
K
ECSMC = || = II
Σ
w
- ||/||-
(10)
K
The CSMC method is only applicable for sampling toplogies to marginalize over the space of phylogenetic trees, leading us to adapt VCSMC to perform VI.
Variational Combinatorial Sequential Monte Carlo. Expanding on the foundation laid by CSMC, Moretti et al. (2021) introduces Variational Combinatorial Sequential Monte Carlo (VCsMC) approach

---

Published in Transactions on Machine Learning Research (12/2024)
to learn distributions over phylogenetic trees. VCSMC employs CSMC means to create unbiased estimator for the marginal likelihood:
LCSMC =B [csMcc
(11)
In the same work, Moretti al. (2021) | introduces Nested Combinatorial Sequential Monte Carlo (NosMc), an efficient proposal distribution, providing an eract approrimation to the intractable locally optimal proposal C The VI algorithms VCSMC and VNCSMC that utilize the estimators 2CSMC and 2NCSMC each introduce a structured approximate posterior that exhibits factorization across rank events. Each state, denoted as 8,, is uniquely characterized by its topology, a collection trees forming a forest, and the corresponding branch lengths. To facilitate reparameterization, discrete terms are either removed from gradient estimates or transformed into Gumbel-Softman random variables.
for CSMC. We provide review of NGsMC in Appendix
3 Methods
7(5) II T(1) GGB
Scction 3.1 adapts the CSMC approach to perform inference on jet tree structures. Section 3.2 refor- mulates VCSMC for inference on global parameters. Section 3.2.1 utilizes VCSMC methodology to learn parameters as point estimates. Section 3.2.2 defines prior on the model parameters to construct a vari- ational approximation on both global and local pa- rameters. The resulting approach is interpreted variational geeudo-marginal method establishing connections between pscudo-margina, methods (An- drieu & Roberts 2009) and Variational Combinato- rial Sequential Monte Carlo [Naesseth et al. 2018
= 11 F.)
:e(A...))
= FA. = FA
E
11
Figure 4: The likelihood FA for a sub-tree defined on leaf nodes D, E, F and G is defined as the recursive product of splitting likelihoods FB and Fc. The in- termediate target T(33) for the partial state 83 also includes the probability of singletons H and I denoted Fh and F1.
Moretti et al. 2021).
3.1 Inference on Tree Structures
Sequential Monte Carlo (SMc) methods (Naesseth et al., 2019 Chopin Papaspiliopoulos 2020) are designed to sample from a soquence of probability spaces, where the final itcration converges to the target distribution. However, adapting CSMC for the Ginkgo model requires crucial modification: ensuring that the splitting likelihood at the final coalescent event reflects the dependence on the entire sub-tree splitting history, not just the most recent split. This dependence is essential for accurately capturing the recursive structure of the tree. Recall that the splitting likelihood is defined as the product of all splitting likelihoods (see Eq. 4..
To achieve this, we reformulate the splitting likelihood to explicitly account for previous splits, as depicted in the recurrence relation:
1 .1 (11
F(tL, tR. CRL,,RRA,ten tewt) t.tuet)
(1-Fs(teun) Fs tp)) x E(tL, X, teuts tb) e(tRR X, teuts t/t
(12)
4T
x F(tri, LLR, tett ) F((RI, tRR. X, tett
In the above, the pair i,j E (L, (I. RI) R) defines ty as the mass squared of the j child of the current coalescent node and f (lr, LR, X, leut for leaf nodes simply evaluates to 1. Eq. 12 represents the tree (sub- tree) likelihood (with root node having squared mass tp) as the recursive product of splitting likelihoods. Each term brings its normalization, and the overall normalization is correctly expressed as the product of individual ones. In practice we use dynamic programming to maintain a running sum of cumulative log probabilities across rank events. The CSMC resampling step illustrated in Fig. 3 is now dependent upon the full sub-tree splitting history, reflecting the recursive nature the coalescent process.
In a slight abuse of notation, the probability (s))) of partial state sf (recall E {1,... .....NN- coalescent event and à denotes the Monte Carlo sample) is defined as the product of the probabilities of all
NN 1] 11 denotes the
8

---

Published in Transactions on Machine Learning Research (12/2024)
disjoint Trees in the forest 8: (8)) II (G:)). An illustration of likelihood for sub-tree along with GE
the likelihood of a partial state provided in Fig. 4 The NCSMC algorithm defined in Morctti et al. (2021)
can similarly be adjusted ensure compatibility with this modified framework. The resulting estimators are unbiased and consistent, for proofs see Wang et al. (2015) and Moretti et al. (2021).
3.2
Inference on Global Parameters
In Section 3.2.1 we define a variational approximation on topologies, while in Section 3.2.2 we outline fully Bayesian inference using a variational pseudo-marging framework.
3.2.1 Maximum Likelihood
We introduce a variational approximation on т, using CSMC and the AEVB framework to optimize X. Using
Eq. and Eq. 10 to define the weights and estimator respectively, along with Eq. 11 and Eq. 12 to evaluate I and form the ELBO, we define a variational family:
K II 1=11
Qá.ú (si:k I:K (1:t-11
(s)) x
IIIII
W
(13)
E w
The full factorization of go. (s. (s/"") choice of CSMC or NCSMC as an inference algorithm. We utilize our adaptation of CSMC and NCSMC to form the two objectives Lesxe and LNCSMC.
( ) is written in Eq. 22 and Eq. 23 of the Appendix and depends on the
3.2.2 Fully Bayesian Inference Using a Variational Pseudo-Marginn Framework
Simulators rooted in quantum chromodynamics are frequently calibrated to align with the data. However; training a simulator separately from the inference process can lead to inefficiencies. To address this, we propose a modification of the posterior distribution defined Eq. 5 to include a prior on along with variational parameters to learn the proposal distribution in Eq. 13 We define a lormnormal distribution over A log VAlp, Σ) so that can be marginalizod along with т. The target distribution can now be specified follows:
P(T, P(X P)
(14)
P(X)
The generative model parameters can be defined as = A = {u, }} or as the output of a neural network and the proposal paramcters Φ = {a. ΣJ used in the variational approximation to Eq. 14 can be shared or separately trained.
The marginll framework is designed to sample from a posterior distribution such as the one defined in Eq. 5 when the marginal likelihood p(X[A) cannot be evaluated directly. We would normally be intcrested in computing the posterior distribution over splitting topologies and decay parameters defined in Eq 5) however, the marginal likelihood p(X(X) is intractable. Given access to function g(u; X, X) accepting random numbers ~ r(u) that can be evaluated pointwise, assume (u; X, X) returns a non-negative unbiased estimate of P(X[A):
E [0 X, A)] r(u)
0(u; X, p((I))
(15)
In our setup, gw: X, A) defined in Eq. 13 and the auxiliary random variables u = (s a1:t-1) are generated via the CSMC or NCsMC algorithm. Let p(A,u) u) be joint target distribution over and u:
p(),u) u) g(u; X, A)r(u)((A)
(16)
(
The integral of the expression above equals one, and its marginal distribution corresponds to the posterior distribution:
=(A)= = / g(u; X, A)r(u)p(A)
du p(X(A)p(A)
(17)
p(X)
x(X)
9

---

Published in Transactions on Machine Learning Research (12/2024)
-30
-55
-55
VNCSMC
VNCSSIC
-60
PXXIT A)
-00
-05
-70
VNCSMC
-80
-75
-75
-75
-35
-3
-55
-00
100
200
Greedy
Beam
Epoch
Figure 5: Left: Scatterplot comparing log-conditional likelihood of VNCSMC with K, M (256. Greedy Search and Center: Scatterplot comparing log-conditional likelihood of VNCSMC with K,M (256, 1) 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 = {256) (and M = 1) samples averaged across 5 random seeds. VNCSMC achieves convergence in fewer epochs than VeSMC and yields higher values, all while maintaining lower stochastic gradient noise. Additional experiments demonstrating the effect of K appear in Fig. of the Appendix.
This implies that by employing nearly any approximate inference technique to Eq. 16 the resulting marginal
distribution of that approximation will serve an estimate of the actual target distribution. The pseudo- marginal framework by Andrieu & Roberts (2009) designs a Markov Chain (0', u²) with Eq. 16 as its target distribution.
Given the conditional likelihood p(X(A,T) we could run MCMC only on the parameter p(A). Instead, we take the approach of sampling K topologics nh p(T[A) so that
1 K
p(X(A)p(A),
where an explicit approximation to p(T[A) is defined by the modified CSMC algorithm. This is a form of variational pscudo-marginal setup where we are interested approximately marginalizing out all jet structures. Our distribution estimator which marginalizes interpreted analogously.
4 Experiments
The industry standard in particle physics uses agglomerative clustering techniques, which are greedy CCaciiri et al. 2012). Beam Search provides straightforward and significant improvement. Thus, we consider both Greedy and Beam Search standard, relevant and efficient baselines, also applied in cited works Greenberg et al.
2021 We simulated 100 jets using Ginkgo running comparisons with Greedy Search, Beam Search and Cluster Trellis.
2020
4.1
MAP estimate
Fig. (left) provides scatterplot comparing log-conditional likelihood values. Across 100 simulated jets, VNCSMC with K, M = (256, 1) returns a higher likelihood on all 100 cases against Greedy Search (left) and 99 cases against Beam Search (center). Notably, VNCSMC simultaneously conducts inference and À learning, feature lacking in Greedy Search and Beam Scarch, which rely on the user providing À values. Furthermore, VNCSMC yields probability distributions over topologies, while Greedy Search and Beam Search yield single topologies.
10

---

Published in Transactions on Machine Learning Research (12/2024)
VNCSMC VCSMC
Inferred vs Exact Likcllhood Ratio
Running Time (5) In Log Scale
133
Estimateed Likellhood
-2
Beam Search
VCSMC
VCSMC
Greedy
Gready Search
Beam
Cluster Trellis
Number of Monte Carlo Samples
Cluster Treilis MAP
Number of Observed Particles (Leaf Nodes)
Figure 6: Left: CSMC and NCSMC converge the exact marginal likelihood K increases (log scale). Center: Log conditional likelihood returned by VCSMC compared with the exact MAP clustering returned by the cluster trellis Greenberg et al. (2020), Greedy Search and Beam Search, for dataset of 100 simulated jets generated with Ginkgo. Even with K = 256 and N = 12, VCSMC closely approximates exact cluster trellis values for specific jets. Right: A comparison of the running times, highlighting that VCSMC on N = 64 significantly outperforms cluster trellis on = 12. As N increases, the cluster trellis quickly becomes impractical due to its exponential complexity.
Fig. 5 (right) shows the log conditional likelihood To for VCSMC (blue) and VNCSMC (red) with K = 256 (and M 1 for VNCSMC) samples averaged across 5 random sceds. VNCSMC converges in fewer epochs and achieves higher likelihood values with lower gradient noise. Additional experiments illustrating the effect of K appear in Fig. g in the Appendix. Larger K yields higher log fog A) values and reduces stochastic gradient noise.
4.2 Running Time and Complexity.
We generated jets with N = [4. .64 leaf nodes and profiled the running time of VCSMC, Cluster Trellis, Greedy Search and Beam Search averaged across 3 random seeds. All experiments were performed on Google Cloud Platform n1-standard-4 instance with an Intel Xeon CPU vCPUs and 15 GB RAM without leveraging GPU utilization.
Fig. 6 (left) illustrates convergence of the CSMC and NCSMC conditional likelihood to the exact cluster trellis marginal likelihood as K increases. Fig. plots the inferred Log-Normal Pseudo-Marginal Distribution (H1, p2) of the Heavy Resonance Jet, estimated through NCsMC. Contours of the
for the parameters log-conditional likelihood are shown with stochastic gradient steps taken on LNGSMC highlighted in rod.
Next, we compare Fig. G (center) the log conditional likelihood returned by VCSMC, Greedy Search and Beam Search with the exact Maximum Posteriori (MAP) clustering value calculated with the cluster trellis technique described in GGrenneer et al. (2020), for dataset of 100 jets. We see that VCSMC returns high quality hierarchies, with log likelihood values close to the exact ones.
Fig. 6 (right) reports the rumning times on a log scale in seconds. VCSMC is an order of magnitude faster than Beam Search on N = 20 leaf nodes. Beam search entails managing a list of log-likelihood pairs at each level and for each beam size The given list is sorted, iterated over, and selectively only a single topology is retained at a time. This incurs a complexity of O N log + bN3 log N). Typically b > N, in which case the complexity can be prohibitively slow, but for b < N² log N, the complexity becomes O(b log N). In contrast, NCSMC is O(KN9 and the CSMC is OKKNM), where K, N and M denote the number of Monte Carlo samples, the number of leaf nodes (observed particles), and the number of subsamples.
5 Discussion
11

---

Published in Transactions on Machine Learning Research (12/2024)
In all experiments, feut is an exogenous variable, and in full physics simulations its value is determined by the energy scale where there is a qualitative change in the theory that describes the process, i.e. we switch from a showering to a hadronization process, described by different models of physics. The difference between the energy of the initial state particle (root of the binary tree) and this energy scale (feue) affects the number of final state particles the Monte Carlo sampling. Given that Ginkgo does not include full quantum chromodynamics (QCD) modeling of the splitting likelihood, we choose lent control the distribution on the number of final state particles and mimic full physics simulation processes. As such, the exact value is not relevant and the developed algorithms capabilities are independent of it.
Simulators rooted in QCD present significant chal- lenges due to their complex likelihoods. Rewriting these simulators a substantial endeavor, demand- ing considerable expertise and effort, with special- ists often dodicating entire carcers mastering the intricacies of QCD. Real-world data utilization ne- cessitates strict adherence to the models embedded within these simulators, further complicating the task. Moreover, the lack of currently available pow- erful quantum computers adds another layer of dif- ficulty, prompting high-energy researchers, such as those at IRIS- HEP, to rely on approximations like Ginkgo (Cranmer et al. 2019a).
The CSMC method (Wang et al. 2015 and the Variational Combinatorial Sequential Monte Carlo (VCSMC) method (Moretti et al., 2021) both use a specific generative model for biology based on a continuous lime Markov chain. Our paper repre- sents the first adaptation of Sequential Monte Carlo methods for jet reconstruction in particle physics. This is achieved by adapting the Ginkgo model and redefining the parent node splitting likelihood re- cursively (Eqq. [2] to ensure that, at the final CSMC rank event, the support of the proposal distribution aligns with that of the target distribution.
A
Figure 7: Inferred log-normal peooudo-margina distri- bution for the parameters ý (2), of the heavy resonance jet, estimated using NCsMC. Contours rep- resent the log-conditional likelihood, with stochastic gradient steps on LNeSMC highlighted in red.
5.1 Conclusion
Variational and pseudo-marginal methods are broadly applicable across a range combinatorial and continu- ous problems. In Bayesian phylogenetics, these methods facilitate efficient estimation of marginal likelihoods over equivalent factorial search space of tree topologies, (2N - 3)!l, allowing for uncertainty in both tree structure and branch lengths [Moretti et al.) 2021 Zhang & Matsen IV 2019 Zhang al. 2021 Addi-
tionally, in social and biological network analysis, variational techniques are central to hierarchical clustering and stochastic block models for complex, structured latent spaces (Zhou 2015 Zhang & IV 2018 Greenberg et al. 2021) Monte Carlo methods, widely used to model self-avoiding random walks on combinatorial struc- tures, provide insight into diffusion and spatial distribution phenomena that bridge into statistical mechanics and more complex particle interaction models (Madras Slade, 1996] Grosborg ct al. 2006). These tech- niques also extend probabilistic graphical models and Bayesian networks, where factorized approximations reduce computational burden, supporting inference
digh-diiensiionll data (Zhang & IV 2018).
We have introdueed the first adaptation of CSMC for unbiased and consistent jet reconstruction, proposing approximate posteriors and variational inference (VT) techniques for both point and distribution estimators. Our approach significantly improves both speed and accuracy, paving the way for broader adoption of variational methods in collider data analyscs. This work not only provides robust framework for jet reconstruction but also sets the stage for future advancements in the application of advanced approximate inference methods to high-energy physics experiments.
12

---

Published in Transactions on Machine Learning Research (12/2024)
References
Christophe Andrieu and Gareth O. Roberts. The pseudo-marginal approach for efficient Monte Carlo computations. The Annals of Statistics, 37(2):697 725, 2009. doi: 10.1214/07-A0S574 URL
httpp://doi.org/ www.o.gg./oooooooooooo.o....
MohammadHlossein Bateni, Soheil Behnezhad, Mahsa Derakhshan, Mohammad Hajiaghayi, Raimon- das Kiveris, Silvio Lattanzi, and Vahab Mirrokni. Affinity clustering: Hierarchical clustering at scale. Advances Newral Information Processing Systems (NeurIPS), 2017.
Mark A Beaumont. Estimation of Population Growth or Decline in Genetically Monitored Populations. Genetics, 164(3):1139- //doi org/ 10. BetS////eeeeiic//1e0
2003. ISSN 1943-2631. doi: (0000/////eti////ee
URL https:
Johannes Bellm et al. Herwig 7.0/Herwig++ 3.0 release note. Eur. Phys. C, 76(4):196, 2016. doi:
Alexandre Bouchard-Côté, Sriram Sankararaman, and Michael Jordan. Phylogenetic inference via sequential
Monte Carlo. Systematic biology, 61:579-93, 01 2012.
Anja Butter et al. The Machine Learning landscape of top taggers. SciPost Plays., 7:014, 2019. doi: 10.21468/Sccli Scy..sst
Matteo Cacciari, Gavin Salam, and Gregory Soyez. The anti-k, jet clustering algorithm. JHEP, 04:063,
2008. doi: 10.1088/ (003-6778/200804444
Matteo Cacciari, Gavin P. Salam, and Gregory Soyez. Fastjet user manual: (for version 3.0.2). The European Physical Journal C, 72)3,, March 2012. ISSN 1434-6052. doi: 10.14400cpt//s0000---------33 1896-2. URL 11w1/11eec/ae// /1tp1//////////000000000--1-1100 doi. wtg/10......
http://dx.oo http:/xxxdoiii.ooooooo
S Catani, Yuri L. Dokshitzer, M H Scymour, and Webber. Longitudinally invariant & clustering algorithms for hadron hadron collisions. Nucl. Phys. B, 406:187 224, 1993. doi: 10.1016/0550-3211303 90166-M.
Nicolas Chopin and Omiros Papaspiliopoulss. An introduction to sequential Monte Carlo, volume 4. Springer, 2020.
K. Cranmer and T. Plehn. Maximum significance at the lhe and higgs decays to muons. The European Physical Journal C, 51(2):415-422 420, Jun 2007. ISSN 1434-6052. doi: /11100//pp///1000500000---020 URL http: http:ttt://d..o..rrrg...... http:///.....////////////0000000
K. Cranmer, S. Macaluso, and D. Pappadopulo. Toy generative model for jets, 2019a. URLtttts:///ndiii cern. ht//evenn/4410///ccmmm contr www.iioso///32ee0 http:///222222ttttaaat/////222/2 The Institute for Rescarch and Innovation in Software for High Energy Physics (IRIS-HEP) Workshop on MIAJets.
pdf
Kyle Cranmer, Sebastian Macaluso, and Duccio Pappadopulo. Toy Gonerative Model for Jets Package. htt.shhw.wo/Ssttaaetaaaaaaaeeaae
2019b.
Kyle Cranmer, Matthew Drnevich, Scbastian Macaluso, and Duccio Pappadopulo. Reframing jet physics with new computational methods. EPJ Web of Conferences, 251:03059, 2021. doi: 10.1051/epjconl/ 202125103059. URL httpe:://ooooooo///// http:///.......eeeeoooaa///////2
Kyle Cranmer, Matthew Drnevich, Lauren Greenspan, Sebastian Macaluso, and Duccio Pappadopulo. Com- puting the bayes-optimal classifier and exact maximum likelihood estimator with a semi-realistic generative model for jet physics. Machine Learning and the Physical Sciences, NewrIPS, abs abss/2002.1100 2022. URL
Cranmer, Kyle, Drnevich, Matthew, Macaluso, Scbastian, and Pappadopulo, Duccio. Reframing jet physics with new computational methods. EPJ Web Conf., 251:03059, 2021. doi: 10.101///ppoomff222220003
URL
13

---

Published in Transactions on Machine Learning Research (12/2024)
Vu Dinh, Arman Bilge, Cheng Zhang, and Frederick A. Matsen, IV. Probabilistic path Hamiltonian Monte Carlo. volume 70 of Proceedings of Machine Learning Research, pp. 1009-1018, International Convention Centre, Sydney, Australia, 06 11 Aug 2017. PMLR.
Yuri L Dokshitzer, G D Leder, Moretti, and B R Webber. Better jet clustering algorithms. JHEP, 08:1, 1997. doi: 10.1088/ /a////11226678801// 6708/1997/08$ /08/001.
Avinava Dubey, Qirong Ho, Sincad Williamson, and Eric P Xing. Dependent nonparametric trees for dynamic hierarchical clustering. Advances in Neural Information Processing Systems (NeurIPS), 2014.
Kumar Avinava Dubey, Michael Zhang, Eric Xing, and Sinead Williamson. Distributed, partially collapsed mcmc for bayesian nonparametries. In International Conference on Artificial Intelligence and Stalistics. 2020.
Stephen D Ellis and Davison E Soper. Successive combination jet algorithm for hadron collisions. Phys. Reu. D, 48:3160-3166 316616 1993. doi: D0.100//hysRccD.45....3
T. Gleisberg, Stefan. Hoeche, F. Krauss, M. Schonherr, S. Schumann, F. Siegert, and J. Winter. Event generation with SHERPA 1.1. JHEP, 02:007, 2009. doi: 10.1088/1126-600//200000 /009//02/007.
Spence Green, Nicholas Andrews, Matthew R. Gormley, Mark Dredze, and Christopher D. Manning, En- tity clustering across languages. In Conference of the North American Chapler of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLTT), 2012.
Craig S. Greenberg, Sebastian Macaluso, Nicholas Monath, Ji-Ah Lee, Patrick Flaherty, Kyle Cranmer, Andrew McGregor, and Andrew McCallum. Data Structures \& Algorithms for Exact Inference in Hier- archical Clustering. 2 2020. URL
Craig S. Greenberg, Sebastian Macaluso, Nicholas Monath, Avinava Dubey, Patrick Flaherty, Manzil Zaheer, Amr Ahmed, Kyle Cranmer, and Andrew McCallum. Exact and approximate hierarchical clustering using a*, 2021.
A. Y. Grosberg, K. Kremer, and M. Kardar. The physics of polymers: from gaussian chains to polymer gels. Reviews of Modern Physics, 78(2):393-414 2006.
Zhiting Hu, Ho Qirong, Avinava Dubey, and Eric Xing. Largo-scale distributed dependent nonparametric trees. International Conference on Machine Learning (ICML), 2015.
A. Ord J. Stuart and S. Arnold. Kendall's advanced theory of statistics. In Vol 2A (6th Ed.) (Oxford University Press, New York, 1994.
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes, 2013.
Tuan Anh Le, Maximilian Igl, Tom Rainforth, Tom Jin, and Frank Wood. Auto-encoding sequential Monte
Carlo. In International Conference on Learning Representations, 2018.
Gilles Louppe, Kyunghyun Cho, Cyril Becot, and Kyle Cranmer. QCD-Aware Recursive Neural Networks for Jet Physics. JHEP, 01:057, 2019. doi: (HE00TJJJE0120011101
Chris J. Maddison, Dieterich Lawson, George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Whye Teh. Filtering variational objectives. 2017.
Neal Madras and Gordon Slade. The Self-Avoidie Walk. Birkhäuser, 1996.
Nicholas Monath, Ari Kobren, Akshay Krishnamurthy, Michael R Glass, and Andrew McCallum. Scalable hicrarchical clustering with tree grafting. In Knowledge Discovery & Data Mining (KDD), 2019.
Nicholas Monath, Avinava Dubey, Guru Guruganesh, Manzil Zaheer, Amr Ahmed, Andrew McCallum, Gokhan Mergen, Mare Najork, Mert Tezzian, Bryon Tjanaka, et al. Scalable bottom-up hierarchical clustering. Xixiv preprint ar Niv:2010. 1100000.11112 11821, 2020.
14

---

Published in Transactions on Machine Learning Research (12/2024)
Nicholas Monath, Manzil Zaheer, Kumar Avinava Dubey, Amr Ahmed, and Andrew McCallum. Dag- structured clustering by nearest neighbors. In International Conference on Artificial Intelligence and Statistics, 2021.
Antonio Khalil Moretti, Zizhao Wang, Luhuan Wu, Iddo Drori, and Itsik Pe'er. Particle smoothing variational objectives. CoRR, abs/ 1909.9999,,,34, 2019.
Antonio Khalil Moretti, Zizhao Wang, Luhuan Wu, Iddo Drori, and Itsik Pe'er. Variational objectives for Markovian dynamics with backward simulation. European Conference Artificial Intelligence, 2020.
Antonio Khalil Moretti, Liyi Zhang, Christian A. Naesseth, Hadiah Venner, David Blei, and Itsik Pe'er. variational combinatorial sequential monte carlo methods for bayesian phylogenetic inference. In Cassio de Campos and Marloes H. Maathuis (eds.), Proceedings of the Phirty-Seventh Conference on Uncertainty in Artificial Intelligence, volume 161 of Proceedings of Machine Learning Researeh, pp. 971-981, PMLR, 27-30 Jul 2021. URL
mlr wwwpppesssssse /ttess//111/morett22l.... html
C. A. Nacsseth, S. Linderman, R. Ranganath, and D. Blei. Variational sequential Monte Carlo. volume 84 of Proceedings of Machine Learning Reseorch, 2018.
C. A. Naesseth, Lindsten, and T. B. Schön. Elements of sequential Monte Carlo. Foundations and Trends in Machine Learning, 12(3):307 392, 2019.
Stanislay Naumov, Grigory Yaroslavtsev, and Dmitrii Avdiukhin. Objective-based hierarchical clustering of deep embedding vectors. arXiv preprint Xrx:::2012 2.00002..000 2.04466 2020.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximato inference in deep generalive models, 2014.
Fredrik Ronquist, Maxim Teslenko, Paul Mark, Daniel Ayres, Aaron Darling, Sebastian Höhna, Bret Larget, Liang Liu, Marc Suchard, and John Huelsenbeck. MrBayes 3.2: Efficient, bayesian phylogenetic inference and model choice across a large model space. Systematic biology, 61:539 42, 03 2012.
Torbjorn Sjostrand, Stephen Mrenna, and Peter Z. Skands. PYTHIA 6.4 Physics and Manual. JHEP, 05: 026, 2006. doi:
Minh-Ngoc Tran, David J. Nott, and Robert Kohn. Variational bayes with intractable likelihood, 2016.
Shikhar Vashishth, Prince Jain, and Partha Talukdar. Cesi: Canonicalizing open knowledge bases using embeddings and side information. In Proceedings of the 2018 World Wide Web Conference on World Wide Web, pp. 1317 1327. International World Wide Web Conferences Steering Committee, 2018.
Liangliang Wang, Alexandre Bouchard-Cött and Arnaud Doucet. Bayesian phylogenetic inference using a combinatorial sequential Monte Carlo method. Journal of the American Statistioal Association, 01 2015.
Shijia Wang and Liangliang Wang. Particle Gibbs sampling for Bayesian phylogenetic inference, 2020.
Cheng Zhang and Frederick A. Matsen IV. Generalizing tree probability estimation via bayesian networks, 2018. URL https://arxiv. wbwg/bb///aaaa 08800007834
Cheng Zhang and Frederick A Matsen IV. Variational Bayesian phylogenetic inference. Tn International Conference on Learning Representations, 2019.
Yixin Zhang, Vladimir N Minin, and Frederick A Matsen IV. Variational inference in phylogenetics. Tn Proceedings of the 38th International Conference Machine Learning, pp. 12311 12321. PMLR, 2021.
Yuchen Zhang, Amr Ahmed, Vanja Josifovski, and Alexander Smola. Taxonomy discovery for personalized rocommendation. In Proceedings of the 7th ACM international conference Web search and data mining.
ACM, 2014.
Mi Zhou. Infinite edge partition models for overlapping community detection and link prediction. Proceedings of the 18th International Conference Artificial Intelligence and Statistics, 68:1135- 1143, 2015.
15

---

Published in Transactions on Machine Learning Research (12/2024)
A Ginkgo Generative Model
The Ginkgo generative process is outlined below: The 2-body decay in the parent rest frame is defined using
Algorithm Toy Parton Shower Generator
Require: parent momentun fă, parent mass squared (p, cut-off mass squared leur, rate for the exponential distribution X, binary tree tree
1: Function NoDEPRoCESSINC(yy); tp, leut, X, tree)
2: Add parent node to tree.
tp teut then
Sample t, and ig from the decaying exponential distribution.
Sample unit vector from uniform distribution over the 2-sphere. Compute the 2-body decay of the parent node in the parent rest frame.
7: Apply a Lorentz boost to the lab frame to each child.
8: call NODEPROCESSING(p"" L. teut, X, tree)
9. call NoDEPROCESSING((py; tR, teut, X, tree)
10: end
momentum pyp = P + = (V3,0,0,0). Due to merrg-momeente conservation the child energies are given
by
V8
tr
t
E1 = 1
(18)
2
$
ER = V
1 tR
t
(19)
2
and the magnitude of their 3-momentum defined
V$g 3 2
v V
2(tz + +tR)
(tz tR)²
(20)
s2
The left and right child momentum are given by 1 = (FL, p) and P4 (En. pp) in the parent rest frame. The Lorentz boost E, and 7B 1pll/ Fp- For more information see Cranmer, Kyle et al. (0021
Cranmer et al. 2019b
16

---

Published in Transactions on Machine Learning Research (12/2024)
B Combinatorial Sequential Monte Carlo
We provide an overview of the CSMC algorithm from Wang et al. (2015
B.1 Partial States and the Natural Forest Extension
Definition 1 (Partial State). A rank E {0, , N partial state, symbolized as (ti, X;), represents collection of rooted trees and adheres to the following three conditions:
1. Partial states of different ranks are disjoint, mcaning that for any two distinct ranks, and there is no overlap between the sets of partial states, written as Vr + 8, S, ns, =
2. The set of partial states at the smallest rank consists of only one element, denoted So L.
The set of partial states at the final rank == corresponds the target space X.
The likelihood, as represented in Eq. 12 and the probability measure и are specifically defined within the scope of the target space, denoted Sn X. It's important to note that these dcfinitions apply exclusively to the targel space of trees and not to the broader sample space encompassing partial states, denoted SpcR, which consists of forests containing disjoint trees. The Sum-Product algorithm is primarily utilized to derive a maximum likelihood estimate for a tree. However, partial states are explicitly characterized as collections of these disjoint trees or leaf nodes. To extend the target measure т to encompass the sample space Sren, a practical approach is treat all elements of the jump chain as trees, as elaborated Wang ett al. 2015
Definition 2 (Natural Forest Extension). The natural forest extension, denoted NFE, expands the target mcasure into forests by forming product over the trees contained within the forest:
(s) II (t))
(21)
(N....)
One notable advantage of the NPE is its ability to transmit information from non-coalescing elements to the local weight update.
Algorithm 2 Combinatorial Sequential Monte Carlo
Input: Y = (Y,,.. Y1s) NaM
1: Initialization. wk, st 1, wh < 1/K.
2: for == =0 to A = N - do
3: for = to K do RESAMPLE
w
P(a. = =
Σ1=1 w
5: EXTEND PARTIAL STATE
8 g(-18,111) +
6: COMPUTE WEIGHTS
0y w
( ( (
イ(,
end for
8: end for
Output: 1:K
I:K
17

---

Published in Transactions on Machine Learning Research (12/2024)
C Nested Combinatorial Sequential Monte Carlo
We provide a review of the NCSMC algorithm from Moretti et al. 2021 The NCSMC method performs standard RESAMPLE step (line 4), similar to CSMC methods, iterating over rank events. In each iteration, NCSMC explores all possible one-step ahead topologies (") and samples sub-branch lengths for each of them (line 5-7). Importance sub-weights or potential functions are evaluated for these sampled look-ahead states (line 8). The aucestral partial state is then extended to a new partial stale by choosing a topology and branch length based on their respective weights (line 11). The final weight for each sample is calculated by averaging over the potential functions (line 12). For visual representation of this procedure, please refer to Fig. 8
Algorithm 3 Nested Combinatorial Sequential Monte Carlo Input: Y -Y,,... Y E QNNM 1: Initialization. Vk, s0 {-,, w + 1/K.
2: for to R = N - do
3: for k == to K do
w
4: RESAMPLE P(a = i) = Σ Di= #r-1
for i i 1 to L (N=7) do
for =1 to M do
FORM LOOK-AHEAD PARTIAL STATE
A.m 1 ~
((-181111
8:
COMPUTE POTENTIALS
x (5 (s
k.
(s
a
q( t.
[i]] )
end for
10:
end for
11:
EXTEND PARTIAL STATE
6 [1]
We [
PP(T =
k.m
-
12:
COMPUTE WEIGHTS
w ML
w k
13: end for
14: end for
1:K Output: 8 W1.R
D Approximate Posteriors
The proposal distribution for point estimator adaptation of CSMC and the corresponding approximate posterior for VCSMC corresponding to Eq. 5 can be written explicitly follows:
K
R
K
)
( (TA1TT²TT)
Qo (Ti.i 1:
1 = II
90 (T²))
To
(22)
k=1
-=
w
1-1
18

---

Published in Transactions on Machine Learning Research (12/2024)
ENUMERATE TOPOLOGIES SUBSAMPLE BRANCH LENGTHS
COMPUTE POTENTIALS
A
A
4
B
B
C
C
C
]
D
D
ABG CC
A
A
4 A B
]
B
B
C
C
C
D
D
D
D
B
A
A
B 7
B
C
D 1
D
D
Figure 8: Illustration of the NCSMC framework: In NCSMC, all possible one-step ahead topologies, which amount to (N=") in total, are systematically enumerated. For the state A, B, C, D, the enumerated topologies include (top): A, B, C, D, (center); A, B, C, D, and (bottom): B, A, D. Subsequently, M = sub-branch lengths are stochastically sampled for cach odge. Following this, the sub-weights or potentials are computed (right), and a single candidate randomly selected proportional to its sub-weight (or potential) to create the new partial state.
In the above, ah_, denotes the ancestor index of the resampled random variable and the partial state s = T is sampled by proposing forest T, (.(-TT,-1 from UNIFORM distribution. Similarly, the proposal
distribution for the global posterior defined in Eq. 14 is expressed where X go (-12,_111) = log (- E):
= i
K (T))) go (Af)
R
K
Qa. (TH
I:K IKK ARR
1)
W 1= 1
90 (
r=
k=
The NCSMC method detailed in Algorithm can also be used to form an unbiased and consistent estimator of the log- marginall likelihood ENOSMG and variational objective which we refer to as LNOSMC:
R
LNOSMC = log 2NCSMC]
2 ENCSMC =
W
(24)
K
19

---

Published in Transactions on Machine Learning Research (12/2024)
E
Log Conditional Likelihood P(X[T, X)
-70
VCSMC
VCSMC K=16
VCSMC K=32
K=8
-90
VNCSMC
VNCSMC
VNCSMC
K, M=(881)
K, M=(16,1)
K, M=(32,1)
100
200
-70
-80
VCSMC
VCSMC K=128
VCSMC K=256
K=64
VNCSMC
VNCSMC
VNCSMC
M=(64.1)
M=(128,1)
M=
-100
0
100
200
100
200
100
200
Figure 9: Log-conditional likelihood P(X|T, A) values for VCSMC (blue) with K = {8, 16, 32, 64, 128, 256} samples and VNCSMC (red) with K = {8, 16, 32, 64, 128, 256} and M samples averaged across random soods. Greater values K result more constrained ELBO and higher log-likelihood values while reducing stochastic gradient noise. VNeSMC with К
explores higher probability spaces than the likelihood returned by the simulator, depicted by the green trace for reference. VNCSMC achieves convergence in fewer epochs than VCSMC and yields higher values, all while maintaining lower stochastic gradient noise. Notably, even VNCSMC with (K, = (8, 1) in the top-left plot (red) outperforms VCSMC with K = 256 in the bottom-right plot (blue).
F Implementation Details
Invalid Partial States when Coalescing Particles
Physics imposes several constraints on which pairs of particles are impossible to coalesce. We must consider these constraints as we are building trees from leaf root, coalescing particles (represented as nodes) every iteration. The following are the conditions
t >> for any node.
2. 4p > lewt for all inner nodes.
3. 4, > max(11,4) 4)) for all inner nodes.
Recall that the ELBO is a function of the weight matrix, which is of dimensions (R, K), and contains all weights the K particles across R iterations. Each entry in the matrix represents the corresponding weight of some partial state.
In VCSMC and VNCSMC, resampling ensures that we not only extend upon partial states of valid non-zero probability, but we also arrive at K valid trees at the final rank event. We note that both Greedy Search and Beam Search often fail to find any valid trees because they reach a set of partial states where no viable tree can be constructed.
20


Reviews:
[{"review_id": "qPtlT5JAjY", "reviewer": "Reviewer_7rqG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "++:\n1. The paper is clear, with a logical flow of ideas and in-depth explanations. The introduction and background sections effectively ground the reader in both particle physics and Bayesian methods.\n2. The adaptation of CSMC methods for jet reconstruction in particle physics represents a significant contribution, as these techniques have not been widely applied in this area before. The combination of SMC with variational inference provides a fresh perspective on jet physics problems.\n3. The experimental section convincingly demonstrates the efficiency and accuracy of the proposed methods. Results show improvements in speed and precision compared to existing techniques, especially in handling large, complex datasets in collider physics.\n4. The unification of the generative model and inference process using the pseudo-marginal framework is particularly appealing for Bayesian approaches to particle physics.\n\n - -:\n1.  While the topic is undoubtedly important for particle physics, it may have limited appeal for a broad machine learning audience. The application is highly specialized, and the novelty of the methodology might not be sufficiently emphasized for researchers focused on more traditional ML domains.\n2.  The broader applicability of the proposed methods beyond particle physics is not discussed in depth. It would strengthen the paper if the authors could relate the techniques to other fields, especially in machine learning where hierarchical models and Bayesian inference are widely applicable.\n\nIn summary, the paper is well-written, and the contributions are substantial for the field of particle physics. However, to increase its relevance to a broader ML audience, I recommend expanding the discussion of potential applications in other domains, emphasizing the methodological contributions over the specific physics problem, and clarifying the computational complexities.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a novel approach to jet reconstruction in particle physics using CSMC. The authors propose a variational inference algorithm for estimating jet structures and learning parameters, integrated with a pseudo-marginal framework for a fully Bayesian approach. They demonstrate the method’s effectiveness through experiments using simulated collider data, showing improvements in speed and accuracy over existing techniques.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "Just several general comments:\n\n1. To appeal to a wider machine learning audience, the authors should explore how the proposed methods could be generalized to other domains. For example, highlighting applications in hierarchical clustering, Bayesian inference, or NLP could make the contribution more impactful for ML researchers.\n2. While the authors mention the improvements in speed, more detail on the algorithm’s scalability in practice, especially when applied to real-world collider data, would be valuable.\n3. The paper could benefit from additional motivation regarding why ML researchers should care about jet reconstruction in particle physics. How does the advancement in this specialized domain contribute to broader ML challenges?", "broader_impact_concerns": "The work’s broader impact on machine learning, particularly in terms of advancing techniques for Bayesian inference and hierarchical clustering, is not fully explored. The methods presented could be of interest to a wide array of ML fields, but the paper does not emphasize this potential enough. The paper’s impact would be much stronger if it provided explicit connections to other relevant ML areas beyond particle physics.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "kmUuhGFQNc", "reviewer": "Reviewer_2M2k", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- The paper is well-written and mostly easy to follow.\n- I appreciated the discussion in section 2.1, which is essential for those who are not experts in this literature and not familiar with the Ginkgo generative model.\n\n### Weaknesses\n- Difficult to separate the contribution here from past work. I am not an expert in this line of work, however, I found it difficult to figure out the pieces that are new in this work and that was already done without going back to read the work of Wang et. al. (2015) and Moretti et. al. (2021). For example, the fact that the estimators here are consistent and unbiased basically defaults to results from both of the previous papers. Overall, it looks like it section 3.2.2 is really the core of the new inference procedure. \n\n- In Figure 6, and 5, I would've like a comparison to a regular CSMC approach, which doesn't take into account the new likelihood, so that we can evaluate the improvements. I assume greedy search (figure 6) is a standard baseline in this lien of work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a modified combinatorial sequential monte carlo approach for inferring the latent structures of jets. The generative model assumed is based on Ginkgo (from Cranmer et. al.), which models a jet as a binary tree, and splits the a parent node based on a cut off determined by the square of its invariant mass. The standard baseline here are methods based on sequential monte carlo or variational variants. The key insight in this work is to change the way the likelihood is defined in the generative model. Specifically, the likelihood proposed in this work, for a final coalescent event, is forced to depend on the entire sub-tree splitting history (eqn 12). The paper then presents various experiments demonstrating the effectiveness of the proposed modified inference procedures against baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- It would be great for the authors to discuss whether the modified VNCSMC that they present can be useful beyond the Gingko generative model. It seems like the gains we see here comes from the specific changes that was made to the model's likelihood.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "fWMzyjnt2F", "reviewer": "Reviewer_zLT2", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "As a disclaimer, the referee is unsure whether some selling points of the paper have been missed. Fort his reason, it'd be greatly appreciated if the authors could highlight any potential misunderstanding in the upcoming review of the manuscript. \n\n\n# **Strenghts**\n\n- The structure of the paper is well done, and the very detailed introduction makes it accessible to a broad audience. I believe this is an important strength of the paper. \n- The manuscript is well-written, and English is appropriate. I could not identify too many typos or problems associated with the writing. \n-  The visuals from figure 1 and 2 are particularly useful to gain an intuitive understanding of the problem. \n- The numerical experiments clearly show an outperformance of the method discussed in this work compared to established methods.\n\n# **Weaknesses**\n- To my understanding, the contributions and novelty of this work sound somehow limited. For instance, in paragraph 1.6, where the paper's contribution are listed, bulletpoint 1 sounds as if an existing algorithm has been applied to the field of Jet reconstruction. I found the other bulletpoints slightly confusing and it was hard for me to disentangle the actual contributions of this work. I believe that reformulating this paragraph by clearly highlighting the novely and contributions of this work would greatly improve the storyline. \n- I find the writing hard to follow at times especially in the method (Section 3.1 onwards) part where the reader gets lost in the details and cannot appreciate the actual contribution of this work. Also, some of the math seems wrong or, at least, not clear or consistent with the definitions. \n- The problems the authors are trying to solve, and how,  aren't very clear until the very end of the paper. \n- Most of the paper focuses on revising previews works, i.e., 8 pages out of 12. While this is good for the broader audience, it makes it hard to asses the actual contribution of the work unless those clearly stand out and clearly emerge from the discussion.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper provides a new framework to reconstruct the latent hierarchy of jets in particle physics. To reconstruct the structure of a jets in particle collision is an extremely challenging task with a super-exponential growth of possible jet topologies when the number of observed particles increases. Because of this unfavorable scaling, brute force and exact methods are intractable and doomed to fail for large numbers of observed particles. \nIn this work, the authors propose to use a Variational Combinatorial Sequential Monte Carlo approach to infer the topologies of jets. Furthermore, they also propose to learn the system parameters and also propose a pseudo-marginal framework to combine inference and generation of possible topologies. \n\nIn essence, this work suggests to leverage both advantages of generative models to estimate likelihood of proposed topologies and combine this with efficient Sequential Monte Carlo algorithms to achieve better log-likelihood than the standard approaches, e.g., Greedy, Beam search.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "# **Requested Changes**\n- Given the content discussed in paragraph 1.4, I'd recommend to rename this \"Related work\" or similar. \n\n- There are several open questions which require some clarification. The authors could refer to the following section for a list (perhaps not exhaustive) of such questions. \n\n- While I like the graphics from Figure 2, I'd strongly recommend to change the child particles to R (right) and L (left) (instead of C and B) and the parent particle from A to P (parent). This is just a simple modification but would make the understanding of each formula with the corresponding subscript much more intuitive. \n\n- Below equation (2) on page 5, the authors write \"(note that $t^L_P = t_P$ , and $t^R_P = (\\sqrt{t_P} - \\sqrt{t_L} )^2)$\". In this notation, what is, in practical terms,  $t^L_P$ and why does that equal $t_P$? Is this attributing all mass initially from the parent particle to the left one as it is the first one to be sampled in the generation process? More elucidations on this step would be appreciated. \n\n- I find Figure 3 hard to parse. In particular, I cannot straightforwardly related this to the text. In particular I am wondering: what is the 'data set' mentioned in the caption? What are input and output of the algorithm? what are the three different rows? What is the output of this algorithm (right-most part of the figure)? I encourage the other to rethink this figure or make it more self-explanatory.\n\n- Small comment. At the end paragraph 2.3.1 the authors write \"complex posterior distributions\". While it might be obvious to understand what this means, in high energy physics, it is not unusual to find probability distributions with complex term, i.e., chemichal potential in path integral formulation in Lagrangian formalism. I'd advise to change complex to \"non-trivial\" or \"complicated\".\n\n- Below equation (14) the authors introduce the model parameters and the proposal parameters (respectively $\\theta$ and $\\phi$). However, it is highly non intuitive to see where the dependence of $\\phi$ comes into place in equation (14) and moreover, the variational family in equation (13) also depend on some parameters $\\psi$ which have not been defined. \n\n- In equation (17) the resulting distribution depends on $\\theta$ but there's not such dependence on the rhs of the equality.\n\n- I believe the first paragraph of Section 4.1 is a key point for this work and needs to be highlighted more (perhaps also more in the beginning where contributions are listed)!\n\n- MAP abbreviation does not seem to have been defined. \n\n- I'd recommend to split the 4 plots in Figure 6 into a 2x2 layout to enhance readability. I'd also advice to substantially increase fonts and perhaps use different linestyle in the 2 middle plots to aid the visualization for colorblind people. \n\n- For the left-most plot I'd recommend using different markers to, again, enhance readability.\n\n- I'd recommend to be consistent with the colors in the different plots, e.g., VNCSMC should always have the same color. Same holds for the other methods. \n\n- I find the discussion on beam size $b$ in section 4.2 slightly confusing. As a referee with no prior on the topic, I don't know what a beam size is in this context and as far as I can tell this has not been discussed in the manuscript. Moreover, the equation on the 8th line of the paragraph, namely $b\\log < N^2\\log N$ seems incomplete as the $\\log$ on the left hand side seems to have an argument missing. I'd recommend the author to revise this paragraph. \n\n- I found the rightmost plot in Figure 6 to not add much to the discussion. Of course since the performance has been validated using other metrics one would expect a nice convergence of the loss within SGD, but as far as I can tell this does not seem to add any additional insights. Is my understanding correct or am I overlooking something here? On the other hand, I think it'd be interesting to instead compare CSMC and NSMC with their variational counterpart if any metric is suitable for this task. I perceive that the distinction between such algorithms and their variational counterparts has not been higlighted enough and I believe one additional paragraph discussing this would greatly improve the clarity of the paper.\n\n# **Questions to the Authors**\n\n-  On page 2, paragraph 1.3, the authors write: \"While these formulations are helpful conceptually, they are not practical [...] we consider Ginkgo [...]\". I am unsure of what the authors mean here. Does that mean they do not consider real data at all, but only data generated by Ginkgo? Or what is the purpose for this claim? If so, why is this the case? Would it be possible to utilize the proposed framework with real world data, e.g., not synthetically generated? Upon carefully reading the manuscript I guess *No* but I'd be great to have some comments from the authors on this. \n\n- On page 2, end of paragraph 1.3, the authors write: \"Within the analogy between jets and NLP, Ginkgo can be considered as ground-truth parse trees with a known language model.\". I am not quite sure if I fully understand this sentence. I'd appreciate if the authors could elaborate on this. \n\n- The parameter $\\lambda$ is introduced in Paragraph 2.1.1. I found it unclear whether this is a learnable parameter or some fixed parameter which is known a priori, e.g., from the physics of the event, type of detector etc. This is important to know in order to understand if that is a variational (learnable) parameter or just a fixed one.\n\n- At the end of paragraph 2.1.1 the authors write \"For a set of observed energy-momentum vectors [...] parameters ${\\theta}$, and a tree topology $\\tau$, the likelihood of a splitting history can be evaluated \u0000efficiently.\". What are the parameters $\\theta$? Are they model weights? Or are they physical parameters which characterise the physics of the process, e.g., $\\lambda, \\lambda_1, \\lambda_2$? I'd recommend to provide additional explanation around this point in order to avoid any confusion. \n\n- In the Bayesian Inference paragraph, the posterior distribution of a tree topology is expressed as a conditional distribution on the Data X and some parameter $\\lambda$. It is not clear to me, as higlighted before, if this is a learnable parameter or if it can be in principle fixed a priori given the physics of a particular event. I believe this is a crucial point of the work and needs more explanation. \n\n- In equation (7) the authors define the ELBO for the problem as a ratio between $P$ and $Q$. To my understanding Q represent the variational sampler conditioned on the data. On the other end, I do not understand what P is. It'd be great if the authors could be more detailed here and explicitly describe each term. Also, the expectation with respect to Q means that parameters $\\lambda$ and topoligies $\\tau$ are sampled with Monte Carlo? How is this ELBO computed exactly?\n\n- Up to Paragraph 2.3.3 the paper performs a simple review of existing methods (e.g., Ginkgo) and focuses in particular on CSMC and NCSMC and their variational counterpart. I kindly request the authors to correct me in case I have missed any contribution of the manuscript up to page 8. \n\n- From the beginning of Section 3.1, it sounds as if the training is performed by matching CSMC and the Ginkgo output. Yet, it is unclear to the referee what are the parameters optimized in achievieng such a converging criteria. \n\n- Moreover, from section 3.1 I understand that the authors uses the likelihood from the GInkgo model to use the CSMC scheme not only on the previous node but on the full sub-tree spliting history. Is this intuition correct?\n\n- Section 3.2.1: What is AEVB? I think it has not been stated clearly that $\\lambda$ is the parameter which needs to be optimized (see previous questions and comments)\n\n- Is $p(X\\vert\\lambda)$ as defined in equation (15) just a monte carlo estimator of the probability $P(X\\vert\\lambda$? Why is the latter $P$ capitalized?\n\n- I wonder if the reason why the authors do not compare to CSMC is that there is no straightforward way to obtain a log-likelihood by using such methods.\n\n- Last paragraph os Sec. 4.2 the authors write \"Fig. 6 (center right) illustrates convergence of the conditional Csmc and Ncsmc likelihood to the cluster trellis marginal likelihood as K increases.\". Judging from the plot this shows convergence of the **variational** CSMC and NCMSC. Am I misunderstanding something here or is there just a typo in the text?", "broader_impact_concerns": "I could not see any concerns regarding the ethical implications of the submitted work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: For my part, I appologize for the time this paper took in its review process. It was very difficult securing reviews, I think in part because of the topic of the paper.  I want to thank the reviewers and authors for their participation and patience.
Decision: Accept (poster)