PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Graph Energies for Finite-Time Generation

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

arXiv:2610.19946v1cs.LGcs.AI

TL;DR

An energy that is accurate at equilibrium need not give accurate samples when an MH sampler stops after a fixed number of steps. The paper learns an acceptance energy for the sampler’s finite-horizon output, reporting improved terminal TV on finite-state targets and lower MOSES FCD than scale-calibrated GEM.

  • Problem

    An energy that represents the target at equilibrium need not minimize sampling error after a sampler runs for a fixed number of steps.

    Source: §1 · p. 1

  • Method

    The method learns only an MH acceptance energy to optimize completed-sample discriminator rewards, while initialization, proposals, and schedule remain fixed.

    Source: §1 · pp. 1–2

  • Results

    Learned energies lower median terminal TV than equilibrium-correct energies in every tested short-horizon setting and lower MOSES FCD by 3.7–10.5% versus scale-calibrated GEM at matched steps.

    Source: §1 · p. 2, §6 · p. 8

  • Takeaways & Limitations

    The finite-state gains concern the finite-run distribution rather than correction of an inaccurate equilibrium density.

    Source: §6 · p. 8

    The molecular results do not establish superiority over same-family coefficient selection, and the original K = 250 comparison does not establish validity non-inferiority.

    Source: §6 · p. 8

Abstract

from arXiv · show

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.

1 INTRODUCTION

Finite-horizon energy learning adapts an MH sampler’s acceptance energy to its K-step output while holding initialization, proposals, and schedule fixed. The paper reports gains on finite-state targets and MOSES.

Motivation

  • An equilibrium-correct energy need not minimize sampling error when an MH generator stops after a fixed number of transitions.The terminal law depends on initialization, proposals, acceptance decisions, and the annealing schedule.Source: §1 · p. 1
  • A two-state example shows that changing acceptance can make a one-step output uniform even when the resulting equilibrium distribution is not uniform.With a flip proposal probability of 3/4, raising b’s energy by log(3/2) lowers acceptance to 2/3 and makes the one-step output uniform, while equilibrium assigns b probability 2/5.Source: §1 · p. 1

Approach

  • Finite-horizon energy learning optimizes the acceptance energy using discriminator rewards on completed K-step samples while keeping initialization, proposals, and the annealing schedule fixed.A coupled continuation of accepted and rejected branches supplies an unbiased gradient estimator for a fixed discriminator under stated conditions.Source: §1 · pp. 1–2

Finite-state results

  • At K ∈ {4, 8, 16}, learned energies achieve lower median terminal TV than equilibrium-correct energies in all 18 reference and 18 new-target settings, and beat scalar-tuned controls in 14 and 15 settings, respectively.The scalar control is a numerically tuned multiple of the equilibrium-correct energy.Source: §1 · p. 2

MOSES results

  • 3.7–10.5% lower FCD than scale-calibrated GEM at matched refinement steps, with budget-dependent validity trade-offs, is reported on MOSES.Generation uses one trajectory without discriminator evaluations or additional continuations.Source: §1 · p. 2

2 RELATED WORK

The related work spans short-run energy learning, learned sampler dynamics, estimators for discrete randomness, and reward adaptation for graph generation. This paper distinguishes its approach by fixing the pretrained sampler and learning only its acceptance energy.

Training energies with finite-run samplers

  • Short-run EBMs treat a fixed number of MCMC steps as part of the generator, while contrastive divergence and GEM learn energy from short-run negatives held fixed during differentiation (Balcerak et al., 2026).Here initialization and proposals are fixed, and the objective is defined on the K-step output.Source: §2 · p. 2

Adapting proposals and transition dynamics

  • ALOE, EB-GFN, A-NICE-MC, reinforcement-learning methods, and DeFoG adapt proposals or sampling dynamics, whereas this work keeps the proposal kernel and schedule fixed.The cited approaches include joint energy-and-sampler learning, adversarial proposal training, kernel adaptation, and separating denoising from sampling dynamics.Source: §2 · p. 3

Differentiating through discrete randomness

  • Coupled stochastic derivatives handle discrete random choices, and their MH adaptations couple accepted and rejected continuations (Arya et al., 2022; Arya et al., 2023, 2024).This work applies the estimator to terminal discriminator rewards, subsamples decisions with inverse inclusion probabilities, and compares it with a score-function estimator.Source: §2 · p. 3

