New method uses diffusion models to design energy landscapes that control molecular states
Scientists introduce GB-FESO, a method that aims to design the free-energy surface of a molecular system. Free-energy surfaces describe which molecular shapes (called metastable states) are most likely and how hard it is to move between them. GB-FESO stands for Gradient-Based Free Energy Surface Optimization and uses a trained generative model to guide changes to the system so the overall distribution of shapes matches a desired target.
The core idea is to train a conditional diffusion model — a kind of generative neural network known as a denoising diffusion probabilistic model (DDPM) — to produce ensembles of molecular configurations given a set of design variables. Once trained, the neural network is frozen and treated as a differentiable surrogate for the ensemble distribution. The researchers then change the input design variables and use gradient-based optimization to make the model produce an ensemble whose free-energy surface matches a prescribed target.
To measure how close the generated ensemble is to the target, the team uses a distribution-level loss built from the Kullback–Leibler divergence (KL divergence), a standard way to compare probability distributions. Because the model does not give an analytic formula for its output distribution, the KL divergence is approximated with kernel density estimation (KDE), a simple method that builds a smooth probability estimate from sample points. Gradients of this loss are backpropagated through the deterministic diffusion sampling steps so the design variables can be updated directly.
The paper tests GB-FESO on two kinds of problems. First, on simple one-dimensional Gaussian ensembles they show the method can recover target distributions. They report it can handle both continuous design variables and relaxed discrete variables, and in those tests it could even reach targets outside the model’s original training domain. Second, they apply the method to a small, four-particle Lennard–Jones toy peptide that has multiple metastable conformations. In that physically motivated test the method succeeded in reproducing target free-energy landscapes in the majority of cases. Optimization was done either over the full internal-coordinate description or over a reduced set of collective variables.