Title: TreeGen: A Bayesian Generative Model for Hierarchies
Venue: NeurIPS 2025
Accepted: yes
Keywords: generative models

hierarchies

jet clustering

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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nJet 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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.\n\nQUESTIONS:\nAt 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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nHow 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.\n\nSTRENGTHS AND WEAKNESSES:\nTreeGen 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\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n1. 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)?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 1, "DATE": "2025-07-03"}