Graph generation and reward adaptation

  • Graph reward-adaptation methods update denoisers or flow networks, or apply rewards during sampling, while this approach keeps pretrained networks fixed and changes only the acceptance energy.The cited approaches include reward fine-tuning (Zhu et al., 2026) and sampling-time reward or conditioning (Phunyaphibarn and Sung, 2026).Source: §2 · p. 3

3 FINITE-HORIZON ENERGY LEARNING

Finite-horizon energy learning defines its objective on an MH sampler’s terminal law, scores completed samples with a fixed discriminator, and estimates acceptance gradients using paired continuations. Training alternates energy updates and discriminator refitting; generation uses the learned energy in the original sampler.

3.1 Sampler and Terminal Law

  • The sampler’s K-step terminal law is formed from an initialization distribution and step-specific MH kernels, with the learned parameters entering only through acceptance energy.Initialization, proposals, and the inverse-temperature schedule remain fixed; annealed kernels need not share a common invariant distribution.Source: §3.1 · p. 3

3.2 Fixed-Discriminator Objective

  • The objective maximizes expected discriminator reward on completed samples, with discriminator parameters held fixed during each energy update and refitted between updates.Unlike equilibrium likelihood objectives, it scores the returned samples; the optimal-classifier divergence interpretation requires stated conditions.Source: §3.2 · p. 3

3.3 Gradient Estimation with Paired Continuations

  • The gradient weights each local acceptance-probability derivative by the difference in expected terminal reward between accepting and rejecting; paired continuations estimate that difference.Both branches use shared future random numbers, and only the branch that differs from the base trajectory needs simulation.Source: §3.3 · pp. 3–4

3.4 Training and Generation

  • Training alternates energy updates using selected decisions from simulated chains with discriminator refitting; generation replaces the original acceptance energy in the same K-step sampler.At generation, discriminator evaluations and paired continuations are absent; per-step cost can still depend on the energy.Source: §3.4 · p. 4

4 EXACT FINITE-STATE STUDY

The finite-state study tests whether adapting only the acceptance energy can improve short-run sampling when initialization, proposals, and schedule stay fixed. It compares learned energies with an equilibrium-correct energy, a scalar-tuned control, and training and horizon variants.

4.1 Targets, Sampler, and Comparisons

  • The study uses small pairwise binary targets with exact terminal laws and holds initialization, proposals, and the horizon-scaled annealing schedule fixed across compared energies.For L = 6 or 8, the sampler starts from target one-variable marginals and proposes a uniformly selected spin flip.Source: §4 · p. 4, §4.1 · pp. 4–5

4.2 Comparison with the Equilibrium-Correct Energy

  • At K = 4, 8, and 16, learned energies have lower median TV than the equilibrium-correct energy in all 18 reference target–horizon settings.The same pattern holds in all 18 new-target settings; at K = 4, 8, and 16, new-target median TV falls from 0.553 to 0.518, 0.413 to 0.307, and 0.240 to 0.069.Source: Figure 2 · p. 5, §4.2 · p. 5

4.3 Comparison with Scalar Tempering

  • Learned energies beat scalar-tuned βH in 14 of 18 reference settings and 15 of 18 new-target settings, with all three new-target exceptions on New-8-2.The scalar-tuned control is an oracle using the target and exact terminal probabilities; its median optimal multiplier on reference targets falls from 9.04 at K = 4 to 2.00 at K = 8 and 1.22 at K = 16.Source: §4.3 · pp. 5–6

4.4 Training Budget, Estimators, and Long Horizons

  • Median recovery over the restricted energy rises from 52.4% at N = 32 to 90.4% at N = 512 with the fitted discriminator, reaching 96.8% with the oracle discriminator.At N = 512, the fitted discriminator leaves a median TV gap of 0.024 to the numerical reference.Source: §4.4 · p. 6
  • Beyond K = 64, learned energies are less accurate than the equilibrium-correct energy, although the numerical reference remains near zero TV.Because the learned family contains both energies, the shortfall is attributed to sampled training rather than family capacity; the study does not distinguish candidate causes.Source: §4.4 · p. 6

5 ADAPTING PRETRAINED GEM POTENTIALS ON MOSES

The MOSES study adapts only GEM’s acceptance energy while holding initialization, proposals, schedule, and refinement-step budgets fixed. The learned energy lowers FCD at all three tested budgets, with validity trade-offs and controls that limit what the comparison establishes.

