# Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation

arXiv: https://arxiv.org/abs/2603.23398 · HTML: https://mlsanity.com/paper/2603.23398/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 30 of its 46 references that are in the map.
- First posted on arXiv 24 Mar 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 30)

- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [DiGress: Discrete Denoising diffusion for graph generation](https://mlsanity.com/paper/2209.14734/) (2022)
- [Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design](https://mlsanity.com/paper/2402.04997/) (2024)
- [Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models](https://mlsanity.com/paper/1811.12823/) (2018)
- [Crystal Diffusion Variational Autoencoder for Periodic Material Generation](https://mlsanity.com/paper/2110.06197/) (2021)
- [Pre-training Molecular Graph Representation with 3D Geometry](https://mlsanity.com/paper/2110.07728/) (2021)
- [Fr'echet ChemNet Distance: A metric for generative models for molecules in drug discovery](https://mlsanity.com/paper/1803.09518/) (2018)
- [Optimal Transport for structured data with application on graphs](https://mlsanity.com/paper/1805.09114/) (2018)
- [Graph Inductive Biases in Transformers without Message Passing](https://mlsanity.com/paper/2305.17589/) (2023)
- [Machine learning for molecular simulation](https://mlsanity.com/paper/1911.02792/) (2019)
- [Stochastic interpolants with data-dependent couplings](https://mlsanity.com/paper/2310.03725/) (2023)
- [Oops I Took A Gradient: Scalable Sampling for Discrete Distributions](https://mlsanity.com/paper/2102.04509/) (2021)
- [Energy-Based Diffusion Language Models for Text Generation](https://mlsanity.com/paper/2410.21357/) (2024)
- [GraphEBM: Molecular Graph Generation with Energy-Based Models](https://mlsanity.com/paper/2102.00546/) (2021)
- [Variational Flow Matching for Graph Generation](https://mlsanity.com/paper/2406.04843/) (2024)
- [DeFoG: Discrete Flow Matching for Graph Generation](https://mlsanity.com/paper/2410.04263/) (2024)
- [Energy-based Out-of-Distribution Detection for Graph Neural Networks](https://mlsanity.com/paper/2302.02914/) (2023)
- [A Langevin-like Sampler for Discrete Distributions](https://mlsanity.com/paper/2206.09914/) (2022)
- [Cometh: A continuous-time discrete-state graph diffusion model](https://mlsanity.com/paper/2406.06449/) (2024)
- [Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling](https://mlsanity.com/paper/2504.10612/) (2025)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [First-order Conditions for Optimization in the Wasserstein Space](https://mlsanity.com/paper/2209.12197/) (2022)
- [Energy-based Epistemic Uncertainty for Graph Neural Networks](https://mlsanity.com/paper/2406.04043/) (2024)
- [Molecule Joint Auto-Encoding: Trajectory Pretraining with 2D and 3D Diffusion](https://mlsanity.com/paper/2312.03475/) (2023)
- [Controlled Generation with Equivariant Variational Flow Matching](https://mlsanity.com/paper/2506.18340/) (2025)
- [Autoregressive Language Models are Secretly Energy-Based Models: Insights into the Lookahead Capabilities of Next-Token Prediction](https://mlsanity.com/paper/2512.15605/) (2025)
- [Follow the Energy, Find the Path: Riemannian Metrics from Energy-Based Models](https://mlsanity.com/paper/2505.18230/) (2025)
- [Variational Analysis in the Wasserstein Space](https://mlsanity.com/paper/2406.10676/) (2024)
- [Graph Diffusion that can Insert and Delete](https://mlsanity.com/paper/2506.15725/) (2025)

## Built on by (most cited first, 1 of 1)

- [Learning Graph Energies for Finite-Time Generation](https://mlsanity.com/paper/2610.19946/) (2026)

## Abstract (arXiv)

Generative modeling of discrete data, such as graphs, underpins many scientific and industrial applications, including molecular discovery and materials design. In these domains, probabilistic inference is particularly valuable, as it enables composable generation and principled incorporation of desired constraints, such as structural or functional properties. Energy-based models naturally support this goal by capturing relative likelihoods and enabling composable inference by directly enforcing constraints during inference. 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, resulting in a fidelity gap compared to discrete diffusion models. To address this gap, we introduce Graph Energy Matching (GEM), a discrete generative framework inspired by the Jordan-Kinderlehrer-Otto (JKO) transport-map optimization perspective. GEM learns a permutation-invariant potential energy that simultaneously guides discrete transport from noise toward high-likelihood graph regions and refines samples within these regions. We further introduce a sampling protocol leveraging an energy-based switching strategy, seamlessly bridging rapid, gradient-guided transport and a local mixing regime for effective exploration. On molecular graph benchmarks, GEM matches or surpasses strong discrete diffusion baselines on most reported metrics. Beyond improving generation quality, GEM's relative likelihood modeling enables targeted exploration, facilitating compositional generation, property-constrained sampling, and interpolation between graphs. Project page: https://michalbalcerak.ai/graph-energy-matching/.