Energy models improve fixed-step graph generation accuracy
Learning Graph Energies for Finite-Time Generation
Machine LearningArtificial Intelligence
Summary
Generating complex structures like molecular graphs by sampling can be tricky because the sampling method might not produce good results if it stops too early. The authors introduce a way to learn energy functions that better guide the sampling when it is limited to a fixed number of steps. They use a pretrained model to suggest starting points and moves, then learn an energy to decide which moves to accept, improving the quality of the generated graphs. Testing showed this approach produces samples that are closer to the desired distribution, especially under tight sampling limits.
What this means in practice
- •For computational chemistry teams: Generate more accurate molecular graphs within fixed computational budgets to support drug discovery and materials design.
- •For machine learning engineers: Improve the quality of discrete structure sampling in finite time for applications such as synthetic graph creation or generative modeling.
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.