5.1 Two-Coefficient Energy and Controlled Comparison

  • Each budget uses its own calibrated baseline and trained coefficients, and the comparison measures FCD, validity, and VUN on matched refinement steps.FCD uses ChemNet feature means and covariances; validity counts outputs that decode and sanitize, and VUN counts valid, unique outputs absent from training.Source: §5.1 · p. 6, §5.2 · p. 7
  • The learned acceptance energy combines frozen GEM potentials using scale coefficient c and difference weight ω, while a training-only discriminator supplies rewards.At K = 500, discriminator screening failed and no energy was trained, so that budget is missing rather than an evaluated failure.Source: §5.1 · p. 6, §5.4 · p. 8

5.3 Results

  • FCD fell 5.7% (2.508 to 2.366), 10.5% (2.044 to 1.829), and 3.7% (1.561 to 1.504) at K = 125, 250, and 1000; the joint criterion passed at 125 and 1000 only.At K = 250, the validity lower confidence bound was −1.25 points, so validity non-inferiority was not established; this was the original primary comparison.Source: Table 2 · p. 7, §5.3 · p. 7, Figure 3 · p. 8
  • At K = 125, validity and VUN rose by 5.42 and 5.29 points, while at K = 250 they fell by 0.78 and 0.32 points; at K = 1000 they changed by −0.02 and +0.15 points.The K = 125 yield gains persisted under strict scoring, with validity +4.77 and VUN +4.65 points; these checks reuse the generated samples.Source: §5.3 · p. 7, Figure 3 · p. 8

5.4 Same-Family Controls and Protocol Conditions

  • The baseline calibrates one coefficient while learning adjusts two, so the comparison cannot isolate gradient-based learning from the added coefficient.Static two-coefficient controls do not establish an FCD advantage over coefficient selection or finite-difference training; no corresponding three-seed static evaluation was reported at K = 1000.Source: §5.4 · p. 8

6 DISCUSSION AND CONCLUSION

The paper adapts an MH sampler’s acceptance energy to improve its finite-run distribution without changing initialization, proposals, or schedule. Its finite-state and MOSES results show short-horizon gains, while the molecular evidence leaves important comparisons and longer-horizon training open.

6 DISCUSSION AND CONCLUSION

  • Learned energies improve median terminal TV over the equilibrium-correct energy in every tested short-horizon setting and over its numerically tuned scalar multiple in most settings.The MOSES energy combines two frozen GEM potentials and lowers FCD by 3.7–10.5% against scale-calibrated GEM at matched refinement steps.Source: §6 · p. 8
  • The MOSES results do not establish superiority over same-family coefficient selection, and the original K = 250 comparison does not establish validity non-inferiority.Accordingly, lower FCD should not be read as uniformly improved sample quality or reduced computational cost.Source: §6 · p. 8
  • The procedure trains a separate energy for each horizon, and sampled training falls behind the correct energy at long horizons in the finite-state study.Training across budgets and understanding the long-horizon shortfall remain future-work directions.Source: §6 · p. 8

AI use statement

  • The authors report using generative AI tools to assist with drafting and text revision, while taking responsibility for the final manuscript and its claims.Source: AI use statement · p. 9

A GRADIENT DERIVATION AND IMPLEMENTATION

The estimator differentiates local MH acceptance probabilities while accounting for their downstream effects through paired terminal-reward continuations. Molecular optimization adapts only the acceptance energy, and the alternative-estimator comparison finds no general efficiency advantage for coupling.

A.1 Gradient Identity and Estimator

  • Paired continuations estimate the downstream value difference between accepting and rejecting a proposal, while only the local acceptance probability is differentiated.The base chain can supply one branch when decision selection is independent of its future randomness; shared random numbers change the coupling, not the expected reward difference.Source: §A.1 · p. 11
  • The subsampled estimator is unbiased when each nonzero acceptance-gradient term has positive selection probability and selected terms are averaged over all base chains.Dividing by the number of selected decisions instead introduces bias; sampled states, inclusion weights, and discriminator rewards are held constant during differentiation.Source: §A.1 · p. 11

