TreeGen: A Bayesian Generative Model for Hierarchies
Marcel Kollovieh, Nils Fleischmann, Filippo Guerranti, Bertrand Charpentier, Stephan Günnemann
NeurIPS, 2025

Links

[OpenReview]

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.

Citation

If you use this work in your research, please cite this paper:

@inproceedings{kollovieh2025treegen,  

title={TreeGen: A Bayesian Generative Model for Hierarchies},  

author={Kollovieh, Marcel and Fleischmann, Nils and Guerranti, Filippo and Charpentier, Bertrand and G{\"u}nnemann, Stephan},  

booktitle={Neural Information Processing Systems (NeurIPS)},  

year={2025},  

url={https://openreview.net/forum?id=d2EouMhAAq} }