TreeGen: A Bayesian Generative Model for Hierarchies
Authors: Marcel Kollovieh
Nils Fleischmann
Filippo Guerranti
Bertrand Charpentier
Stephan Günnemann
Venue: NeurIPS 2025 poster ()
Keywords: generative models
hierarchies
jet clustering
flow matching
Primary area: deep_learning
TL;DR: A novel generative model for hierarchies based on Bayesian Flow Networks.

Abstract: In this work, we introduce TreeGen, a novel generative framework modeling distributions over hierarchies. We extend Bayesian Flow Networks (BFNs) to enable transitions between probabilistic and discrete hierarchies parametrized via categorical distributions. Our proposed scheduler provides smooth and consistent entropy decay across varying numbers of categories. We empirically evaluate TreeGen on the jet-clustering task in high-energy physics, demonstrating that it consistently generates valid trees that adhere to physical constraints and closely align with ground-truth log-likelihoods. Finally, by comparing TreeGen’s samples to the exact posterior distribution and performing likelihood maximization via rejection sampling, we demonstrate that TreeGen outperforms various baselines.

TreeGen: A Bayesian Generative Model for Hierarchies
Marcel Kollovieh!:a 1:2.3
Nils Fleischmann"
Filippo Guerranti
Bertrand Charpentier
Stephan Günnemann 1233,44
School of Computation, Information and Technology. Technical University of Munich Munich Center for Machine Learning Pruna AI
Munich Data Science Institute
Correspondence to: m. www.ooieeettmmmdde.
Abstract
In this work, introduce TreeGen, novel generative framework modeling distributions over hicrarchies. We extend Bayesian Flow Networks (BFNs) to enable transitions between probabilistic and discrete hierarchies parametrized via categorical distributions. Our proposed scheduler provides smooth and consistent entropy decay across varying numbers of categories. We empirically evaluate TreeGen on the jet-clustering task in high-energy physics, demonstrating that it consistently generates valid trees that adhere to physical constraints and closely align with ground-truth log-likelihoods. Finally, by comparing TreeGen's 's samples the exact posterior distribution and performing likelihood maximization via rejection sampling, we demonstrate that TreeGen outperforms various baselines.
1 Introduction
Hierarchies are central representing complex relationships across diverse domains. Consequently, hierarchical clustering algorithms are core components in well-established machine libraries such as scikit-leam [36], and find applications in a wide range of fields, spanning from phylogenetics and particle physics to web and citation network analysis. In phylogenetics, for example, clustering algorithms group organisms or genetic sequences based on similarity, helping to infer evolutionary relationships [18]. Further, real-world systems, from neural representation [22] to citation networks and web graphs [37] often exhibit hierarchical organization. Importantly, in high-energy physics, agglomerative linkage algorithms are indispensable for jet clustering and provide insights into the substructure of particle collisions [7].
Traditional hierarchical clustering algorithms operate primarily in an unsupervised setting. Agglomer- ative methods greedily merge clusters with the shortest distance until only a single root remains [21]. Divisive algorithms work in a top-down approach and start with all points in single cluster and iteratively split them into smaller groups [43, 14]. More recent methods optimize global cost functions via continuous relaxations, yet the setting remains unsupervised [11, 34, 9, 48, 30]. This is sensible many applications, as ground truth hierarchies are often unavailable for supervised training.
Jet clustering in high-energy physics is prime example where supervised simulations are available. yet traditional algorithms remain unsupervised. In particle accelerators, high-energy collisions produce unstable particles that successively decay and split into more particles. This resulting spray of particles, also known as a jet, detected by particle detectors. The task in jet clustering to reconstruct the latent hierarchy that describes the splitting process from the observed constituents, i.e., leaves of the hierarchy [33]. Simulators are extensively used in this field to generate collision events based on a physical model, enabling the creation of realistic hierarchy datasets for jet clustering [3, 4,
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

12]. Yet, despite this availability, prevailing jet-clustering pipelines often still rely on unsupervised algorithms [15, 46, 7, 17].
To address this problem, we introduce a novel generative model for hierarchies. Building on Bayesian Flow Networks (BFNs) [19] and Bayesian Sample Inference (BST) [31], our model explicitly models the posterior distribution over hierarchies.
Our key contributions are summarized as follows:
Generative Model for Hierarchies. We introduce TreeGen, a novel generative model tailored tree-structured data. The model transitions between probabilistic and discrete hierarchies parametrized via categorical distributions.
Entropy Scheduler for BFNs. We propose an alternative entropy scheduler for BFNs that provides smooth transitions of categorical distributions from uniform discrete states across varying numbers of categories.
Application to Jet Clustering. We demonstrate the practical effectiveness of our generative model by applying it to the jet clustering task using data from the GINKGO simulator [12].
2 Background
2.1 Probabilistic Hierarchies
Let 7 be rooted hierarchy, also called tree, with n leaves V = (V1,..., Va and n' internal nodes 7 (71: (71:-.., 7nr), J. where 7 is the root. We repre- sent the hierarchy via two binary adjacency matrices:
2/2
2
2z
23
p(2/22,)
Za
A E {0, 13mxn', BE 30,0, 10'''''
(1)
where the entry Aij encodes an edge from leaf va to internal node Zj, and Bij encodes an edge from internal node z; internal node Zj.
p(zlv,))
We enforce both matrices to be row-stochastic and additionally constrain B be upper triangular to the binary
V1
V2
V33
ensure they define valid tree-structure. case, this means that each row contains exactly non-zero entry, while the last row of B is all zeros
Figure Example of probabilistic hierarchy, The matrices represent A and B, respectively,
because the root has no parent. These constraints guarantee that (Å, B) encodes valid rooted tree.
Following Zügner al. [48], we relax these matrices to probabillistic assignments:
A E [0, 1jnx"',
B 10. 1 xn'
(2)
while retaining the row-stochastic and upper-triangul constraints. We interpret each row of A and B as a categorical distribution over parent assignments:
Aij p( v:). By = p(;; 2z),
denoting the probability that internal node z, is the parent of leaf v; and the probability that z, is the parent of internal node z;, respectively. Together, the matrices describe probabilistic hierarchy T= (A, B). By interpreting each row of a probabilistic hierarchy as categorical distribution, we can sample valid discrete hierarchies. An example of a probabilistic hierarchy is shown in Fig. 1.
2.2 Bayesian Flow Networks
Bayesian Flow Networks (BFNs) form a class of generative models that refine a belief over a target sample via successive Bayesian updates [19]. Recently, Lienen et al. [31] proposed Bayesian Sample
Project page: es. wtw.tumm.dd/dmm/tteeeenn

---

Inference (BSI), including BFNs in a more general framework, which we adopt here for notational simplicity. BSI starts with an initial belief over RD using Gaussian prior with precision X,
( \N,, E VER.. X>
which is updated in a Bayesian manner: y(æ y) ply x)p(2)) At each step, we observe a noisy measurement
y Na a
0,
where a denotes the precision of the measurement. We compute the closed- form Gaussian posterior
p y) -
Autoo, (A (A+a)= A+a
which converges towards the true a as the precision increases by a with each update. As we do not have direct access the sample noisy measurements y of æ train neural network & fo A) to approximate æ. The model is trained similarly to diffusion models [23] and minimizes the squared error between è and the ground truth However, unlike diffusion models, BFNs and BSI operate distributions rather than single samples. While BSI has considered the Gaussian case, it is not limited it, and we will consider categorical distributions in Sec. 3, done by BFNs.
3 TreeGen: A Generative Model for Hierarchies
t 1/2
Figure 2: Overview of the TreeGen generation process. At each continuous time E [o, 1]. the model maintains probabilistic hierarchy Tr, applies a Bayesian-inspired update via a neural network and noisy sample, and progressively reduces entropy until converging discrete hierarchy T 1.
In this section, present main contribution: TreeGen, Bayesian generative model designed for hierarchies. TreeGen starts from a probabilistic hierarchy T. which is iteratively updated via Bayesian updates of observed noisy hierarchies until the entropy of the hierarchy is sufficiently low and converges toward a discrete hierarchy. To obtain noisy observations, we design neural network trained predict ground-truth discrete hierarchies. As such hierarchy simply consists of collection of categorical variables (see Sec. 2.1), we first describe the framework for arbitrary categorical distributions. While we derive update equations within the BSI framework [31]. the resulting Bayesian updates match those of BFNs [19] but follow from a simpler derivation. Compared with standard BFNs, TreeGen instantiates a categorical BSI that uses the BFN update with explicit precision control while introducing different loss, omits auxiliary distributions, and extends the parametrization trees. We discuss these differences and provide detailed derivations in App. A.
Objective and prior. Our goal is to infer an unknown sample x E {l,..., K}. We encode our current belief about that sample via categorical prior, i.c..
p(x= = k) = Thy π = (#I....,Tk) E EA--1
where AK-1 is the K K-simplex, i.e., Te 0 and Σx T 1.
Noisy observation. Assume we have access to a noisy measurement vector y (y1,... ... vk) sampled from multinomial distribution with m Σn U/ trials and class probabilities 0, which we model as a mixture between the ground-truth Dirac distribution e, and uniform noise:
f + (6-ww w) 1. 05w51
(8)
3

---

Cr is the one-hot encoding of x and w defines noise trade-off: 1 yields perfectly informative counts, while w gives pure uniform noise.
The conditional likelihood for x: = k
ply k. w)
m!! 11, Øj.
m gk
1-
+
Bayes update.
After observing y, can update our belief on i.e., parameters π via Bayes' rule:
ply k, mT p(x y.w) Σ, ply x = j.w) m By canceling the common multinomial coefficients, we obtain softmax-ty taype expression:
(10)
2h n
kK
p( = k y, w) Σ, v Us TT
n =
(11)
1-
We can express this equivalently in vectorized form:
post softmax (log π + y logm')
(12)
This form shows that the observation contributes an additive vector y log n to our belief before the softmax normalization. This form also allows simple aggregation of sequential updates.
Sequential updates. rior update simplifies to:
With multiple independent observations {y J and constant w, the poste-
(V) softmax π (0) + log Σ C
(13)
Thus, the class probabilities are aggregated by addition in log- spg-space.
Gaussian interpretation. For large observation count m, the Central Limit Theorem implies that the rescaled sum of counts is approximately Gaussian. This allows to rewrite the Bayesian update using a Gaussian-distribute. variable and in continuous time:
T softmax (log .;-1 2):
24 N( ( (or at-1)
(14)
Due to the additive aggregation in log-space (see Eq. (13)), we can summarize multiple Bayesian updates into single one, i.e., the step from To to NT;
softmax (z:).
Ze ~ N(aaKaaa
(15)
where o₁ represents the accumulated signal. Note that we omitted log To as it is constant and the softmax operation is shift invariant. Consequently, the log-evidence in the softmax follows Gaussian and is parametrized with og. The recovered updates equal those of Graves et al. [19].
Distribution of πi: Following the categorical BFN formulation, represent the belief vector time as the softmax of latent Gaussian logits, n. = softmax(z:) conditioned on the ground-truth class This results in the distribution:
= f [ softmax (z:)) p(zz x, arg (1-1) dz
p( TT-
(16)
Furthermore, we can model the marginal of n,, directly without intermediate steps:
p( x) - f softmax(())) p(z, ax) dzz-
(17)
Both distributions are trivial to sample by drawing from the corresponding Gaussian and applying the softmax. During generation, the true class æ is unknown, so we replace it with proxy So(T,, (T,, !) predicted by our neural network and approximate the one-step transition by pa(m; TT-1) p(T, fo (t,t)) ()), which recovers the distribution of л, when fo( (Te, 1) successfully predicts Unlike the BEN update, this allows to control the precision directly.

---

Training of the neural network fe t). Our goal is to train the network to recover the original sample z from categorical distribution и,, Therefore, at each training step we draw ground sample z from the dataset D and timestep (~ 1), then sample the corresponding distribution according to Eq.(1 The network receives л, and as inputs, and outputs a categorical distribution, aiming to concentrate all probability mass on the true class We optimize the parameters A by minimizing the expected cross-entropy loss:
(()) Bg-
CELoss (fa t), x)]
(18)
where CELoss denotes the cross-entropy loss between the predicted categories and the ground-truth label. The nested expectation is approximated with Monte-Carlly samples of t, and π.; in each minibatch, and gradients are propagated through the network to update
Entropy scheduler. During generation, we want the entropy of T. to decrease smoothly from its maximum (uniform) value to zero (one-hot). If we draw
K =
K ===
K 16
K = 32
Quadratic
Logarithmic
R//4 Peddor mt)
N
0.5
TT. = softmax(Z)).
then the entropy of is monotone func- tion of a₁ While BFNs [19] chose ag = Ct the constant C' requires retuning when the number of classes K changes. We in- stead propose
0.0
0.2
0.4
1.0
0.0
0.2
0.4
0.6
1.0
Figure 3: Comparison of entropy schedulers. The plots show the expected normalized entropy of the categor- ical belief IT over time t, for varying class counts K. Left: the quadratic schedule used in standard BFNs. Right: our proposed log-based schedule, which yields an approximately linear decay across different K.
a2 = C+ Caalogg2 K
In(1-t), t), (19)
K
for E [0, 1]. which yields an approxi- mately linear decay of expected normalized entropy (see Fig. 3) while ensuring To is uniform and π e.. Because our experiments span hierarchies with varying numbers of classes (see Sec. 2)1), this scheduler enables us to use the same hyperparameters across all sizes of hierarchies.
Adaptation to bierarchies. Extending the categorical generation from a single random variable an entire hierarchy means treating every row of (A, B) its own categorical distribution, while updating all rows jointly so that the resulting sample is a valid hierarchy. Concretely, we model the n rows of and the n' rows of the apper-triangul class represents a parent choice for a node. Note that K, i.e., the number of possible parents, varies with hierarchy size and across rows of B due the apper-trianuulat structure. This parametrization induces distribution over rooted trees (hierarchies) and substantially reduces the search space compared arbitrary graph generation. Accordingly, So takes the current probabilistic hierarchy T, (A,, B,) and outputs a discrete proposal T. This prediction is then used to update the current belief, i.e., the probabilistic hierarchy, using Eq. (15). We show algorithms for sampling and training in Algs. and 2, respectively.
matrix B as categorical distributions, where each
Algorithm 1 Sampling with TreeGen
Algorithm 2 Training TreeGen
input: Neural network fe. entropy schedule ag for t E [0, 1]. sampling steps N
input: Dataset D, neural network fo. entropy schedule a for t [0, 1]. training steps N for n 1 to N do
output: Discrete hierarchy Ti
1: To =1/K
2:
T ~ D. ~ u(0,1) 1)
2: for in (I/N,..., 1) do
3:
Z: ~ N(a, K T, a,K1)
> Eq. (15)
3:
t + fo fT 1/N)
Evaluate fo Eq. (15)
4:
T + softmax((z)
& Eq. (15)
4: 2, Na, To,KKT,kKT KT)
5:
T < fe(Te, t)
b Evaluate fa
5: Te + --ftmmaxxxz
b Eq. (15)
6:
L. + CE(T, T
Compute loss
6: end for
7:
0 0= VoL
Gradient step
7: return T
8:
end for
5

---

Experiments
In this section, we present our empirical results on the jet-clustering task. Our primary goal is to demonstrate the effectiveness of TreeGen in generating valid hierarchies adhering to the true data distribution. We compare TreeGen against various baselines and test how well the generated hierarchies approximate the ground-truth posterior distributions.
We evaluate TreeGen on hierarchies derived from the GINKGO dataset, focusing on QCD and W jets. More specifically, our experiments focus on the conditional generation p(T X), where X are leaf features. Finally, we conduct ablation studies to analyze the impact of key design choices in our method in App. C.
4.1 Experimental Setup
Datasets. Our evaluation uses five datasets obtained with the GINKGO jet shower generator [12]: QCD jets: QCD-S, QCD-M, QCD-L, and W-Boson jets: W-S, W-M. All five share the same root four-momentum Pioot and decay rate X, but vary in the shower cut-off Acnt. controlling the final jet size. The datasets have Acut E {4.0², 1.12, 0.6²] and contain hierarchies with most 19 (small), 59 (medium), and 99 (large) nodes. Each dataset consists of 100,000 hierarchies. Split into 98,000 hierarchies for training and 1,000 for validation and testing. In our setup, the leaves of the hierarchies correspond to the observed particles with features X, while internal nodes are latent and to be inferred via the sampled hierarchy T. Fig. 8 shows the leaf-count distribution for the three QCD datasets. We provide more details about the datasets and task in App. B.2.
Baselines. We benchmark the generative capabilities of TreeGen against various baselines. We include CatFlow [16] and (standard) Bayesian Flow Networks (BFNs) [19] representing sfate-of-the- art generative models for categorical data. Furthermore, we compare to greedy clustering approaches and the ground-truth posterior obtained via cluster trellis [33]. Finally, we compare to the CA [15, 46), kr [17], and anti-kr [7] algorithms in App. C, representing traditional jet clustering algorithms. All generative baselines share the tree assumption (see Sec. 2.1), while the agglomerative algorithms are tailored to jet clustering and restricted to binary hierarchics. We describe the baselines in more detail in App. B.5.
Evaluation metrics. To assess the performance of the generative model, we employ two metrics. First, we compute the valid hierarchies fraction of generated hierarchies that satisfy dataset- specific constraints, i.e.. required physical properties. Second, we assess how well the sampled hierarchies match the data distribution by comparing their liglliiilsoosss to those of the ground truth hierarchies. Specifically, we compute the ratio of log-likelihoods of each generated hierarchy to its corresponding ground-truth counterpart and report the average ratio over the test set. A ratio close to one indicates that the model closely approximates the true data distribution. We discuss further details about both metrics in App. B.6.
Practical considerations. TreeGen requires architecture that predicts hierarchy given probabilistic hierarchy as input. This structure makes suitable for Graph Transformer [38]. We fix n' depict the high-level architecture in Fig. and provide more details in App. B.3 and discuss node and edge features in App. B.4. We train all models using Adam [28] with learning rate of 0.0001 and gradient clipping set to 0.5 for 50 epochs. For generation, all models use 1000 steps, except for the ablation in Fig. where all use 100. We provide an overview of all hyperparameters in App. B. Finally, we report the mean and standard deviation of four random seeds for all experiments to ensure reproducibility.
i.e., one parent for every non-root node, we operate on binary hierarchies. We
4.2 Results
Hierarchy generation. We report the fraction of valid trees and log- likeeiooosss in Tab. 1. TreeGen generates nearly perfect trees every case, achieving > 0.93 validity and likelihood ratios close to one, consistently outperforming baselines across all datasets. By contrast, CatFlow and BFN deteriorate rapidly as jet size grows: on the medium QCD set, their valid-tree fractions drop to 0.24 and 0.65, respectively. We attribute CatFlow's inferior performance to its flow leaving the

---

Feature Generation
Graph Transformer
Prediction Head
Momeuilum
Conservation
Angular
Features
Figure 4: TreeGen model architecture. Node and edge features (structural and physics-inspired) are derived from the current probabilistic hierarchy, then processed by Graph Transformer to predict the parameters of the final hierarchy Ti.
probability simplex: intermediate states can violate the categorical constraints that define valid jet hierarchy. The Bayesian- based approaches BFN and TreeGen remain the simplex throughout generation by design. These results show that TreeGen's successfully scales to larger, more complex jets, consistently maintaining both the highest validity and the best log-likelihood fractions.
Table I:
Evaluation metrics for different models datasets. Best scores in bold.
CatFlow
BFN
TreeGen
Dataset
Valid Frac. (t)
LLH Frac.
Valid Frac. (1)
LLH Frac.
Valid Frac. (1)
LLH Frac.
QCD-S
0.752=0.014
0.972+#.000
0.941 000
0.982:00 000
0.997:0.00 #0.000
1.003
QCD-M
0.236=0.016
0.882#0.000
0.645=0.040
0.936 002
0.977 40.0000
0.994:0.002
QCD-L
0.416=0.036
0.898 002
0.943+0.0016
0.000
W-S
0.604:0.014
0.93740.007
0.8014
0.965 001
0.99441..00 002
1.006
W-M
20.014
0.886=0.00m 002
0..30040.244
0.980
Furthermore, Fig. 5 visualizes the per-tree ligllikei:oodd:: each dot compares the log-likelihood of the ground-truth hierarchy (x-axis) with the log-likelihood of the hierarchy generated by our model (y-axis). TreeGen's points cluster along the diagonal, indicating alignment with the true distributions. This confirms that TreeGen not only produces more valid hierarchies but faithfully reproduces the likelihoods of the ground-truth.
QCD-S
QCD-M
QCD-L
19
55
91
Number of Nodes
log 1og (ThoeGenn
1
17
-47
777
15
-63
16
13
31
11
335
1.5
21
70
60
50
40
30
150
100
200
150
100
50
log p( (Trmn)))
log p (TTrith)
log p( TTruts)
Figure 5: Scatter plot log-likelihoods. Each point corresponds to hierarchy: the z-axis shows the log-likelihood of true hierarchy, and the y-axis shows the log-likelihon hierarchy generated by TreeGen, demonstrating close alignment across varying hierarchy sizes.

---

Posterior distribution. While our reported quan- titative metrics demonstrate that TreeGen produces reasonable hierarchies, they do not directly assess how well it approximates the true posterior p(T X). To complement our evaluation, we employ the cluster trellis data structure [33] to draw exact sam- ples from the posterior. This enables us to com- pare these samples with those generated by TreeGen. Since the computational complexity of cluster trel- lis scales exponentially as ((3 ) this evaluation is only feasible for the QCD-S dataset, which contains hierarchies of up to 20 nodes. For each of the 1000 test hierarchies, we draw sample from the true posterior via cluster trellis and one from TreeGen.
True Posterior
TreeGen (Ours)
-30-
-00
log (())
-50
-60-
-70
Figure 6: Histogram of posterior log- likelihoods. Distributions of log X) for true hierarchies versus TreeGen-sampled hierarchies, demonstrating that TreeGen suc- cessfully approximates the full posterior.
We visualize and compare the empirical distributions over the log-likelihoods using histograms in Fig. 6. The empirical distribution produced by our genera- tive model successfully captures the key characteris- ties of the true posterior distribution. This further validates TreeGen's 's ability to accurately learn and represent hierarchical structures.
Likelihood Maximization. While sampling from TreeGen provides diverse plausible reconstructions, many applications require single high-likelihood hierarchy. Our model's non-deterministic nature resample hierarchies multiple times. To this end, we perform approximate MAP infer- ence via rejection sampling. More specifically, for each jet event X. we draw N independent hier- archies {T.) PT X)., X). compute their log- likelihoods log p(TT), and select the sample with the highest value. We benchmark this against a deter- ministic greedy agglomerative linkage algorithm that merges clusters by maximizing local likelihood gains. In Fig. 7, we show how the log- sgdlikelihood grows with the number of samples N compared to the generative baselines and the greedy algorithm.
TreeCien (Ours)
CatHlow
BFN
Greedy
46
allows
logp(T)
-17-
-48
10
20
30
40
50
Number of Samples
Figure 7: Likelihood maximization. Maxi- mum log-likelihood obtained by drawing N candidate hierarchies from TreeGen and the log-likelhhood of the deterministic greedy re- construction. TreeGen outperforms the base- line within few samples.
We observe that the log- ligelikel obtained by Tree- Gen exceeds that of the greedy agglomerative base- line after only 16 samples and consistently outper- forms CatFlow and BFN, demonstrating that our samples not only concentrate around siobbb-probbbility hierarchies but also capture diversity beyond a single greedy pass. Note that both this rejection-sangee sampling MAP procedure and the greedy likelihood clustering assume access to the exact likelihood, which often unknown or intractable for realistic collision data. Nevertheless, this experiment highlights the practical benefit of our model's ability to sample multiple hierarchies, enabling improved reconstruction quality whenever quantitative evaluations available.
Scheduler comparison. We compare our scheduler, pro- posed in Sec. 3, with the BFN scheduler [19] across differ-
Table 2: Scheduler ablation on validity and LLH fraction.
ent C. Following Graves et al. [19], we include C' = 0.75
Scheduler
Valid Frac.
LLH Frac.
and C 3.0. Additionally, we test C 6.0 and C 9.0.
The results on QCD-S are shown in Tab. 2.
Ours
0.997 0.001
1.003 + 0.002
0.7.5
0.410 0.019
0.990 0.002
As we observe, the scheduler of Graves et al. [19] is highly sensitive to C. While C = 0.75 highly degrades validity,
3.0.0
0.989 0.001
1.002 1. 0.002
C=
0.994 + 0.001
1.003 4: 0.002
C
0.992 + 0.004
1.002 + 0.001
larger values result in matching performance to ours.
8

---

5 Related Work
5.1 Hierarchical Clustering
Agglomerative algorithms iteratively merge clusters based on the shortest distance, building a hierarchy from individual points until converged to single cluster. Different choices define the inter-cluster distance give rise to various linkage algorithms [21]. Divisive algorithms proceed in reverse and build hierarchies by splitting single cluster, containing all points, into smaller and smaller clusters, c.g. by repeatedly applying the k-means algorithm [43] or spectral clustering [14].
Recently proposed cost functions, such the Dasgupta cost [14] and the Tree Sampling Divergence (TSD) [10], have enabled global optimization across the entire hierarchy. While these cost functions remain infeasible to optimize directly, various continuous relaxations have been proposed to allow for gradient- gradient-based optimization [11, 34, 9, 48, 30J. These methods then optimize relaxed objectives.
Similar to hierarchical clustering algorithms, TreeGen aims to infer hierarchies and shares the same parametrization for probabilistic hierarchies as Zügner et al. [48] and Kollovieh al. [30]. However, we focus on supervised setting and learn distribution of hierarchies, whereas the former focuses on inferring hierarchies based on heuristics or objective functions.
5.2 Jet Clustering
Jet clustering is a well-studied task in high-energy physics. Initially, cone algorithms [5], which select the most energetic particles and group all particles within cone around them, were employed for this task. However, these algorithms were sensitive to low-energy particles [29]. This limitation led to the development of sequential recombination clustering algorithms, which are instances of agglomerative hierarchical clustering. Different algorithms in this class, such as the Cambridge/A Aachhdg algorithm [15, 46], the ky algorithm [17], and the anti-ky algorithm [7], differ in how they combine angular proximity and energy of particles into a distance measure.
Recently, the development of the GINKGO simulator [12] introduced new class of algorithms that use this simulator to assess the likelihood of splits in the hierarchy, aiming to find the hierarchy with the maximum likelihood [6, 20, 33, 13). Among those, cluster trellis [33] is most similar to method as it enables sampling from the exact posterior distribution for small hierarchies. However, our model learns directly from simulator- -enttrreeeeeeeeedd data without requiring likelihood evaluations inference time, which allows it to be applied with any black-box simulator. Finally, Yang et al. [47] propose a variational-infereni method for jet clustering on the GINKGO dataset.
5.3 Generative Models
Our work builds on recent advances in generative modeling, particularly diffusion and flow-matching frameworks. Diffusion models add random noise data and then learn to reverse this diffusion process to generate samples starting from noise [39, 23, 40, 41]. Flow matching generalizes this idea by directly learning an ODE that transforms data to noise via regression of vector fields [32]. Furthermore, Graves et al. [19] introduced BFNs, another related generative framework that maintains and updates a probability distribution during generation in Bayesian manner rather than acting directly on the sample. Atkinson et al. [1] proposed an ODE-based BFN sampling algorithm that replaces aggregated previous predictions with the most recent prediction, also providing explicit precision control, similar to ours.
Generative models for discrete data are particularly relevant to our approach since hierarchies are discrete structures. One class of such models maintains discreteness throughout the generative process. Various approaches have adapted diffusion models to discrete data [24, 2, 8]. Recent flow matching approaches confine the generative process in continuous space. For example, Dirichlet Flow Matching [42] constrains the process on the probability simplex, while CatFlow [16] removes this constraint, allowing values outside the simplex. Our approach belongs to the former category, as we build upon the BFN framework, transitioning between categorical distributions.
Prior work, such as the junction tree autoencoder [25], modeled trees in an autoregressive fashion for molecule generation. Other approaches used diffusion and flow matching for general graph generation [35, 26, 45, 16). Unlike these, however, our framework explicitly parametrizes valid hierarchies (rather than distributions over all graphs), substantially reducing the search space.

---

Conclusion
In this work, we introduced TreeGen, novel generative model learning hierarchies from distributions building upon the BFN [19] and BSI [31] framework. By proposing an entropy scheduler, we are able to model categorical variables smoothly across a varying number of classes. The intermediate states of the generation process are probabilistic hierarchies, which we used to derive meaningful features, improving generative performance.
We assessed TreeGen on high-energy jet simulator benchmark and found that close to 100% of the generated hierarchies adhere their corresponding physical properties. Moreover, by comparing the likelihoods of conditionally generated hierarchies to those of the ground truth, we demonstrated that our model able to successfully approximate the posterior distribution.
Limitations and future work. A core component of our generative model is graph transformer whose computational complexity increases quadratically with the number of nodes. This issue arises from the densely connected probabilistic hierarchies during generation.
To this end, have evaluated TreeGen on GINKGO [12] putations. Immediate extensions include evaluating TreeGen on different simulators [3, 4] and experimental, i.e., real-world jet datasets, where measurement noise and detector effects introduce additional complexity. TreeGen naturally extends to non-binary trees by adjusting the number of internal nodes n' (for binary trees n 1). Potential strategies include sampling n' (as in molecule generation) [45]; predicting n' using a classifier similar to Kerrigan et al. [27]; or setting a larger n' and pruning unused internal nodes, as done in hierarchical clustering [48]. Finally, one could explore alternative hierarchy parameterizations enforcing strict binary-tree constraints or improving both computational efficiency and likelihoods.
it provides analytical likelihood com-
Contributions
MK developed the project idea, derived the theory, implemented the core method, and wrote the manuscript. NF implemented most of the experimental pipeline (data loading, features, backbone, baselines, training, and evaluation) and assisted with the manuscript. FG supervised the neural network architecture design. BS contributed to conceiving the initial idea. SG contributed to the method design and provided overall scientific guidance. All authors discussed results, both theoretical and empirical, and revised the manuscript.
Acknowledgments
We thank Johanna Sommer for valuable feedback the idea and Marten Lienen for help with the background section.
10

---

References
[1] Timothy Atkinson, Thomas D Barrett, Scott Cameron, Bora Guloglu, Matthew Greenig, Char- lie B Tan, Louis Robinson, Alex Graves, Liviu Copoiu, and Alexandre Laterre. Protein sequence modelling with bayesian flow networks. Nature Communications, 16(1):3197, 2025.
[2] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34: 17991-1799999 2021.
[3] Johannes Bellm, Stefan Gieseke, David Grellscheid, Simon Plätzer, Michael Rauch, Christian Reuschle, Peter Richardson, Peter Schichtel, Michael H Seymour, Andrzej Siódmok, al. Herwig 7.0/herwig++ 3.0 release note. The European Physical Journal C, 76:1-8, 2016.
[4] Christian Bierlich, Smita Chakraborty, Nishita Desai, Leif Gellersen, Ilkka Helenius, Philip Ilten, Leif Lönnblad, Stephen Mrenna, Stefan Prestel, Christian Tobias Preuss, al. A comprehensive guide to the physics and usage of pythia 8.3. SciPost Physics Codebases, page 008, 2022.
[5] Gerald Blazey, Jay R Dittmann, Stephen Ellis, V Daniel Elvira, Frame, S Grinstein, Robert Hirosky, R Piegaia, H Schellman, R Snihur, al. Run ii jet physics: proceedings of the run ii qed and weak boson physics workshop. arXiv preprint hep-ex/0005012, 2000.
[6] Johann Brehmer, Sebastian Macaluso, Duccio Pappadopulo, and Kyle Cranmer. Hierarchical clustering in particle physics through reinforcement learning. arXiv preprint arXiv:201 1.08911, 2020.
[7] Matteo Cacciari, Gavin P Salam, and Gregory Soyez. The anti-kt jet clustering algorithm. Journal of High Energy Physics, 2008(04) 2008.
[8] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 14-28266-28279, 2022.
[9] Ines Chami, Albert Gu, Vaggos Chatziafratis, and Christopher Ré. From trees to continuous embeddings and back: Hyperbolic hierarchical clustering. Advances in Neural Information Processing Systems, 33:15065-15076 2020.
[10] Bertrand Charpentier and Thomas Bonald. Tree sampling divergence: an information-thhorreie metric for hierarchical graph clustering. In IJCAI-19, 2019.
[11] Giovanni Chierchia and Benjamin Perret. Ultrametric fitting by gradient descent. Advances in neural information processing systems, 32, 2019.
[12] Kyle Cranmer, Sebastian Macaluso, and Duccio Pappadopulo. Toy generative model for jets. URL https: hete:://gttttthtt htt..///aoaasttaaaaaaalsssoo httpr///aaaaaaaaaogogoooa/////// /gi.g////blo//mmtt///oootooo coyshower_v4.pdf.
[13] Kyle Cranmer, Matthew Drnevich, Sebastian Macaluso, and Duccio Pappadopulo. Reframing jet physics with new computational methods. In EPJ Web of Conferences, volume 251, page 03059. EDP Sciences, 2021.
[14] Sanjoy Dasgupta. cost function for similarity-based hierarchical clustering. In Proceedings of the firtt-eighth anmual ACM symposium Theory of Computing, pages 118-127, 2016.
[15] Yu L. Dokshitzer, GD Leder, S Moretti, and BR Webber. Better jet clustering algorithms. Journal of High Energy Physics, 1997(08):001, 1997.
[16] Floor Eijkelboom, Grigory Bartosh, Christian Andersson Naesseth, Max Welling, and Jan- Willem van de Meent. Variational flow matching for graph generation. Advances in Neural Information Processing Systems, 17:11735-11744, 2024.
[17] Stephen Ellis and Davison Soper. Successive combination jet algorithm for hadron collisions. Physical Review D, 4867):3160, 1993.
[18] J. Felsenstein. Inferring phylogenies. Sinauer Associates, 2003.

---

[19] Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow networks, arXiv preprint arXiv:2308.07037, 2023.
[20] 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. In Uncertainty in Artificial Intelligence, pages 2061-2071. PMIR, 2021.
[21] Trevor Hastie, Robert Tibshirani, Jerome H Friedman, and Jerome H Friedman. The elements of statistical learning: data mining, inference, and prediction, volume Springer, 2009.
[22] Geoffrey Hinton. How to represent part-whole hierarchies in neural network. Neural Computation, 35(3):41-442, 2023.
[23] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
[24] Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in neural information processing systems, 34: 14:1244--1246 2021.
[25] Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In International conference machine learning, pages 2323-2332. PMLR, 2018.
[26] Jachyeong Jo, Scul Lee, and Sung Ju Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. In International conference on machine learning, pages 10362-10383, PMLR, 2022.
[27] Gavin Kerrigan, Kai Nelson, and Padhraic Smyth. Eventflow: Forecasting temporal point processes with flow matching. arXiv preprint arXiv:2410.07430, 2024.
[28] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14Xiv:1412.6980, 2014.
[29] Roman Kogler, Benjamin Nachman, Alexander Schmidt, Lily Asquith, Emma Winkels, Mario Campanelli, Chris Delitzsch, Philip Harris, Andreas Hinzmann, Deepak Kar, et al. Jet substruc- ture at the large hadron collider. Reviews of Modern Physics, 91(4):045003, 2019.
[30] Marcel Kollovich, Bertrand Charpentier, Daniel Zügner, and Stephan Günnemann. Expected probabilistic hierarchies, Advances in Neural Information Processing Systems, 37:13818-13850,
2024.
[31] Marten Lienen, Marcel Kollovieh, and Stephan Günnemann. Generative modeling with bayesian sample inference. arXiv preprint urXiv:2502.07580, 2025.
[32] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
[33] Sebastian Macaluso, Craig Greenberg, Nicholas Monath, Ji Ah Lee, Patrick Flaherty, Kyle Cranmer, Andrew McGregor, and Andrew McCallum. Cluster trellis: Data structures & algorithms for exact inference in hierarchical clustering. In International Conference on Artificial Intelligence and Statisties, pages 2467-2475. PMLR, 2021.
[34] Nicholas Monath, Manzil Zaheer, Daniel Silva, Andrew McCallum, and Amr Ahmed. Gradient- based hierarchical clustering using continuous representations of trees in hyperbolic space. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 722, 2019.
[35] Chenhao NIu, Yang Song, Jiaming Song, Shengjia Zhao, Aditya Grover, and Stefano Ermon. Permutation invariant graph generation via score-based generative modeling. In International conference artificial intelligence and statistics, pages 4474-4484, PMLR, 2020.
12

---

[36] Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit- learn: Machine learning in python. Journal of machine learning research, 12(0ct):2825-2833) 2011.
[37] Erzsébet Ravasz and Albert Laszló Barabási. Hierarchical organization in complex networks. Physical review. E, Statistical, nonlinear, and soft matter physics, 67 2 Pt 2:0261 12, 2002.
[38] Yunsheng Shi, Zhengjie Huang, Shikun Feng, Hui Zhong, Wenjin Wang, and Yu Sun. Masked label prediction: Unified message passing model for semi-supervised classification. arXiv preprint arXiv:2009.03509 2020.
[39] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages
pmlr, 2015.
[40] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[41] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 13456, 2020.
[42] Hannes Stark, Bowen Jing. Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to dna sequence design. In International Conference on Machine Learning, pages 46495-46513. PMLR, 2024.
[43] Michael Steinbach, George Karypis, and Vipin Kumar. comparison of document clustering techniques. 2000.
[44] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[45] Clément Vignac, Igor Krawezuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation. In ICLR, 2023.
[46] Markus Wobisch and Thorsten Wengler. Hadronization corrections to jet cross sections in deep-inelastic scattering. arXiv preprint hep-ph/99072280, 1999.
[47] Hamming Yang, Antonio Khalil Moretti, Sebastian Macaluso, Philippe Chlenski, Christian A. Naesseth, and Itsik Pe'er. Variational pseudo marginal methods for jet reconstruction in particle physics. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
[48] Daniel Zügner, Bertrand Charpentier, Morgane Ayle, Sascha Geringer, and Stephan Günnemann. End-to-ent learning of probabilistic hierarchies on graphs. In International Conference Learning Representations, 2021.
13

---

A Theory
In the following, we will provide derivations for equations Eqs. (13) to (15). Both theorems build on the theory of Bayesian Flow Networks and summarize the derivations from Graves et al. [19].
Theorem A.1. Let x [I....,K] K have prior (0)) E AK-1- For each update i i I....,............. N we observe y( (vi)) .....UK! (0) UK E GNK, with total count m = Σx Uk sampled
(a)
m!!
VE
ply x = k,w)
1-
w
w E (0, 1).
(20)
1I, Mj
Then the posterior after N sequential Bayesian updates is
π (N) softmax (log TO y log (ogm')
(21)
where x/:=====11 1 ok I-
and y the sum of N samples, i.e., y Ei= y(0)
Proof. We prove by induction on N.
Base case (N = 1). By Bayes' rule for single observation y1),
T (1) = plx- = k =11 (())
(22)
p(y(1) =k=kw (0)
(23)
Ex=1p(ye p(y(1) x== W (0)
m!
1-w m-- U/
w/ uk (0)
11 My!
Tk
(24)
ml!
1-w
L- + w
1-w
wK (1
1-
(25)
1- m Σx= k==1
(1 +
WK
24 (0)
1-
T
7:' T/ T
(1)
(0)
(26)
Σ क
(0)
k
softmax (log π (0) + y
(1) log (ogm))
(27)
Thus, (1)
softmax T (0) + (2) log r') matching the formula as EL= y(0) = (1))
Inductive step.
Assume the statement holds for some N 2 1:
T (N)
sofumax (log ( (0) + ENN y(t) log 1og
After receiving y(+1) we update:
(zr))ak (N+) (N)
(N+1)
T
(28)
E, (oroo (N))
14

---

We insert the inductive hypothesis for (M); mb
(1 (3)) -2) (N)
(N+1]
m Σ (T))))) (N)
(29)
#
( ( exp(log T + log n Ei=1
(1))
Vi:
N Σ
Σ
exp exp e + log π
(1)
Vi
i=1
(30)
E(rjju exp(log
n (0) log x Zi=1 N
3/
exp(log Te
(0) +log T'
exp( k (0) + log x
N
(
(
ii=
2k
(31)
exp(log ;j (0)
+ +log n'
M/
(1)
+ 1/ ((+1I))
i=1
exp(log 1 (0) log Σ
N+1 (2)
i=1
y/
(32)
N+1
Σ exp(log ; + loog π
(0)
y3
Because Ee y/ 1/2
(6)
(N+1) Ei=l y/
N+1 (6) the numerator and denominator are exactly the softmax
with the sum taken to N +1. Hence
(N+1) softmax (0)
(0) + EN y(1) log n'
(33)
Thus, a(N+11 softmax (lg (0) + [EiNT y(0) logn") completing the induction.
Theorem A.2. Fix K 2 2 and w E (0, 1). Let x E fI,..., K} have prior To E AK-1- At each time step E N observation y(() E NK with Σx Ui m is drawn from
(1)
mn!
m-v(
p(y(0) x== = k:w) = I1
1-w K
+ w
ug K
(1)
WK 1-w and Dai :a::000 ar = mww
Define n² :=
Then the following posterior approximations hold for m + oc:
Single-step update. Given n; m. can be approximated using a Gaussian random vector:
T. x softmax (log T:-1 z1):
NA Aa,KI).
(34)
2. Aggregated update. ne can be approximated using a Gaussian random vector directly:
~ softmax((2)) ~ N( Kes, a, KT).
(15)
Proof. (i) Single step. Conditioned x, the vector y() is multinomial with parameters and probabilities
0 we. + (1 - w) 1
(35)
The Central Limit Theorem yields:
(())
N (mme,, + +(1 (1- w-) R m (inia())) -00)
(36)
m-
15

---

By applying mean-field approximation diag(O) 00 = I and noting that log x' R wk since - 0. we obtain:
log nt
Nmm Ke, + w01 w)1, mu² KT).
(37)
Since the softmax operation is shift-invariant, "N(m K Ke,, KI). Plugged into the update, we obtain softmax(log nT- 2,)
drop the constant w(1 w)1 and define
matching Eq. (14).
(ii) Aggregation. The vectors Z1. 2 are conditionally i.id. Gaussians; their sum is Gaussian with mean E === Aa, Ke Ke, and covariance Σ Ao; KT = KT. Because a constant vector can be subtracted inside the softmax, the factor log Ty drops out, yielding æ softmax(z)) with the stated distribution.
A.1 Differences to BFNs.
While there is a different theoretical perspective between TreeGen and BFNs, both models build on the same Bayesian update. The key differences mostly occur in the sampling and training procedure. TreeGen performs the following update during sampling:
A f 1)
(38)
(
(39)
softmax(zi).
(40)
In contrast, BFNs perform:
म [T 1)
(41)
~ N da: dt
day KI) dt
K
(42)
n. = softmax (TT + 2t).
(43)
Furthermore, BFNs optimize their framework using the following objective:
((0) (...
da.
( t) ex
(44)
2dt
while TreeGen the loss presented in Sec. 3.
16

---

B Experiment Details
B.1 Hyperparameters
We provide overview of the training and model hyperparameters in Tab. 3.
Table 3: Hyperparameters of TreeGen.
Training
Sampling
Graph- Transforrner
Upscaler
Prediction Head
Epochs
Grad. Clip
EMA
Steps
Layers
Dim
Heads
Node Layers
Edgo Layers
Layers
Value
All experiments are conducted on A100 GPUs.
B.2 Datasets
We simulate five jet-shower datasets using the GINKGO generator [12]. Three contain QCD jets (QCD-S, QCD-M, QCD-L), and two describe W-boson jets (W-S, W-M). All sets share common decay rate and differ only in the shower cut-off Acm E (4.0², 1.12, 0.6²). which controls the tree sizes: up to 19 nodes for S, 55 for M, and 99 for L. Each dataset contains 10 binary hierarchies, split into 98,000 training examples and 1,000 each for validation and test. Fig. 8 plots the leaf-count distributions of the three QCD samples, illustrating how the cut-off controls the final hierarchy size. Inside every hierarchy, the nodes, i.e., an energy manngy-mmmmnuum vector specifies particles:
p" (E,, Pay Py, P=):
(45)
which contains the energy E and momentum (Pez Py Pz). In our setting, we only observe the leaves and infer the hierarchy, i.c. the structure of internal nodes.
QCD-S
QCD-M
QCD-L
0.06
0.25
0.08
Density
0.20-
0.06
0.04
0.15
0.04 -
0.10-
0.02
0.05
0.02
0.00
0.00
0.00
10
10
20
30
40
Number of Leaves
Number of Leaves
Number of Leaves
Figure 8: Histogram of the number of leaves for the three different QCD datasets.
In our setup, the tree corresponds to a splitting process. A parent node splits into its children. The root corresponds to the initial particle, while the leaves represent the observed particles. We aim to reconstruct the splitting process, i.e., find the edges of the tree by inferring the parent-child edges given the leaves.
B.3 Architecture
The input to our neural network is a probabilistic hierarchy with n + n'' nodes, representing leaves and internal nodes, i.e., T [0, [(====')x(((+))))))" Our architecture consists of three components: (f) the feature generation (discussed in App. B.4). (ii) graph transformer, and (iil) prediction head.
Graph Transformer. Given h-dimensional node and edge features, denoted as X E R (n+n')x and E mRmxh respectively, we iteratively update node features (() at layer given (1-1), E, and the edges of the hierarchy transformer proposed by Shi et [38] as our backbone with four layers, two attention heads, and hidden dimension of 64 (Tab. 3).
adjacency matrix A E {0. (m
We use the graph
17

---

Prediction head. Our goal is to predict parent probabilities, i.e., edge probabilities from the node representations of the graph-transforme Therefore, we apply a prediction head each child -paeent pair (v,, 74) to obtain a logit for the corresponding edge and apply a softmax over all outgoing edges of the child.
The prediction head concatenates the latent representations of the child and the parent together with the corresponding edge features. These are then passed through a small MLP to obtain the logits as depicted in Fig. 9.
V1
Batch Norm
Batch Norm + Leaky ReLU
Linear Layer
Linear
2
Linear Layer
Layer
Leaky ReLU
(V1, Zj
Figure 9: Prediction head. The prediction head processes the representations child-parent pairs
(Ve, 2,), together with their corresponding edge features, and predicts single logit.
B.4
Feature Generation
To model probabilistic hierarchies effectively, we extract diverse set of features, We first introduce features encoding structural and uncertainty-related information such node type, entropy, and ancestor probabilities (App. B)..1). We then present chysically-inspieee features, which capture domain-spenn -peeiffc properties like momentum, rapidity, and invariant mass (App. B.4.2). Finally, we project raw features into the hidden space used by the network (App. B.4.3).
B.4.1 Features for Probabilistic Hierarchies
In this section, we present node and edge features that can be derived from probabilistic hierarchies and aid the network in obtaining global features. We use the equations and sampling procedure from Züigner et al. [48].
Node Type. If nodes have children, they categorized as internal nodes; otherwise, leaves. We represent this with a binary feature. Furthermore, we introduce a binary feature identifying the root, as it is the only node without parent.
model each node's parent assignment as a categorical distribution over its candidate parents and derive two features: (i) the number of candidate parents and (ii) normalized entropy quantifying uncertainty. Specifically, we compute:
Node Entropy.
H(v,) = log(n')
Aij log Aij.
(46)
j=1
n'
H(zz)) = log(n' (i+ 1))
Σ
Biy log Bij-
(47)
j=it1
The logarithmic normalizers match the number of candidate parents in each case, scaling the entropy to [0, 1].
Expected Children. For each internal node Zμ, we include the expected number of children under the tree sampling procedure (Á, B) PA,B feature. Under PA,B, this equals the total
18

---

probability that ž₁ is chosen as the parent of each leaf and each earlier internal node j < k:
E Tehild (2+)]
Σ p(zk
vo)
p(zz 22))
(48)
1=1
A
BjRn
(49)
1=1
where A, and B, denote the parent of leaf v, and internal node ,,, respectively. The equalities follow from the linearity of expectation and [[[[ k]] Pr(X k)).
Sibling Probability. We encode the sibling relation, i.e., whether two nodes share the same parent, by its probability under the rre-sampling procedure as edge features. For leaves v, and Vj, we compute the probability by:
p (v; and vy are siblings)
Σ
p(zz
Vol Zk vz))
(50)
Σ p(zz vo) p(zz vz)
(51)
Σ k=1
Auk-
(52)
where used the independence of parent draws in the sampling procedure. For internal nodes, compute this quantity using B, summing only over feasible parents.
Ancestor Features. We use two ancestor- ancestor-based features. (i) Ancestor probability as an edge feature for an internal node and a node vi, let Pane == p(zz anc(v:)) denote the probability- -under the tree-sampling procedure that 2ạ is an ancestor of v,. (ii) Expected number of ancestors node feature, which equals the node's s expected depth.
For leaf va, we compute the expected number of ancestors as:
Teane (v:))
Σ I|z anc(vi)]
(53)
p(zz E anc(v:))
(54)
k=
n'
panc
(55)
The computation for internal nodes follows analogously, using the matrix punc
For pair of nodes, we also compute the (expected) number of shared ancestors as a similarity measure. As cannot derive an exact expression due the dependence between the ancestor probabilities p(z+ E anc(v,)) and p(zz anc(v,)), we use an independence approximation:
n
E[ (Ve, v,)] 2
pane panc
(56)
ik
The computation of pane and June follows Zügner et al. [48].
Hierarchy size. the hierarchies can vary in size within the datasets, we include the hierarchy size. i.c., the number of nodes, as feature.
19

---

Time embeddings. In addition to the probabilistic hierarchies and their derived features, we include the time E [0. 1] as imput the network. We use 10-dimensional sinusoidal positional encoding as done in the transformer architecture [44]. We concatenate these embeddings to the node features.
B.4.2 Physically Inspired Features
In addition to the lask-agnostic features described above, we include physically inspired features tailored to jet clustering. Each node in the hierarchy corresponds to a particle and is described by its four-momentum vector
pd (E, Per Py, Ps)
(57)
which serves the basis for our physics-derived features. The momentum ø (px, Pys Ps) is defined in three spatial directions: the z-axis is aligned with the beam (beamline), and the plane spanned by the two axes orthogonal is the transverse plane.
Pi
Figure 10: Illustration of the transverse momentum pr and the azimuth & (left), and the angle between the momentum vectors of two particles (right).
Transverse Momentum. The transverse plane is crucial for analyzing the dynamics of particle interactions and scatterings. A key quantity in this context is the transverse momentum:
Pr :- 12 + p.y
(58)
which we include as a node feature.
Angular Features. Following the anti-ky algorithm, we compute the azimuth & and the rapidity y. which define the pairwise distance measure used by the algorithm. The azimuth is measured in the transverse plane and is computed as:
Px Pa
Φ = arctan
(59)
The rapidity y is related to the angle of a particle relative to the beamline. It is especially useful in high-energy physics as remains approximately invariant under boosts along the beam axis:
E+Pp
= 2
log E --pp/
(60)
We include and as node features and, analogously to the anti-kr algorithm, use their differences edge features:
Ay y, y As additional edge feature, the angle 0, which accounts for all momentum components and is given by:
A Φi o2
and
(61)
i arccos
(62)
where p= (Pe, Po, P2) the momentum vector of a particle. We illustrate the angles e and 0 in Fig, 10.
20

---

Invariant Mass. Finally, we use the squared invariant mass, which represents the intrinsic mass of a particle independent of its motion or reference frame:
A = E -
(63)
This property plays crucial role in the splits within our hierarchies. As discussed in Section B.2, the invariant mass of parent particle must be greater than the sum of the invariant masses of its children. Hence, for any internal node 2 with children V, and Z32-
VAn V... Aey.
(64)
In the probabilistic setting, we relax this condition to
h-1
Van Σ A.VA. + BAVAS
(65)
We use the difference between the invariant mass of the parent node and the weighted sum of the invariant masses of its children as a node feature. Additionally, we include binary variables whether A,, exceeds the cut thresholds and whether it's larger than 0.
We also derive edge features from this constraint. For each pair of nodes, we compute the energy- momentum vector of their potential parent and calculate the difference between the invariant mass of this potential parent and the sum of the invariant masses of the two children. This edge feature effectively quantifies the plausibility of the nodes being siblings.
B.4.3 Feature Upscaling
To align inputs with the hidden size h and stabilize training, normalize node features with batch normalization, apply a linear layer, and then a LeakyReLU. Because the backbone does not update edge features, we process them using three successive blocks. The upscaling block is shown in Fig. 11. These serve as input to our backbone (see App. B.3).
Batch Norm
Linear Layer
Leaky ReLU
Figure 11: Feature Upscaler. The feature upscaler processes the input features and maps them the hidden dimension
B.5 Baselines
In the following, we give a short description of the baselines. Both generative models operate on probabilistic hierarchies, just like TreeGen. We keep architecture, features. training, and sampling parameters consistent.
Bayesian Flow Networks (BFNs). The most related baselines are standard BFNs. We discuss details and differences possible to TreeGen. While we use our proposed entropy scheduler as explained in Sec. 3, we perform training and sampling described by Graves et al. [19].
model in App. A. For the implementation, we keep it similar as
CatFlow. CatFlow [16] is flow- matching model that evolves samples on the probability simplex by integrating an ODE whose vector field parametrized by a neural network. Training uses cross-entropy loss to predict the ground-truie hierarchy from noisy ones, similar to TreeGen. Two key differences remain: (i) CatFlow's trajectories are not constrained to stay on the simplex during integration. We use Gaussian prior with mean and standard deviation of 0.25, clipped at 10-1 to enable the computation of all features, and (ū) the generative process is deterministic once the prior sample is fixed.
21

---

anti-k; jet clustering algorithm. The anti-k; is a well-established algorithm [7]. It iteratively merges the two entities, i.e., particles or pseudojets, with the smallest pairwise distances dıy The dis- tances are computed based on the transverse momentum, azimuth angle, and rapidity (see App. B.4).
B.6 Metrics
We report two complementary, physics-aware metries.
Valid Hierarchy Percentage (Valid Fraction). A hierarchy is valid if it satisfies three conditions:
(i) it is binary decay tree
(ii) at every internal node the invariant mass exceeds those of the sum of the children
(iii) the squared invariant mass of each internal particle exceeds the datasct-specific cut-off.
For every generated hierarchy, we check these conditions and return the percentage that are valid.
Ratio of Log-Likelihoods (LLH Fraction). Physical plausibility does not guarantee that samples follow the data distribution. Leveraging the Ginkgo simulator, we compute the log- oog-likeihhodd of each generated hierarchy (en)) and of its ground-truth counterpart (fme). We report the average ratio fme/tga over the test set. A value close to indicates that the generative model aligns with the target distribution.
22

---

C Additional Results
C.1
Feature ablation
In the following, we ablate core features discussed in App. B.4 on QCD-S. We start with a base model that neither uses features from the ancestor nor physical properties. The results are shown in Tab. 4.
Table 4: Ablation of ancestor and physically inspired features.
Model
Valid Frac. (1)
LLH Frac.
Base model
0.563110.004
0.937.10.000
Ancestor Features
0.56410.001
0.939
4-Momentum Angular Features
0.645#t0.02
0.982:+0.00
Invariant Mass
0.997
1.003++0.00z
We observe that the base model only provides 56.3% of valid trees. While including the ances- tor features only provides marginal improvements, the physically inspired features yield major improvements.
C.2 Comparison to clustering algorithms.
In Tab. 5, compare TreeGen the CA [15, 46), k, [17], and anti-k; jet clustering algorithm. As we observe, TreeGen consistently outperforms the traditional clustering algorithms in terms of
Table 5: Comparison between clustering algorithms and TreeGen. Best scores in bold.
CA
k+
anti-ke
TreeGen
Dataset
Valid (?)
LLH
Valid (t)
LLH
Valid
LLH
Valid (f)
LLH
QCD-S
0.416
0.989
0.302
1.029
0.840
0.873
0.997 10.001
1.003 =0.002
QCD-M
0.088
0.987
0.050
1.050
0.552
0.752
0.977 10.010
0.994
QCD-L
0.025
0.978
0.008
1.065
0.440
0.636
0.943 20.016
0.975 =0.001
valid hierarchies. In terms of likelihood, the k, algorithm achieves slightly higher scores.
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We propose TreeGen in Sec. and demonstrate its performance empirically on the jet clustering task in Sec. 4
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We include separate limitations paragraph in Sec. 6 and discuss potential future work. All assumptions are stated when applicable.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
24

---

Answer: [Yes]
Justification:
We discuss theoretical properties and assumptions in App. A.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We describe our algorithm in Sec. 3 and the experimental setup in Sec. 4. Furthermore, we specify our architecture and features in App. B.3 and App. B.4. All hyperparameters are listed in App. B.1
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
25

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
The code is available on the project page and contains instructions to reproduce
the main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
cereiissssoooolioio Policy)ly for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. cc/
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All experimental details are specified in Sec. 4 and App. B.I.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Our results include standard deviations computed over four seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
26

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We specify the compute resources in App. B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https / http///eerrpp..eeeee teabblii/Ethissuuiidee
Answer: [Yes]
Justification:
We read the code of ethics and our paper conform with it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper proposes a generative model for hierarchies empirically evaluated on the jet clustering task that does not have direct societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
27

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper does not pose risks for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the Ginkgo simulator.
Guidelines:
The answer NA means that the paper does use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi wwwwde.thoe...... com conddatasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
28

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[NA]
Justification:
We do introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not include research with human subjects crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not include research with human subjects or crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
29

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: Our method does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy (https: http///rppsscccoofffeeeeeeeeeeee for what should or should not be described.
30


Reviews:
[{"review_id": "eplqVna9gZ", "reviewer": "Reviewer_yzpo", "summary": "The paper develops a probabilistic generative framework for hierarchies, named TreeGen. TreeGen is built on top of two earlier frameworks: Bayesian Flow Networks (BFNs) and Bayesian Sample Inference (BSI), after adapting for the tree-input data type. Moreover, a logarithmic entropy scheduler is proposed that allows keeping hyperparameters constant across all sizes of hierarchies. Finally, the performance achieved by the new framework is showcased on Jet Clustering experiments, yielding superior performance compared to CatFlow and standard BFN baselines.", "questions": "1. The contribution that the paper makes on top of the existing BSI and BFN frameworks seems incremental, and I'm not convinced this validates acceptance as a conference paper. As spelled out in A1, differences from BFN amount to slightly different dynamics over sampling and a different loss over training. It would be helpful that differences from BSI/BFN are clearly articulated over the derivation of TreeGen. \n2. Despite including the entropy scheduler as part of the main contributions of the work, I could not find ablation studies that empirically reinforce its usefulness on TreeGen and the considered baselines.\n3. The likelihood maximization experiment contains a single unsophisticated baseline. Could authors also include the achieved performance by the probabilistic baselines considered in the rest of the experiments (CatFlow, standard BFNs)?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 1, "strengths_and_weaknesses": "**Strengths**\n1. The significance of the paper is timely for the Jet Clustering problems. The experimental results are mostly convincing, showing improvement upon state-of-the-art methods.\n2. The paper is well-written in general. A good job is done at motivating the problem and presenting related work.\n\n**Weaknesses**\nPlease see below.", "was_revised": true, "final_justification": "Thanks to the authors for their copious response and further experimentation. I have now increased my score accordingly.", "rebuttal": "We thank the reviewer for their thorough feedback. We address their remarks in the following.\n\n**Comment:** It would be helpful that differences from BSI/BFN are clearly articulated over the derivation of TreeGen.  \n**Response:** We would like to reiterate our contributions and differences from previous works. First, we derive a categorical of BSI, which uses the Bayesian update proposed by BFNs [1], but obtain a different loss and update rule (see App. A.1). Second, we simplify the theoretical framework by omitting the input distribution, output distribution, sender distribution, receiver distribution, flow distribution, and update distribution. Third, we extend this categorical framework to hierarchies by connecting it to the hierarchical clustering literature [2], which has not been done before. Finally, we present a probabilistic view of jet clustering and demonstrate that TreeGen empirically outperforms prior methods by a substantial margin.\n\nWe will highlight these similarities and differences in section 3 of the updated manuscript.\n\n**Comment:** The proposed entropy scheduler is not ablated.  \n**Response:** We agree with the reviewer that this ablation could be interesting and have conducted additional experiments. However, our entropy scheduler was not designed to outperform the scheduler by Graves et al. [1], but rather to eliminate the need to retune it for different experiments. \n\nIn the following, we show experiments conducted with their original scheduler and varying C. In [1], the authors used C=0.75 and C=3.0 in different experiments. Additionally, we tested C=6.0 and C=9.0. The following table shows the results on QCD-S.\n\n\n| Scheduler   | Valid Fraction    | LLH Fraction      |\n|-------------|-------------------|-------------------|\n| Ours        | 0.997 ± 0.001     | 1.003 ± 0.002     |\n| C = 0.75    | 0.410 ± 0.019      | 0.990 ± 0.002      |\n| C = 3.0     | 0.989 ± 0.001     | 1.002 ± 0.002     |\n| C = 6.0     | 0.994 ± 0.001     | 1.003 ± 0.002     |\n| C = 9.0     | 0.992 ± 0.004     | 1.002 ± 0.001 |\n\nThe scheduler of Graves et al. [1] is largely influenced by C and reaches matching performance after tuning. We will include these results in the updated manuscript.\n\n\n\n**Comment:** The likelihood maximization experiment contains a single unsophisticated baseline. Could authors also include the achieved performance by the probabilistic baselines (CatFlow, standard BFNs) considered in the rest of the experiments?  \n**Response:** We follow the reviewer's suggestion and include the experiments in Figure 7 in the updated manuscript. In the following table, we present the maximum log-likelihoods for the different methods after $N$ samples:\n\n\n| Model | 1 | 5 | 10 | 15 | 20 | 25 | 30 | 35 | 40 | 45 | 50 |\n| -- | --  | --  | --  | --  | --  | --  | --  | --  | --  | --  | --  |\n| CatFlow | -48.267 $\\pm$ 0.072 | -47.973 $\\pm$ 0.026 | -47.219 $\\pm$ 0.085 | -46.879 $\\pm$ 0.072 | -46.669 $\\pm$ 0.056 | -46.504 $\\pm$ 0.061 | -46.405 $\\pm$ 0.057 | -46.323 $\\pm$ 0.047 | -46.256 $\\pm$ 0.044 | -46.194 $\\pm$ 0.047 | -46.148 $\\pm$ 0.046 |\n| BFN | -49.151 $\\pm$ 0.100 | -47.224 $\\pm$ 0.022 | -46.597 $\\pm$ 0.027 | -46.324 $\\pm$ 0.030 | -46.153 $\\pm$ 0.025 | -46.038 $\\pm$ 0.024 | -45.946 $\\pm$ 0.027 | -45.881 $\\pm$ 0.026 | -45.822 $\\pm$ 0.029 | -45.778 $\\pm$ 0.028 | -45.741 $\\pm$ 0.028 |\n| TreeGen | **-48.531** $\\pm$ 0.084 | **-46.476** $\\pm$ 0.054 | **-46.011** $\\pm$ 0.067 | **-45.810** $\\pm$ 0.059 | **-45.697** $\\pm$ 0.045 | **-45.610** $\\pm$ 0.046 | **-45.548** $\\pm$ 0.045 | **-45.496** $\\pm$ 0.041 | **-45.460** $\\pm$ 0.038 | **-45.429** $\\pm$ 0.036 | **-45.402** $\\pm$ 0.036 |\n\n\nGreedy: -45.746\n\nAs shown, TreeGen consistently outperforms the two baselines, BFN and CatFlow. Notably, BFN requires 50 samples to match the log-likelihood of the greedy baseline, while CatFlow remains below its performance. This again highlights the advancements of TreeGen over BFNs.\n\n\nWe again thank the reviewer for their feedback and are happy to address any remaining concerns.\n\n[1] **Graves, Alex, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez.** \"Bayesian flow networks.\" arXiv preprint arXiv:2308.07037 (2023).\n\n[2] **Zügner, Daniel, Bertrand Charpentier, Morgane Ayle, Sascha Geringer, and Stephan Günnemann.** \"End-to-end learning of probabilistic hierarchies on graphs.\" In International Conference on Learning Representations. 2021."}, {"review_id": "oMRfuFPU32", "reviewer": "Reviewer_Za3S", "summary": "How should we model distributions over hierarchies? This long addressed question is approached with Bayesian Flow Networks and a novel scheduler for smooth and consistent entropy decay across varying numbers of categories. The authors describe their probabilistic framework, present the application for clustering particles in jet physics. The proposed method uses bayesian updates and gradient-based learning for the feature encoder.", "questions": "* Can you describe more why some of the past work such as \"Hierarchical clustering in particle physics through reinforcement learning\" do not provide ways to sample hierarchies?\n* Can you say more about the assumptions of the given distribution over trees? Should other models be considered?\n* What happens if no simulator is available? How should we think about OOD data?", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "TreeGen is an interesting paper, which explores the intersections of probabilistic methods for discovering hierarchical structure, feature learning, and applications in niche and intriguing areas. \n\nStrengths of the paper include: \n* Quality: The paper reads in a complete way, experiments and hypotheses are reasonable and explored adequately. \n* Clarity: The paper reads clearly, explains this in an understandable way. \n\nWeaknesses include: \n* Quality: It is not clear if there should be more comparison & discussion around the assumptions of the distribution over trees given by treegen's model compared to other families of models considered. \n* Originality: The differentiating factors seem to be relatively minor (though crucial for performance) (at least as I understand) compared BFNs and past work. It is also not a new application of machine learning. I see this as the key weakness. The methodology seems to be only minor deviations from the past work, the applications are not especially novel either. \n* Significance: There are certainly people who are deeply focused on the kinds of problems in this paper, however, it is not clear how significant the results are to a broader audience. The advances on BFNs seem relatively minor\n\nMissing references: \n* [Variational Pseudo Marginal Methods for Jet Reconstruction in Particle Physics](https://openreview.net/forum?id=pCapRF2vFf) - seems to be a key reference", "was_revised": true, "final_justification": "The authors address much of my concerns. Given the support of other reviewers and the addressing of comments, I am quite torn about whether or not to raise my score. However, I am not able to champion the paper myself and so I leave my original rating. That said, I am OK if the paper is accepted.", "rebuttal": "We thank the reviewer for their valuable feedback. In the following, we address their comments.\n\n**Comment:** Quality: It is not clear if there should be more comparison & discussion around the assumptions of the distribution over trees given by treegen's model compared to other families of models considered.  \n**Response:** The only assumption of the distribution is that the structures are directed acyclic graphs, i.e., trees/hierarchies. Both CatFlow and BFN impose the same restrictions. In contrast, the agglomerative baselines are tailored to the jet clustering task and only produce binary hierarchies. We will clarify this in the updated manuscript.\n\n\n**Comment:** Originality: The differentiating factors seem to be relatively minor (though crucial for performance) compared to BFNs and past work. It is also not a new application of machine learning.  \n**Response:** We respectfully disagree with the reviewer and would like to reiterate the contributions and novelty of our paper to highlight differences from previous works. \nFirst, our model results from a simplified theoretical perspective and derivation. Bayesian Flow Networks are derived by introducing an input, output, sender, receiver, flow, and update distribution. TreeGen, however, does not require these and is derived via a simple Bayesian update, yielding a different training loss and update equations. \nSecond, we adapt generative models to hierarchies by drawing connections to the clustering literature [9], introducing an inductive bias ensuring valid hierarchies with a novel scheduler balancing the entropy (see Sec. 3).\nFinally, unlike dominant greedy approaches, we present a probabilistic perspective on jet clustering which learns the full posterior distribution over trees and demonstrate that TreeGen faithfully learns the posterior distribution with substantial performance gains over previous works.\n\n\n**Comment:** It is unclear how significant the results are to a broader audience. The advances on BFNs seem relatively minor  \n**Response:** Although the resulting model resembles a similar update compared to BFNs, they are derived from a simplified perspective and yield substantial performance improvements, i.e., on the larger W-Boson jets dataset, the valid fraction increases by nearly 60%. \n\nOur results are important to:\n1. High-energy physics, where TreeGen improves over traditional jet-clustering algorithms. \n2. Hierarchical modelling, since our method is not limited to jets but can be applied to any hierarchical data. \n3. Generative modelling, as our derivation simplifies BFNs and allows for a broader range of modelling choices with substantial empirical improvements.\n\n\n**Comment:** Can you describe more why some of the past work such as \"Hierarchical clustering in particle physics through reinforcement learning\" do not provide ways to sample hierarchies?  \n**Response:**\nMost existing methods search for the maximum-likelihood tree $\\mathcal{T}^* = \\operatorname{arg} \\max_{\\mathcal{T}} p(\\mathcal{T}\\mid X)$. Algorithms such as [6,7,8] frame this as a clustering problem and build the tree bottom-up. Often, these algorithms are deterministic and focus on the small subset of high likelihood trees. For this reason, the constructed trees are not samples from the posterior $p(\\mathcal{T} \\mid X)$.\nSampling from the posterior $p(\\mathcal{T}\\mid X)$ is hard because the number of candidate trees grows doubly exponentially with the number of leaves, limiting earlier approaches such as Cluster Trellis [4] to trees with fewer than 20 elements. Making use of generative modelling, TreeGen is able to produce samples from the approximated posterior for trees in the order of 100 elements.\n\n**Comment:** Can you say more about the assumptions of the given distribution over trees? Should other models be considered?  \n**Response:** \nOur parametrization is tailored to directed acyclic graphs. We do not place any assumptions or constraints on the trees in the distribution, and they can have varying sizes and branching factors.\n\n**Comment:** What happens if no simulator is available? How should we think about OOD data?  \n**Response:** In high-energy physics, there are typically various simulators available [1,2,3]. However, if no simulators are available, one can first gather a small set of samples from the ground-truth posterior if the probability of splits is known [4], serving as a training set. Alternatively, if neither of these is available, one could apply reweighted importance sampling, which only requires an unnormalized energy function [5] (Appendix D.2). In the case of OOD data, TreeGen can be finetuned on real observations, which we leave for future work.\n\n\nWe again thank the reviewer for their valuable feedback and have included the missing reference in the updated manuscript. We are happy to address any remaining and upcoming concerns.\n\n[1] **Bellm, Johannes, Stefan Gieseke, David Grellscheid, Simon Plätzer, Michael Rauch, Christian Reuschle, Peter Richardson et al.** \"Herwig 7.0/Herwig++ 3.0 release note.\" The European Physical Journal C 76, no. 4 (2016): 196.\n\n[2] **Bierlich, Christian, Smita Chakraborty, Nishita Desai, Leif Gellersen, Ilkka Helenius, Philip Ilten, Leif Lönnblad et al.** \"A comprehensive guide to the physics and usage of PYTHIA 8.3.\" SciPost Physics Codebases (2022): 008.\n\n[3] **Agostinelli, Sea, John Allison, K. al Amako, John Apostolakis, Henrique Araujo, Pedro Arce, Makoto Asai et al.** \"Geant4—a simulation toolkit.\" Nuclear instruments and methods in physics research section A: Accelerators, Spectrometers, Detectors and Associated Equipment 506, no. 3 (2003): 250-303.\n\n[4] **Macaluso, Sebastian, Craig Greenberg, Nicholas Monath, Ji Ah Lee, Patrick Flaherty, Kyle Cranmer, Andrew McGregor, and Andrew McCallum.** \"Cluster trellis: Data structures & algorithms for exact inference in hierarchical clustering.\" In International Conference on Artificial Intelligence and Statistics, pp. 2467-2475. PMLR, 2021.\n\n[5] **Tong, Alexander, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio.** \"Improving and generalizing flow-based generative models with minibatch optimal transport.\" arXiv preprint arXiv:2302.00482 (2023).\n\n[6] **Greenberg, C.S., Macaluso, S., Monath, N., Dubey, A., Flaherty, P., Zaheer, M., Ahmed, A., Cranmer, K. &amp; McCallum, A.**. (2021). “Exact and approximate hierarchical clustering using A*“.\n\n[7] **Brehmer, Johann et al.** “Hierarchical clustering in particle physics through reinforcement learning.” ArXiv abs/2011.08191 (2020).\n\n[8] **Greenberg, Craig S., Sebastian Macaluso, Nicholas Monath, Ji-Ah Lee, Patrick Flaherty, Kyle Cranmer, Andrew McGregor, and Andrew McCallum**. \"Data Structures & Algorithms for Exact Inference in Hierarchical Clustering.\" arXiv preprint arXiv:2002.11661 (2020).\n\n[9] **Zügner, Daniel, Bertrand Charpentier, Morgane Ayle, Sascha Geringer, and Stephan Günnemann.** \"End-to-end learning of probabilistic hierarchies on graphs.\" In International Conference on Learning Representations. 2021."}, {"review_id": "JRk9eQMXpx", "reviewer": "Reviewer_tb6N", "summary": "Jet clustering in particle physics is an algorithmic process that links low energy hadronic detector outputs to high energy partonic states. While the standard usage relies on unsupervised clustering methods which is useful for discovery and is also agnostic to the required analysis, it is known that different signals require bespoke setting of parameters such as jet radius.  These are often tested on simulated showered events in order to extract physically meaningful quantities from these parameter-dependent jets.  This paper seeks to extract more meaningful patterns by learning a generative model for probabilistic trees to represent stochastically generated showers.   They base this on prior work on flow networks with Bayesian conditional updates with an entropy based accumulated signal, which is a scaling of a Gaussian random variable that approximates a discrete sum over categorical variables.  The conditional update over the discrete random variable is extended to the hierarchical setting required for tree generation.  The evaluation is done using Ginkgo datasets for QCD and W boson showers. For the small datasets a direct comparison of the posterior distributions is feasible and shown.  Per generated tree likelihoods are reported as well as maximum a posteriori tree likelihoods, a property enabled by the generative nature of the learned tree models.", "questions": "At the risk of repeating the comments from above, here are some questions.  \n\n(1) Could you add comparisons with Cambridge-Aachen and the $k_T$ algorithms?  \n\n(2) From the numbers in Table 3, it is clear that the invariant mass plays a key role in outperforming anti-$k_T$.  Could you provide some explanations as to why this might be the case?  \n\n(3) Why do the angular features not guarantee l. 968-9 (ii-iii) is satisfied, whereas an unsupervised method appears to do so by virtue of the agglomerative step?  \n\n(4) Does this not invalidate the method as a reliable probe of jets that are typically extracted via anti-$k_T$ at the LHC and later processed via C-A and $k_T$ algorithms for substructure and tagging studies, and for pileup removal?  \n\n(5) Please provide potentially physically meaningful insights that this, otherwise elegant, approach might help extract.  Given the concerns I have raised, being able to appreciate what this way of processing jets might bring to typical tasks undertaken would enhance my assessment of the paper's contribution.", "limitations": "Yes.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The paper is very well written.  It takes the updates of Graves et al from the perspective of Lienen et al and combines these update rules with a proposed entropic accumulator, following which this is adapted to the hierarchical construction of Zuegner et al.  This facilitates the application to the particle physics context, which makes this paper an interesting contribution.  While the results depict what is naturally available given a probabilistic model -- likelihoods and posterior averages -- it would have helped to see a set of conceptual insights that get opened up as a consequence of these powers.  \n\nThe Ginkgo trees are generated with an angular distance $d_{ij}=\\Delta R^2_{ij}$ between two particles $i$, $j$ in mind.  The comparison shown in C.2 Table 4 is against anti-$k_T$ which uses a different $d_{ij}$ using higher values of $p_T$ to act as a seed to agglomeratively create hierarchies.  This is not a fair comparison.  It would be helpful to see the comparisons with Cambridge-Aachen and the $k_T$ algorithm.  Moreover, from the numbers in Table 3, it is clear that the invariant mass plays a key role in outperforming anti-$k_T$.  Could you provide some explanations as to why this might be the case?  Why do the angular features not guarantee l. 968-9 (ii-iii) is satisfied, whereas an unsupervised method appears to do so by virtue of the agglomerative step?  Does this not invalidate the method as a reliable probe of jets that are typically extracted via anti-$k_T$ at the LHC and later processed via C-A and $k_T$ algorithms for substructure and tagging studies, and for pileup removal?  While the application of the generative framework to the particle physics context is elegant, it is not clear what insights this helps extract that is physically meaningful.", "was_revised": true, "final_justification": "I have read the comments by the other reviewers and the way the authors have engaged with them.  I have upped my scores accordingly.", "rebuttal": "We thank the reviewer for their thorough and thoughtful feedback. In the following, we address their concerns and comments.\n\n**Comment:** Could you add comparisons with Cambridge-Aachen and the $k_t$ algorithms?  \n**Response:** We follow the reviewer's suggestion and include the two baselines. We reported the anti-$k_t$ algorithm because the CA and $k_t$ often violated constraints (ii) in App. B.6, leading to low valid fractions.\n\nValid fraction:\n\n| Dataset | CA     | $k_t$   | anti-$k_t$ | TreeGen           |\n|---------|--------|-------|----------|------------------------|\n| QCD-S   | 0.416  | 0.302 | 0.840    | **0.997 $\\pm$ 0.001**       |\n| QCD-M   | 0.088  | 0.050 | 0.552    | **0.977 $\\pm$ 0.010**       |\n| QCD-L   | 0.025  | 0.008 | 0.440    | **0.943 $\\pm$ 0.016**       |\n\nLLH fraction:\n\n| Dataset | CA     | $k_t$   | anti-$k_t$ | TreeGen           |\n|---------|--------|-------|----------|------------------------|\n| QCD-S   | 0.989  | **1.029** | 0.873    | 1.003 $\\pm$ 0.002       |\n| QCD-M   | 0.987  | **1.050** | 0.752    | 0.994 $\\pm$ 0.002       |\n| QCD-L   | 0.978  | **1.065** | 0.636    | 0.975 $\\pm$ 0.001       |\n\n\nAs we observe, the CA and $k_t$ algorithms achieve higher LLH fractions than the anti-$k_t$ algorithm but come with a lower validity.\n\n**Comment:** From the numbers in Table 3, it is clear that the invariant mass plays a key role in outperforming anti-$k_t$. Could you provide some explanations as to why this might be the case?  \n**Response:** The invariant mass features play a crucial role in the validity of the trees. This results from the conditions (ii) and (iii) in App. B.6, inducing constraints on the invariant mass of the trees.\n\nA single misconnected branch, i.e., split, is sufficient to violate the validity of a tree. Therefore, the invariant mass features become substantial for the valid fraction, while the LLH only receives a minor gain.\n\n**Comment:** Why do the angular features not guarantee L968-9 (ii-iii) is satisfied, whereas an unsupervised method appears to do so by virtue of the agglomerative step?  \n**Response:** While the angular features are theoretically sufficient, we noticed that the information is more easily accessible if further processed. Furthermore, while valid branches are essential for the evaluation, our training loss does not prioritize these. For illustration, we have added the final training loss in Tab. 3:\n\n|Model|Valid Frac.|LLH Frac.| Training Loss | \n|--|--|--|--|\n| Base model| 0.563 $\\pm$ 0.004| 0.937 $\\pm$ 0.002| 0.362 $\\pm$ 0.003 |\n| + Ancestor Features| 0.564 $\\pm$ 0.001| 0.939 $\\pm$ 0.001| 0.361 $\\pm$ 0.004 |\n| + 4-Momentum + Angular Features  | 0.645 $\\pm$ 0.021| 0.982 $\\pm$ 0.001| 0.334 $\\pm$ 0.002 |\n| + Invariant Mass | 0.997 $\\pm$ 0.001 | 1.003 $\\pm$ 0.002 | 0.319 $\\pm$ 0.002 |\n\nAs we observe, including 4-Momentum and angular features reduces the loss but only slightly improves validity. Including the invariant mass, however, substantially improves the valid fraction with a minor impact on the loss.\n\n\n**Comment:** Does this not invalidate the method as a reliable probe of jets that are typically extracted via anti-$k_t$ at the LHC and later processed via C-A and $k_t$ algorithms for substructure and tagging studies, and for pileup removal?  \n**Response:** TreeGen offers access to the full posterior over trees, something deterministic algorithms lack. Even if later steps use CA/$k_t$, TreeGen can provide multiple high‑probability trees and potential uncertainty estimates for downstream analyses.\n\n**Comment:** Please provide potentially physically meaningful insights that this, otherwise elegant, approach might help extract. Given the concerns I have raised, being able to appreciate what this way of processing jets might bring to typical tasks undertaken would enhance my assessment of the paper's contribution.  \n**Response:** The key insight of our approach is the feasibility of modelling jet clustering as a generative task. While classical agglomerative jet algorithms output one tree per event, TreeGen models a distribution over all potential hierarchies consistent with the observed four‑momentum. This probabilistic framework provides two key advantages over deterministic trees:\n\n1. Deterministic approaches only provide a single tree covering only a minor part of the posterior distribution, as many different splitting processes can lead to the same observed leaves. This is especially important for large trees due to the combinatorial explosion.\n2. Modelling the conditional trees in a distribution provides empirical uncertainty estimates via the empirical sampled distribution. The samples can then be filtered by specific desired properties, e.g., validity or high likelihoods. Our likelihood maximization experiment provides an example (see Fig. 7).\n\nMoreover, TreeGen provides better scalability than exact posterior methods [1], while faithfully approximating the posterior (see Fig. 6). Finally, future work could explore imputation/inpainting tasks inferring only parts of the tree. \n\nWe again thank the reviewer for their feedback and are confident that we were able to improve the manuscript. \n\n[1] **Macaluso, Sebastian, Craig Greenberg, Nicholas Monath, Ji Ah Lee, Patrick Flaherty, Kyle Cranmer, Andrew McGregor, and Andrew McCallum.** \"Cluster trellis: Data structures & algorithms for exact inference in hierarchical clustering.\" In International Conference on Artificial Intelligence and Statistics, pp. 2467-2475. PMLR, 2021."}, {"review_id": "J6eEzjyvZ6", "reviewer": "Reviewer_uAVU", "summary": "The paper presents TreeGen, a novel generative model that extends Bayesian Flow Networks (BFNs) to capture distributions over hierarchical data. In addition to the model, the authors propose a new entropy scheduling strategy for BFNs that performs well across various structures without requiring hyperparameter tuning. TreeGen is designed to solve jet clustering problems, and the experimental results demonstrate that it outperforms existing methods across multiple datasets.", "questions": "1. As someone not familiar with jet clustering, I am somewhat confused by the underlying prediction problem. If I understand it correctly, each node in the tree corresponds to an individual particle and the goal is to learn a graph neural network that generalizes across different particles, i.e., predict the probability of each edge in the tree given the representations of the parent node, the child node, and the edge itself. Is that assessment correct?\n2. In Section 3, each node in the tree is treated as a random variable which is observed multiple times. Does that mean that each node represents a type of particle rather than individual particles, and that is why we can collect repeated observations for each?\n3. Do all trees in the training and test datasets have to share the same structure? That is, would be possible to generalize across different structures provided we have good features for all nodes?", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths\n-----------\n- The main method of the paper, TreeGen, extends Bayesian Flow Networks (BFNs) to hierarchical data and is novel, to the best of my knowledge.\n- The experiments are well designed and the results indicate TreeGen outperforms CatFLow and BFNs on a number of jet clustering datasets.\n- The mathematical development is clear and sound.\n- The paper is very well written and organized.\n\n### Weaknesses\n------------\nI should mention that I am not familiar with the literature on jet clustering, so I am not in a position to assess whether it is a problem of particular relevance to the NeurIPS community. From my perspective, the paper seems more focused on addressing the jet clustering task itself than on introducing a broadly applicable machine learning method. If jet clustering is already well-established as a relevant topic within the community, then this focus makes perfect sense. However, if that is not the case, it might be helpful for the authors to make the paper more approachable to a broader machine learning audience. For example, by providing a clearer explanation of the problem (see questions below) and highlighting how the proposed approach or insights could be useful in other machine learning contexts.", "was_revised": true, "final_justification": "I think the proposed methods are novel and interesting enough to justify acceptance. However, as I mentioned in my review, I would have liked to see the methods applied in other domains or at least presented in a more general fashion, as I see jet clustering as a somewhat niche application within the Neurips community. This is the reason I maintain my score as a 4 for now. Of course, as someone not entirely familiar with the area of machine learning for physics, I am happy to raise my score to a 5 if I have got the wrong impression about the relevance of the topic here.", "rebuttal": "We thank the reviewer for their valuable feedback. In the following, we address their comments.\n\n**Comment:** As someone not familiar with jet clustering, I am somewhat confused by the underlying prediction problem. If I understand it correctly, each node in the tree corresponds to an individual particle and the goal is to learn a graph neural network that generalizes across different particles, i.e., predict the probability of each edge in the tree given the representations of the parent node, the child node, and the edge itself. Is that assessment correct?  \n**Response:** The assessment is correct. Each node represents a particle, and our network predicts the probability of each possible parent-child edge based on the features of the particles and current edges in the tree. The edges themselves encode the splitting process, i.e., the parent particle splits into the two child particles. \n\n**Comment:** In Section 3, each node in the tree is treated as a random variable which is observed multiple times. Does that mean that each node represents a type of particle rather than individual particles, and that is why we can collect repeated observations for each?  \n**Response:** The tree corresponds to a splitting process. A parent node splits into its children. The root corresponds to the initial particle, while the leaves represent the observed particles. We collect repeated observations by making predictions with the neural network, allowing us to refine our beliefs. We added a paragraph in the manuscript clarifying how particles correspond to nodes in our setup.\n\n**Comment:** Do all trees in the training and test datasets have to share the same structure? That is, would be possible to generalize across different structures provided we have good features for all nodes?  \n**Response:** All trees in the datasets are binary trees, i.e., have binary branches but vary in size, shape, and features (4 momentum). We will clarify this in the experiments section.\n\nFinally, we would like to highlight the importance of our work to different audiences:\n1. High-energy physics, where TreeGen improves over traditional jet-clustering algorithms. \n2. Hierarchical modelling, since our method is not limited to jets but can be applied to any hierarchical data. \n3. Generative modelling, as our derivation simplifies BFNs and allows for a broader range of modelling choices with substantial empirical improvements.\n\nWe thank the reviewer again for their feedback and have improved the clarity based on their comments by introducing an additional paragraph in the updated manuscript describing the jet clustering task. We are happy to address any remaining concerns."}]

Author rebuttal:

Meta-review: This paper describes a generalization of Bayesian Flow Networks to model hierarchical data, as well as a new entropy scheduling strategy, and experimental results for jet clustering problems. Informally speaking, the method starts with a probabilistic hierarchy which is gradually reduced in entropy until a discrete hierarchy is found. Reviewers agreed the method was novel, the paper was clearly written, and the experimental results were of high quality, plausibly demonstrating a practical advance.
Decision: Accept (poster)