A.2 Molecular Optimization

  • Molecular optimization keeps both pretrained potentials frozen and trains two acceptance-energy coordinates, with constrained reparameterized updates and a training-only discriminator.The setup uses AdamW, clips the energy coordinates, and refits the discriminator during training; the discriminator is not used at generation.Source: §A.2 · p. 12, Table 3

A.3 Alternative Estimators

  • With matched transition counts, both coupled and leave-one-out score-function learning improve on the correct energy in all 18 settings, but coupling is not generally more efficient.Coupling has lower median TV in 12 settings by 0.0015 and takes 4.66 times the CPU runtime; the fixed-gradient diagnostic finds no general efficiency advantage over alternatives.Source: §A.3 · p. 13

A.4 Computational Resources and Continuation Work

  • Training continuations add simulated transitions beyond the base chains, while deployment needs neither continuations nor discriminator evaluations.With an expected F selected decisions per chain, the stated upper bound is NK + NF(K − 1) transitions; these counts do not measure wall-clock ratios or bound mixing time.Source: §A.4 · p. 13

B.1 Targets and Optimization

The finite-state experiments use independently drawn pairwise binary targets with exactly enumerable laws, then compare a learned energy against an equilibrium-correct energy and a scalar-calibrated control. Optimization uses a richer feature family and a fitted discriminator, while the numerical reference is not certified globally optimal.

B.1 Targets and Optimization

  • The finite-state study evaluates synthetic pairwise binary targets with 6 or 8 variables, using product-marginal initialization, single-spin-flip proposals, and an inverse-temperature ramp.Reference and new-target sets each contain three independently drawn targets of each size, and enumeration gives target probabilities and transition laws.Source: §B.1 · p. 13
  • The learned family includes all unary and pairwise features, while the scalar control tunes a restricted unary-and-ring energy to minimize terminal TV.Population pseudo-likelihood recovers the correct energy up to a constant; a linear discriminator is initialized by an exact population logistic fit and then updated from generated and target samples.Source: §B.1 · p. 13
  • The numerical reference uses exact gradients and bounded optimization but is not certified globally optimal; eight of 42 configurations miss the projected-gradient tolerance, mainly at acceptance kinks.The reference minimizes a discriminator-related objective rather than TV, so its TV recovery fraction is not a known maximum TV gain.Source: §B.1 · p. 14

B.2 Repeated Training and Scalar Controls

Across fresh-seed replications and independently drawn targets, learned energies usually outperform a numerically tuned scalar multiple of the equilibrium-correct energy. The two-state example illustrates why finite-step accuracy and equilibrium accuracy can diverge.

B.2 Repeated Training and Scalar Controls

  • The fresh-seed replication improves all 18 reference settings, while new-target training uses independently drawn targets and transfers no trained energy.The win counts are descriptive because horizons within a target are dependent; Table 4 reports medians and sample standard deviations across three training seeds.Source: §B.2 · p. 14, Table 4 · p. 15
  • Learned medians beat tuned βH in 14 of 18 reference settings and 15 of 18 new-target settings; all three new-target exceptions are on New-8-2.The scalar control uses exact terminal TV; its search stabilizes across all 36 settings, but tail control does not establish global optimality within the search interval.Source: §B.2 · p. 14

B.3 Finite-Time and Equilibrium Distributions

  • In the two-state example, raising b's energy makes the one-step output uniform, although the equilibrium probability of b becomes 2/5.With the correct energy, the one-step probability of b is 3/4; the adjusted energy lowers flip acceptance to 2/3.Source: §B.3 · p. 14, Figure 4 · p. 14

C.1 Sampler, Calibration, and Metrics

The molecular sampler retains the released proposal procedure and acceptance-temperature schedule, while GEM calibration and evaluation use distinct development and final reference sets. Metrics follow fixed validity conventions and are interpreted within the graph encoding.

C.1 Sampler, Calibration, and Metrics

  • Molecular chains use fixed node counts drawn from the MOSES training histogram, greedy initialization, released proposal settings, and an acceptance-temperature ramp.The ramp runs from approximately 0.78 to 17.62 over 500 steps, then holds; the supplement's molecular example does not reconstruct the original development, selection, and final evaluation pipeline.Source: §C.1 · p. 16
  • Exact MH marginalizes forward and reverse proposal probabilities over the same retry stages, whereas GEM's released rule uses a realized forward retry path and a reverse marginal.The released rule is not exact MH for the same target; both rules have parameter-independent proposal terms.Source: §C.1 · p. 16
  • Development FCD uses 16,010 training molecules, while final FCD uses a separate 176,225-molecule scaffold-test reference.The development molecules were not held out from released GEM checkpoint pretraining; graph scores are interpreted within the shared encoding and evaluator because the encoding cannot represent every aromatic-hydrogen configuration.Source: §C.1 · p. 16

