Learning Graph Energies for Finite-Time Generation

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

cs.LGOct 14, 2026

Learning Graph Energies for Finite-Time Generation

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

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: Finite-horizon energy learning for a pretrained GEM graph sampler. Generation (top): one $K$ -step MH trajectory whose initialization and proposals use GEM’s frozen potential $V_{\mathrm{G}}$ ; only the acceptance energy $E_{\theta}$ , with coefficients $\theta=(c,\omega)$ in ( 6 ), is learned. Training (dashed; absent at generation): at a selected decision $t$ , the accepted and rejected branches are continued to step $K$ with shared future random numbers, and the difference of their terminal rewards weights $\nabla_{\theta}\alpha_{\theta,t}$ . The displayed update is ( 5 ) with selection indicators $I_{it}$ and inclusion probabilities $p_{it}$ ; $N$ is the number of base chains.

Figure 2: Finite-time sampling on the reference targets (original runs). (a) Exact terminal total variation (TV) against the number of steps

K

; each learned point is an energy trained separately for that horizon with

N=512

trajectories per update. Curves are medians over six targets after aggregating training seeds. (b) Recovery (%) of the numerical reference’s improvement over the restricted energy, for learning with a fitted or an oracle discriminator, against trajectories per update

N

; 100% is not an upper bound. (c) TV-minimizing multiplier

\beta^{*}

E=\beta H

under the shared annealing schedule, for individual targets (thin) and their median (thick); the dashed line marks

\beta=1

. Shading shows interquartile ranges, not confidence intervals. Horizontal axes are logarithmic, as is the vertical axis in (c).

Targets	$K$	Correct $H$	Learned (ours)
Reference	4	$0.329$	$\mathbf{0.197}$
(original runs)	8	$0.201$	$\mathbf{0.063}$
	16	$0.089$	$\mathbf{0.030}$
New targets	4	$0.553$	$\mathbf{0.518}$
	8	$0.413$	$\mathbf{0.307}$
	16	$0.240$	$\mathbf{0.069}$

Table 1: Exact terminal TV at short horizons (lower is better). Entries are medians over six targets of each target’s median over three training seeds. Reference rows summarize the original runs in Figure 2 a; Table 4 reports per-target results for a fresh-seed replication on the reference targets and for the new targets. Bold marks the smaller median.

$K$	Method	FCD $\downarrow$	Validity $\uparrow$	VUN $\uparrow$
125	GEM (calibrated)	$2.508\pm 0.026$	$87.13\pm 0.14$	$85.15\pm 0.09$
	GEM + learned energy (Ours)	$\mathbf{2.366\pm 0.063}$	$\mathbf{92.55\pm 0.24}$	$\mathbf{90.44\pm 0.25}$
250	GEM (calibrated)	$2.044\pm 0.030$	$\mathbf{95.43\pm 0.10}$	$\mathbf{92.76\pm 0.05}$
	GEM + learned energy (Ours)	$\mathbf{1.829\pm 0.021}$	$94.65\pm 0.33$	$92.45\pm 0.29$
1000	GEM (calibrated)	$1.561\pm 0.018$	$\mathbf{97.84\pm 0.14}$	$93.77\pm 0.12$
	GEM + learned energy (Ours)	$\mathbf{1.504\pm 0.018}$	$97.82\pm 0.03$	$\mathbf{93.92\pm 0.09}$

Table 2: GEM with and without finite-horizon energy learning at matched refinement steps

K

, with exact MH acceptance. Mean

\pm

sample SD over three paired seeds, each with 25,000 requested graphs; validity and VUN are % of requested graphs. Bold marks the better mean in each pair, not significance.

Figure 3: FCD–validity changes at matched budgets. Circles: calibrated GEM; squares: GEM with the learned energy (both exact MH); arrows join matched-budget pairs. FCD decreases at every budget, whereas validity rises at

K=125

, falls at

K=250

, and is nearly unchanged at

K=1000

. Bars show one sample SD over three seeds; each budget has its own calibrated baseline and trained coefficients.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Setting
Energy updates	80 updates; 256 base chains per update; approximately four selected decisions per chain on average; final iterate.
Decision selection	Inclusion probability $\min(1,\lambda_{j}h_{it})$ , with $h_{it}$ an acceptance-gradient norm in rescaled coordinates and guidance weight three. $\lambda_{0}=0.0501$ ; later scales $\lambda_{j}$ use the preceding update’s sensitivity sum.
Optimizer	AdamW, learning rate 0.02, five-update warm-up then linear decay to zero, moments $(0.9,0.999)$ , no weight decay.
Coordinates	AdamW on $(u^{\prime},w)$ with $u^{\prime}=\log(e^{u}+0.57w)$ ; gradient components divided by $(0.676,0.152)$ before clipping to norm one. After each update, clip $u\in[-1.5,1.5]$ and $w\in[-1,3]$ , then raise $w$ where necessary so that $e^{u}+0.57w\geq e^{-1.5}$ .
Discriminator	Rank-64 LoRA adaptation of $V_{\mathrm{G}}$ on all 12 layers and the scalar head ( Hu et al.,, 2022 ) ; approximately 9.7 million trainable parameters; logit $-\beta_{\mathrm{ref}}(V_{\psi}-V_{\mathrm{G}})$ with fixed $\beta_{\mathrm{ref}}\approx 17.62$ .
Initial discriminator fit	25,600 generated examples; 1,000 logistic updates with 128 generated and 128 node-count-matched data graphs per update. AdamW learning rate $2\times 10^{-5}$ , 200-step warm-up, no weight decay, clipping to one.

