🤖 machine learning

Generating Physically Consistent Molecules with Energy-Based Models

The paper introduces EBMol, a novel energy-based model that restores physical inductive bias for 3D molecular generation by learning an atom-additive scalar potential via a flow-inspired objective, achieving state-of-the-art performance on QM9 and GEOM-Drugs while enabling principled quality ranking and controllable zero-shot design.

Original authors: Christoph Griesbacher, Lea Bogensperger, Andreas Habring, Thomas Pock

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Christoph Griesbacher, Lea Bogensperger, Andreas Habring, Thomas Pock

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a computer to design new molecules, like tiny Lego structures made of atoms. The goal is to create molecules that are stable and physically possible, not just random piles of blocks that would fall apart the moment you touched them.

Most current AI methods for this task are like a GPS navigation system. They learn a route from "chaos" (random noise) to "destination" (a valid molecule). They are very good at getting you there, but once you arrive, they don't really understand why the destination is a good place to be. They can't easily tell you if a slightly different version of the molecule is better or worse, and they can't easily combine parts of two different molecules to make a new one without retraining the whole system.

EBMol is a new approach that changes the game. Instead of learning a route, it learns a landscape.

The Energy Landscape: A Mountain Range of Molecules

Think of all possible arrangements of atoms as a giant, 3D mountain range.

Valleys represent stable, real molecules. The deeper the valley, the more stable the molecule.
Mountains represent unstable, impossible, or "broken" structures.

In physics, molecules naturally settle into the deepest valleys (this is called the "Boltzmann distribution"). The problem is that learning to draw this mountain map from data is incredibly hard. Previous attempts either gave up on drawing the map entirely (using the GPS method mentioned above) or tried to draw it but got stuck because the map was too jagged and confusing to navigate.

How EBMol Draws the Map

The authors introduce a new training method called Restoring Field Matching (RFM).

Imagine you have a smooth rubber sheet representing the energy landscape. You place a heavy ball (a real molecule) on the sheet. You want the sheet to dip down exactly under the ball, creating a perfect valley.

Old methods tried to guess the shape of the dip by looking at the ball from a distance, often getting the shape wrong.
EBMol's method is like having a magical wind. It blows from both sides of the ball toward the center. If you are to the left of the molecule, the wind pushes you right. If you are to the right, it pushes you left. If you are above or below, it pushes you down.

By training the AI to create this "restoring wind" that always points back to real molecules, the AI automatically learns to carve out perfect valleys (local minima) where the real data sits. It does this without needing to run expensive physics simulations during the training phase.

Navigating the Map: The Mirror Walker