C.2 Released-Rule Deployment and Uncertainty

Under the prespecified joint rule, released-rule deployment meets the criteria at K = 250 and 1000, but not at K = 125; the comparisons use few seed-level observations.

C.2 Released-Rule Deployment and Uncertainty

  • Table 5 reports FCD, validity, and VUN for 25,000 requested graphs per seed, with means ± sample SD over two prespecified seeds at K = 125 and three elsewhere.Baselines are uncalibrated released GEM at K = 125 and 1000 and calibrated GEM at K = 250; bold marks the better mean, not statistical significance.Source: Table 5 · p. 16
  • Released-rule deployment meets the joint rule at K = 250 and 1000, but not at K = 125, where validity and VUN lower limits are −2.52 and −2.80 points.The rule requires a one-sided paired FCD test at 0.05 and one-sided 95% lower confidence limits above −1 point for validity and VUN.Source: §C.2 · p. 17
  • Seed-level differences, not individual graphs, are the observations for the small-sample t tests; energies are reused across acceptance rules and some discriminator data are shared.Training seed k is paired with generation seed 1911 + k.Source: §C.2 · p. 17
  • One released-GEM K = 1000 comparator seed mixes 63 H100 and 35 L40-class batches, while the two H100-only seeds leave the FCD comparison unresolved (p = 0.11).Runtime estimates exclude packed batches.Source: §C.2 · p. 17

C.3 Protocol Changes and Sensitivity

Sensitivity checks preserve the exact-rule K = 125 improvement, while protocol choices and limited GEM comparisons constrain how broadly the findings can be interpreted.

C.3 Protocol Changes and Sensitivity

  • Protocols were recorded internally, and the six comparisons are not independent prospective replications; no final test score selected the deployed iterate.The original K = 250 comparison followed development results and exploratory pilots, while later budgets followed earlier tests.Source: §C.3 · p. 17
  • At K = 125, the discriminator-pool eligibility threshold rose from 1% to 2% before pool output or training after 22 of 1,833 native-valid graphs mismatched encoding.The completed pool had 368 mismatches among 22,270 valid graphs (1.65%); the change affected eligibility, not evaluation or sample handling, and the original rule would have stopped the experiment.Source: §C.3 · p. 17
  • Strict scoring preserves the exact-rule K = 125 improvement: validity and VUN gains are 4.77 and 4.65 points, while the released-rule joint failure remains.About 0.64 points of native validity gain arise from disconnected output; matching valid-input counts changes the FCD difference from −0.1419 to −0.1379, and fixed subsample comparisons remain negative.Source: §C.3 · p. 17
  • At K = 500, the discriminator screen failed at correlation approximately −0.071, so no energy was trained and no second attempt was made; this is not an evaluated negative generator result.At K = 1000, a revised negative pool containing guided samples passed at the correlation boundary of −0.50.Source: §C.3 · p. 17

C.4 Scope of the GEM Comparison

  • At K = 125, coefficient-selection control FCD is 2.280 versus 2.366 for learning, but validity is 91.10% versus 92.55%; the three-seed FCD difference is unresolved.The K = 250 control FCD values are 1.811 versus 1.829 with exact acceptance and 1.769 versus 1.784 with released acceptance; these comparisons do not establish an FCD advantage for learning over alternatives.Source: §C.4 · p. 17
  • The K = 250 setting has the lowest development FCD among tested settings but is validity-ineligible at 94.97% against a 95.44% threshold; no three-seed static evaluation corresponds to K = 1000.The K = 125 static setting was selected among four development candidates.Source: §C.4 · p. 17
  • At K = 1000, molecular-weight distance worsens despite improved FCD, precluding a claim of uniformly improved molecular properties; planned residual-energy and random-sign placebo tests were not completed.The six-setting coefficient grid costs 3.2 H100-hours at K = 250, versus 4.86 per seed for energy updates before discriminator preparation; schedule compression is inconclusive.Source: §C.4 · p. 17

Related papers

Loading 2610.19946v1…