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

Michal Balcerak; Suprosanna Shit; Chinmay Prabhakar; Sebastian Kaltenbach; Michael S. Albergo; Yilun Du; Bjoern Menze

cs.LGMar 24, 2026

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

Authors: Michal Balcerak, Suprosanna Shit, Chinmay Prabhakar, Sebastian Kaltenbach, Michael S. Albergo, Yilun Du, Bjoern Menze

Organizations: University of Zurich · Harvard University · Kempner Institute

Abstract

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/.

Figures & tables

Panel of Figure 1 — Figure 1: GEM Sampling Overview. Two perspectives on GEM sampling: a probability-distribution view (top) of the two-phase MCMC process, and a samples view (bottom) showing molecular trajectories from MOSES. Sampling alternates between a transport phase, where gradient-informed, greedy proposals rapidly move samples toward regions of high probability, and a mixing phase employing MH acceptance to ensure correct stationary distribution and efficient mixing between modes. Color key: transport ( orange ), MH ( blue ).

Panel of Figure 2 — Figure 2: Proposal Scoring. Local edits are scored by $q_{\mathrm{mixing}}(x\to y)$ ; dot area encodes proposal probability. Gradient-aligned, shorter jumps are favored, and stay proposals are resampled.

Panel of Figure 3 — Figure 3: Energy and Sampling Trajectories. Left: energy evolution for noise- vs. data-initialized chains, with mean $\pm$ 1 std and reference energies $\mathbb{E}_{x\sim\pi_{\mathrm{data}}}[V_{\theta}(x)]$ and $\mathbb{E}_{x\sim\pi_{0}}[V_{\theta}(x)]$ . Noise-initialized chains use greedy proposals to reach the data distribution, with these transport-aligned steps serving as an effective burn-in for the subsequent mixing chain; data-initialized chains use temperature annealing (low initial $\beta_{\mathrm{mh}}$ , then increasing) to recover novelty. Right: validity and novelty trajectories for data-initialized ( teal ) and noise-initialized ( red ) chains. Uniqueness $\approx$ 100%.

Table 1: Unconditional Generation. Results by initialization for (a) MOSES and (b) QM9. Both report V.U.N. and FCD ; higher V.U.N. and lower FCD are better. ∗ Matched DeFoG backbone; EBMs use a small energy head ( Appendix A ). MOSES: 25k generations, 1000 steps; QM9: 10k generations, 500 steps.

Panel of Figure 4 — Figure 4: V.U.N. and FCD vs Steps. V.U.N. (higher is better) and FCD (lower is better) versus inference steps on MOSES for GEM noise initialization (uniform) with greedy warmup proposal, GEM data initialization with annealed proposal, and DeFoG (marginal).

Figure 5: Qualitative Geodesic Paths. Representative MOSES interpolation from Molecule A to Molecule B. Columns show successive samples along each path; GEM uses the learned energy-weighted geometry, while the baseline uses only graph-edit cost.

Panel of Figure 6 — Figure 6: Validity Along Geodesics. Average chemical validity (%) vs. distance (node-matching, $c_{\mathrm{FGW}}$ ) for GEM energy-weighted versus cost-only geodesics. Cost-only ignores the data geometry and interpolates through low-validity regions.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Energy Backbone. Permutation-equivariant backbone with invariant pooling to produce a scalar energy.

Task	GEM $\lambda_{\text{prop}}$	DeFoG $\lambda_{\text{prop}}$
logP $\leq 1.5$	10	140
logS $\geq-2.25$	16	130
QED $\geq 0.9$	193	360
TPSA $\leq 50$	0.03	0.5

Appendix

Table 3: Hyper-parameters for noisy-regressor conditioning on MOSES. GEM uses uniform-noise regressors; DeFoG uses marginal-noise regressors. Conditional results based on 5k samples.

Source of variability	V.U.N. $\uparrow$	FCD $\downarrow$
Training seeds	$0.858\pm 0.013$	$1.51\pm 0.03$
Sampling seeds, fixed checkpoint	$0.855\pm 0.011$	$1.51\pm 0.03$

Appendix

Table 4: MOSES unconditional variability across five seeds.

Variant	V.U.N. $\uparrow$
Full GEM (transport + mixing)	0.856
Transport-only, same checkpoint	0.712
Mixing-only, same checkpoint	0.485
Mixing-only, retrained	0.21

Appendix

Table 5: Transport/mixing ablation on MOSES under noise initialization.

