Home | Publications | Kfg 25

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.

Principal Investigator

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