TreeGen: A Bayesian Generative Model for Hierarchies

MCML Authors

Marcel Kollovieh

→ Group Stephan Günnemann
 Data Analytics & Machine Learning

Filippo Guerranti

→ Group Stephan Günnemann
 Data Analytics & Machine Learning

Stephan Günnemann

Prof. Dr.

Core PI

Data Analytics & Machine Learning

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.

inproceedings KFG+25

---

NeurIPS 2025

39th Conference on Neural Information Processing Systems. San Diego, CA, USA, Nov 30-Dec 07, 2025.

---

Authors

M. Kollovieh • N. Fleischmann • F. Guerranti • B. Charpentier • S. Günnemann

---

Links

URL

---

In Collaboration

 Pruna AI

---

Research Area

 A3 | Computational Models