Once the map is drawn, you need a way to walk through it to find new molecules. Standard walking algorithms (like Langevin dynamics) work fine on flat ground (coordinates), but they struggle when you hit a wall (like the rules that say an atom's probability must add up to 100%).

EBMol uses a clever trick called the Mirror Langevin Algorithm.

Imagine you are walking in a room with curved walls. If you try to walk straight, you might hit a wall.
The "Mirror" method is like putting on special glasses that make the curved room look flat to you. You walk in a straight line in this "mirror world," and when you take off the glasses, you realize you've actually walked along the curve of the wall perfectly.
This allows the AI to update both the position of atoms (where they are) and their type (whether they are Carbon, Oxygen, etc.) simultaneously and correctly, without breaking the rules of chemistry.

Why This Matters: The "Superpower" of the Map

Because EBMol actually learns the energy map (the "why" behind the molecule), it gains superpowers that the GPS-style models don't have:

Ranking and Filtering: You can generate 1,000 molecules and instantly score them. The ones with the lowest energy (deepest valleys) are the best. It's like having a quality control inspector that knows exactly which molecules are physically sound.
Mix and Match (Composable Generation): Because the energy is just a sum of parts, you can take the "energy map" of a specific shape (like a sphere) and combine it with the map of a drug molecule. The AI then walks the path that satisfies both conditions. You can ask for a "spherical drug" and it will generate one, even if it never saw a spherical drug in its training data.
Zero-Shot Linker Design: Imagine you have two halves of a molecule (like two puzzle pieces) and you need to design the middle piece to connect them. EBMol can fill in the missing middle by "walking" the energy landscape to find a stable connection, without needing to be retrained for that specific puzzle.

The Results

The paper shows that EBMol creates molecules that are not only valid but are physically closer to the "perfect" stable state than any previous method that didn't explicitly model chemical bonds. It achieved top scores on standard tests (QM9 and GEOM-Drugs) and, crucially, its generated molecules require much less "relaxation" (physical tweaking) to become stable.

In short, EBMol stops treating molecule generation as a "guess the destination" game and turns it into a "navigate a physical landscape" game, giving the AI a deeper, more intuitive understanding of chemistry.

Technical Summary: Generating Physically Consistent Molecules with Energy-Based Models

Problem Statement

Molecular states at thermal equilibrium are governed by the Boltzmann distribution, $p(x) \propto \exp(-E(x)/(k_BT))$ , where the probability of a configuration is determined by its potential energy. While standard datasets for molecular generation (e.g., QM9, GEOM-Drugs) contain geometry-optimized structures relaxed to local energy minima, they are not drawn from the Boltzmann distribution; their relative frequencies reflect enumeration strategies rather than thermodynamic weights. Consequently, a scalar energy landscape where minima coincide with stable configurations is the natural modeling target.

However, learning and sampling from such landscapes remain open challenges. Traditional Energy-Based Models (EBMs) require Boltzmann-distributed data for training or rely on expensive Markov Chain Monte Carlo (MCMC) methods that struggle to traverse the rugged, unphysical regions between noise and valid structures. Conversely, state-of-the-art diffusion and flow-matching models sidestep these issues by learning time-conditional transport fields. While effective, these models lose the energy inductive bias, providing no explicit energy function to evaluate, rank, or compose configurations at inference time.

Methodology: EBMol

The authors introduce EBMol, an Energy-Based Model for 3D molecular generation that restores the energy inductive bias without requiring explicit simulation during training. The framework consists of two primary components:

1. Restoring Field Matching (RFM) for Training

To learn a tractable energy landscape $E_\theta(x)$ without Boltzmann-distributed data, the authors propose Restoring Field Matching (RFM).

Concept: Instead of treating data points as trajectory endpoints (as in standard flow matching), RFM extends the interpolation interval to $t \in [-1, 1]$ , passing through the data at $t=0$ .
Target Field Construction: The model is trained to match the negative gradient of the energy, $-\nabla E_\theta(x_t)$ , to a target "restoring field" $u$ . This field points toward the data $x_0$ from both sides of the manifold (interpolation $t \in [0, 1]$ and extrapolation $t \in [-1, 0]$ ).
Stationarity: A smoothing function $b(|t|)$ scales the target field such that it vanishes at $t=0$ (the data point), ensuring data points become local minima of the learned energy landscape.
Regularization: A per-atom regularization term anchors the energy values of clean data points toward zero to prevent artificial depth biasing in the sampler.
State Space Handling: The method handles mixed continuous-discrete states ( $x = (c, p)$ , where $c$ are coordinates and $p$ are atom types) by defining separate interpolants: linear trajectories for coordinates (aligned via optimal transport) and reflection-based trajectories for categorical features to remain within the simplex.

2. Mirror-Langevin Algorithm (MLA) for Sampling

To sample from the learned distribution $p_\theta(x) \propto \exp(-E_\theta(x)/\tau)$ , the authors adapt the Mirror-Langevin Algorithm.

Unified Updates: Standard Langevin dynamics push atom types outside the probability simplex. MLA resolves this by performing updates in an unconstrained dual space and projecting back via a mirror map.
- For continuous coordinates, the Euclidean potential reduces MLA to the Unadjusted Langevin Algorithm.
- For categorical atom types, the negative entropy potential maps the simplex to logit space, with the softmax function projecting back.
Parallel Tempering: To address the difficulty of traversing rugged energy landscapes, the authors employ parallel tempering with multiple chains at different temperatures. Periodic Metropolis-Hastings swaps allow "hot" chains to explore broadly and "cold" chains to refine, enabling inference-time compute scaling.

Key Contributions

Restoring Field Matching (RFM): A simulation-free training objective that shapes a scalar, atom-additive energy landscape with data points as local minima. The authors provide theoretical guarantees regarding the existence and uniqueness of the solution and characterize the behavior under varying regularization strengths.
Mirror-Langevin Sampling Framework: The first unified mixed-geometry Langevin framework for joint continuous-discrete deep generative modeling, combining coordinate updates and simplex-constrained atom type updates within a single step, augmented by parallel tempering.
State-of-the-Art Performance: EBMol achieves state-of-the-art results on QM9 and GEOM-Drugs benchmarks, demonstrating that energy-based modeling can match the generation quality of diffusion/flow models while offering additional capabilities.

Experimental Results

The model was evaluated on QM9 (134k small molecules) and GEOM-Drugs (430k drug-like molecules).

Generation Quality: On QM9, EBMol (with 1810 Network Function Evaluations) surpasses all baselines in atom and molecule stability, achieving 99.76% and 97.09% respectively. On GEOM-Drugs, EBMol significantly outperforms the strong baseline SLDM, reaching 98.68% atom stability and 61.42% molecule stability at high compute budgets.
Physical Consistency: The learned energy landscape serves as a principled quality metric. EBMol achieves a median relaxation energy of 1.78 kcal/mol on GEOM-Drugs, an order of magnitude lower than previous bond-implicit models and competitive with bond-explicit models (Megalodon) despite having 18 $\times$ fewer parameters and no explicit bond supervision.
Correlation with Physical Energy: The learned energy correlates significantly ( $r=0.65$ ) with the energy decrease upon GFN2-xTB relaxation, allowing the model to act as a lightweight filter for physically strained geometries.
Controllable Generation:
- Shape Steering: By composing the learned potential with differentiable shape potentials (linear, planar, spherical), the model can steer generation toward underrepresented geometries without retraining.
- Zero-Shot Linker Design: The additive nature of the energy allows for conditional generation by fixing molecular fragments. EBMol achieves a 73.9% success rate in linker design tasks across varying complexities without task-specific training.

Significance and Claims

The paper claims that EBMol demonstrates the potential of energy-based models as a versatile framework for 3D molecular generation. Its primary significance lies in restoring the energy inductive bias lost by time-conditional models, thereby enabling:

Physically meaningful scoring: The ability to rank and filter generated samples based on an explicit energy function.
Composable generation: The ability to steer generation via potential composition (e.g., shape steering) and perform zero-shot conditional generation (e.g., linker design) without retraining.
High-quality generation: Achieving state-of-the-art stability and validity while generating structures that are surprisingly close to physical local minima, rivaling models with explicit bond supervision.

The authors note limitations, including the sensitivity of the sampling algorithm to hyperparameters and the current need for significant compute budgets (NFEs) to achieve optimal results. They suggest future work could focus on energy calibration against physical quantities (DFT/GFN2-xTB) and the integration of more sophisticated samplers.