TreeGen: A Bayesian Generative Model for Hierarchies

We thank the reviewer for their valuable feedback. In the following, we address their comments.

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

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

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

Finally, we would like to highlight the importance of our work to different audiences:

1. High-energy physics, where TreeGen improves over traditional jet-clustering algorithms.
2. Hierarchical modelling, since our method is not limited to jets but can be applied to any hierarchical data.
3. Generative modelling, as our derivation simplifies BFNs and allows for a broader range of modelling choices with substantial empirical improvements.

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

Dear reviewer,

We hope our additional clarifications have addressed your comments satisfactorily. We remain fully available to provide clarifications or answer any further questions they may have.

Again, thank you for your efforts and thoughtful feedback.

We thank the reviewer for their thorough and thoughtful feedback. In the following, we address their concerns and comments.

Comment: Could you add comparisons with Cambridge-Aachen and the $k_t$ algorithms?
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.

Valid fraction:

LLH fraction:

As we observe, the CA and $k_t$ algorithms achieve higher LLH fractions than the anti-$k_t$ algorithm but come with a lower validity.

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

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

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?
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:

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

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

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

1. 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.
2. 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).

Moreover, 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.

We again thank the reviewer for their feedback and are confident that we were able to improve the manuscript.

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

We thank the reviewer for their valuable feedback. In the following, we address their comments.

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

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.
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. First, 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. Second, 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). Finally, 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.

Comment: It is unclear how significant the results are to a broader audience. The advances on BFNs seem relatively minor
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%.

Our results are important to:

1. High-energy physics, where TreeGen improves over traditional jet-clustering algorithms.
2. Hierarchical modelling, since our method is not limited to jets but can be applied to any hierarchical data.
3. Generative modelling, as our derivation simplifies BFNs and allows for a broader range of modelling choices with substantial empirical improvements.

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?
Response: Most 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)$. Sampling 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.

Comment: Can you say more about the assumptions of the given distribution over trees? Should other models be considered?
Response: Our 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.

Comment: What happens if no simulator is available? How should we think about OOD data?
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.

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

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

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

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

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

[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).

[6] Greenberg, C.S., Macaluso, S., Monath, N., Dubey, A., Flaherty, P., Zaheer, M., Ahmed, A., Cranmer, K. & McCallum, A.. (2021). “Exact and approximate hierarchical clustering using A*“.

[7] Brehmer, Johann et al. “Hierarchical clustering in particle physics through reinforcement learning.” ArXiv abs/2011.08191 (2020).

[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).

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

We thank the reviewer for their thorough feedback. We address their remarks in the following.

Comment: It would be helpful that differences from BSI/BFN are clearly articulated over the derivation of TreeGen.
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.

We will highlight these similarities and differences in section 3 of the updated manuscript.

Comment: The proposed entropy scheduler is not ablated.
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.

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

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

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?
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:

Greedy: -45.746

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

We again thank the reviewer for their feedback and are happy to address any remaining concerns.

[1] Graves, Alex, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. "Bayesian flow networks." arXiv preprint arXiv:2308.07037 (2023).

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