TreeGen: A Bayesian Generative Model for Hierarchies

Marcel Kollovieh, Nils Fleischmann, Filippo Guerranti, Bertrand Charpentier, Stephan Günnemann

Keywords: generative models, hierarchies, jet clustering, flow matching

TL;DR: A novel generative model for hierarchies based on Bayesian Flow Networks.

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.

Primary Area: Deep learning (e.g., architectures, generative models, optimization for deep networks, foundation models, LLMs)

Submission Number: 24062