Appendix

Table 3: Molecular training configuration. Both pretrained potentials are frozen; only the two energy coordinates parameterize the acceptance energy. The discriminator is used only during training.

Reference targets: fresh-seed replication
Target ID	$K$	Correct $H$	Tuned $\beta H$	Learned median	Seed SD
Ref-6-1	4	$0.1266$	$0.0613$	$\mathbf{0.0529}$	$0.005$
Ref-6-1	8	$0.0845$	$0.0284$	$\mathbf{0.0275}$	$0.004$
Ref-6-1	16	$0.0593$	$0.0279$	$\mathbf{0.0208}$	$0.004$
Ref-6-2	4	$0.3435$	$0.2319$	$\mathbf{0.2284}$	$0.001$
Ref-6-2	8	$0.2389$	$0.1247$	$\mathbf{0.0805}$	$0.009$

Appendix

Table 4: Exact terminal total variation by target distribution and sampling horizon; lower is better. Each target is a synthetic pairwise binary model with random fields and couplings (Appendix B.1 ); no external dataset is used. Target IDs specify the set ( Ref or New ), number of spins (6 or 8), and instance index (1–3). Each instance has fixed fields and couplings; its three rows evaluate the same target at horizons

K=4,8,16

, the number of MH transitions. The controls are the equilibrium-correct energy

H

and its numerically tuned scalar multiple

\beta H

; because

H

is exact at equilibrium, lower TV than

H

reflects a finite-horizon improvement. Learned columns report the median and sample standard deviation of terminal TV across three training seeds; TV itself is evaluated exactly, so the SD reflects only variation across training seeds. The reference block repeats training with fresh seeds on the original targets, whereas the new-target block trains from scratch on independently drawn target distributions to test whether the gains recur; no trained energy is transferred. Table 1 and Figure 2 use the original reference-target runs. All learned runs use

N=512

and 2,000 updates. Bold identifies the lowest of the three TV values using unrounded results.

Panel of Figure 4 — Figure 4: Finite-time accuracy need not imply equilibrium accuracy. One-step output and equilibrium probability for the two-state example, computed exactly.

$K$	Method	FCD $\downarrow$	Validity $\uparrow$	VUN $\uparrow$
125	GEM (released)	$2.406\pm 0.0003$	$\mathbf{93.25\pm 0.05}$	$\mathbf{90.99\pm 0.06}$
	GEM + learned energy (Ours)	$\mathbf{2.295\pm 0.011}$	$92.45\pm 0.44$	$90.43\pm 0.44$
250	GEM (calibrated)	$1.920\pm 0.012$	$93.48\pm 0.09$	$91.20\pm 0.10$
	GEM + learned energy (Ours)	$\mathbf{1.784\pm 0.033}$	$\mathbf{95.02\pm 0.42}$	$\mathbf{92.97\pm 0.41}$
1000	GEM (released)	$1.529\pm 0.013$	$\mathbf{98.02\pm 0.03}$	$94.02\pm 0.22$
	GEM + learned energy (Ours)	$\mathbf{1.487\pm 0.022}$	$97.95\pm 0.08$	$\mathbf{94.21\pm 0.11}$

Appendix

Table 5: Deployment with GEM’s released acceptance rule. “GEM + learned energy (Ours)” reuses the energies from Table 2 without retraining, with released-rule acceptance replacing exact MH. Metrics and denominators match Table 2 : FCD, validity (%), and VUN (%), with 25,000 requested graphs per seed. Entries are means

\pm

sample SD over two prespecified seeds at

K=125

and three elsewhere. Baselines are uncalibrated released GEM at

K=125,1000

and calibrated GEM at

K=250

. Bold marks the better mean within each pair, not statistical significance.

Energy-Based Models

Energy-based models assign scalar energies to configurations, defining probabilities and guiding generation, inference, prediction, and associative memory.

See all 91 papers

Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation
Mar 24, 2026Michal Balcerak, Suprosanna Shit, Chinmay Prabhakar +4Generative Modeling Energy-Based Models
University of Zurich · Harvard University · Kempner Institute
Generating Physically Consistent Molecules with Energy-Based Models
May 18, 2026Christoph Griesbacher, Lea Bogensperger, Andreas Habring +1Energy-Based Models 3D Molecular Generation
Graz University of Technology · University of Zurich
Energy Generative Modeling: A Lyapunov-based Energy Matching Perspective
May 7, 2026Yixuan Wang, Wenqian Xue, Warren E. DixonEnergy-Based Models
Department of Mechanical and Aerospace Engineering University of Florida

Mar 24, 2026J/K move · Enter open · S save

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

Michal Balcerak, Suprosanna Shit, Chinmay Prabhakar, Sebastian Kaltenbach+3

University of Zurich · Harvard University · Kempner Institute

Generative Modeling Energy-Based Models Molecular Graph Generation Constrained Generative Modeling

Open paper AskPDF 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/.

Abstract

Energy-Based Models

Explore similar work

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

Generating Physically Consistent Molecules with Energy-Based Models

Energy Generative Modeling: A Lyapunov-based Energy Matching Perspective