[2603.23398] Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation

Computer Science > Machine Learning arXiv:2603.23398 (cs) [Submitted on 24 Mar 2026] Title:Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation Authors:Michal Balcerak, Suprosana Shit, Chinmay Prabhakar, Sebastian Kaltenbach, Michael S. Albergo, Yilun Du, Bjoern Menze View a PDF of the paper titled Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation, by Michal Balcerak and 5 other authors View PDF HTML (experimental) Abstract:Energy-based models for discrete domains, such as graphs, explicitly capture relative likelihoods, naturally enabling composable probabilistic inference tasks like conditional generation or enforcing constraints at test-time. However, discrete energy-based models typically struggle with efficient and high-quality sampling, as off-support regions often contain spurious local minima, trapping samplers and causing training instabilities. This has historically resulted in a fidelity gap relative to discrete diffusion models. We introduce Graph Energy Matching (GEM), a generative framework for graphs that closes this fidelity gap. Motivated by the transport map optimization perspective of the Jordan-Kinderlehrer-Otto (JKO) scheme, GEM learns a permutation-invariant potential energy that simultaneously provides transport-aligned guidance from noise toward data and refines samples within regions of high data likelihood. Further, we introduce a sampling protocol that leverages an energy-based swit...