Generating Physically Consistent Molecules with Energy-Based Models

TL;DR

EBMol introduces an energy-based model for 3D molecular generation that leverages physical energy landscapes, achieving state-of-the-art results and enabling controllable molecule design without retraining.

Contribution

The paper presents EBMol, the first energy-based model for 3D molecules that incorporates physical energy inductive bias and achieves superior generation quality.

Findings

Achieves state-of-the-art performance on QM9 and GEOM-Drugs datasets.

Provides a principled energy landscape as a quality metric for molecules.

Enables controllable molecule generation via shape-steered sampling.

Abstract

Molecules in equilibrium follow a Boltzmann distribution, making the underlying energy landscape a physically grounded modeling objective. However, such landscapes are difficult to learn from data and, once learned, hard to sample from. Diffusion and flow-matching models sidestep these difficulties by learning a time-conditional score or transport field between noise and data, losing the energy inductive bias in exchange for a more tractable training objective. We introduce EBMol, an energy-based model (EBM) that restores this inductive bias by learning an atom-additive scalar potential without explicit simulation during training. Our method employs a flow-inspired Restoring Field Matching objective to approximate the energy landscape. We adopt the Mirror-Langevin algorithm for sampling, enabling unified updates of atomic positions and types, and incorporate parallel tempering for inference-time compute scaling. EBMol is the first EBM for 3D molecular generation to achieve state-of-the-art performance on QM9 and GEOM-Drugs. Moreover, we show that the learned energy landscape serves as a principled quality metric for ranking and filtering configurations, and demonstrate controllable generation without retraining through shape-steered sampling via potential composition and zero-shot linker design.