Time (s)	0	1	2	3	6	9	12	15
V.U.N. (DeFoG)	0.000	0.788	0.823	0.820	0.821	0.821	0.823	0.825
V.U.N. (Data-init)	0.000	0.234	0.467	0.692	0.885	0.887	0.887	0.891
V.U.N. (Noise-init)	0.000	0.012	0.056	0.370	0.836	0.855	0.861	0.860
FCD (DeFoG)	5.150	2.495	1.956	2.039	2.097	2.016	2.021	2.036
FCD (Data-init)	0.244	1.060	1.748	1.353	0.799	0.662	0.653	0.655
FCD (Noise-init)	26.16	19.41	16.01	12.61	4.427	2.596	1.895	1.714

Appendix

Table 6: MOSES wall-clock generation results on an NVIDIA H100, batch size 128.

Panel of Figure 8 — Figure 8: Test MAE of time/noise-conditioned property regressors versus denoising time $t$ (0 = fully noised, 1 = clean/denoised). Each point averages 10k test molecules. Solid lines are marginal-noise regressors; dashed lines are uniform-noise regressors. TPSA MAE is scaled by $1/50$ for visualization. Errors are large for noisy inputs and only become reasonable around $t=1$ .

Regressor	logS $\geq$ -2.25	QED $\geq$ 0.9	logP $\leq$ 1.5
	C.V.U.N. $\uparrow$	C.V.U.N. $\uparrow$	C.V.U.N. $\uparrow$
Clean	0.531	0.154	0.552
Noise-conditioned	0.836	0.422	0.846

Appendix

Table 7: Regressor ablations ( GEM ). Clean regressors are time-independent and trained only on valid molecules. Noise-conditioned regressors are time-dependent and trained on uniformly noised graphs. We report C.V.U.N. for each property condition on MOSES.

Generative Modeling

Generative modeling learns data distributions to create new samples, from images and text to molecules, simulations, and structured records.

See all 493 papers

Learning Graph Energies for Finite-Time Generation
Oct 14, 2026Amartya Banerjee, Xingyu Xu, Caroline Moosmüller +1Energy-Based Models Molecular Graph Generation
Efficient Graph Generation via Direct Prediction and Flow Matching
Oct 4, 2026Susie LuFlow Matching Graph Generation
MIT
Equivariant Efficient Joint Discrete and Continuous MeanFlow for Molecular Graph Generation
Apr 9, 2026Rongjian Xu, Teng Pang, Zhiqiang Dong +1Molecular Graph Generation Generative Modeling
School of Software, Shandong University, Jinan, Shandong, China · School of Software, Shandong University, Jinan, Shandong, China; Luke EI, Jinan, China

Oct 14, 2026J/K move · Enter open · S save

Learning Graph Energies for Finite-Time Generation

Amartya Banerjee, Xingyu Xu, Caroline Moosmüller, Harlin Lee

Energy-Based Models Molecular Graph Generation Markov Chain Monte Carlo

Open paper AskPDF arXiv

Energy-based models can represent distributions over discrete structures such as molecular graphs, but an energy that is accurate at equilibrium need not yield accurate samples when the sampler stops after a fixed number of steps. We introduce finite-horizon energy learning, which adapts the acceptance energy of a Metropolis--Hastings sampler to the distribution it returns after a prescribed number of steps, with initialization, proposals, and annealing schedule held fixed. For graph generation, a frozen pretrained model supplies initialization and proposals, and a learned energy governs acceptance. A training-only discriminator scores completed samples, and coupled continuations of the accepted and rejected branches of a proposal give an acceptance-gradient estimator that is unbiased for a fixed discriminator under stated conditions, without equilibrated negative samples or differentiation through discrete states. Generation uses the learned energy in the original single-trajectory sampler, without discriminator evaluations or reranking. On twelve pairwise binary targets with exactly evaluated terminal laws, learned energies have lower median terminal total variation than the equilibrium-correct energy in every tested short-horizon setting, and than its numerically tuned scalar multiple in most settings. On MOSES, a learned two-coefficient combination of pretrained Graph Energy Matching (GEM) potentials lowers FCD by 3.7--10.5% relative to scale-calibrated GEM at matched sampling steps, with validity gains at the shortest budget and a sub-point validity trade-off at higher budgets.

Abstract

Generative Modeling

Explore similar work

Learning Graph Energies for Finite-Time Generation

Efficient Graph Generation via Direct Prediction and Flow Matching

Equivariant Efficient Joint Discrete and Continuous MeanFlow for Molecular Graph Generation