EPO: Diverse and Realistic Protein Ensemble Generation via Energy Preference Optimization

AI summary

Overview

Research area: Machine learning for computational structural biology, specifically generative modeling of protein conformational ensembles.

Technical level: Intermediate. The paper assumes familiarity with diffusion/flow-matching generative models and preference-optimization objectives such as Direct Preference Optimization (DPO), though its central idea and results are describable without that background.

One-sentence scope: The paper introduces Energy Preference Optimization (EPO), an online refinement algorithm that converts a pretrained protein ensemble generator into an energy-aware sampler by ranking its own sampled conformations with a physics-based molecular force field and optimizing it with a listwise preference objective, avoiding the need for additional molecular dynamics (MD) trajectories.

What This Paper Is About

Proteins function by moving between many conformations, so understanding them requires recovering the entire Boltzmann ensemble rather than one static structure. Molecular dynamics (MD) can in principle produce these ensembles, but it is computationally expensive and gets trapped behind high energy barriers, while generative models fine-tuned on MD data tend to revisit the same local regions of conformational space even when trained on well-sampled trajectories. This paper asks whether preference signals derived only from physical energy can steer a pretrained generator toward diverse, thermodynamically consistent ensembles without any new MD simulation.

Key Contributions

1. The EPO framework. An online refinement algorithm that directly aligns a pretrained protein ensemble generator with the target Boltzmann distribution using listwise energy preferences, combining stochastic differential equation (SDE) sampling for exploration with energy-ranked preference optimization.

2. A tractable upper bound for the listwise objective. Because exact log-likelihoods over long continuous-time generation trajectories are intractable in flow-matching models, the authors apply Jensen's inequality to the Plackett–Luce choice probability, approximate the intractable reverse transition with the forward process, and replace a key term with a Mean Squared Error (MSE) proxy. This yields gradients compatible with modern denoising score-matching and flow-matching architectures.

3. A practical comparison of pairwise versus listwise alignment. The paper argues that pairwise DPO is structurally unsuitable for free-energy landscapes because every comparison pushes the model toward the single globally lowest-energy basin, erasing the higher-energy basins that ensemble modeling exists to capture. It implements a simplified pairwise variant (EPO-Pair) and a full listwise variant (EPO-List) and shows the listwise version wins in practice.

4. State-of-the-art empirical results without extra MD. EPO reports new state-of-the-art performance across nine distinct distributional metrics on the Tetrapeptides, ATLAS and Fast-Folding benchmarks, requiring no post-hoc MD simulations.

Main Findings

• Nine-metric state of the art on three benchmarks. EPO establishes new state-of-the-art results across nine evaluation metrics on Tetrapeptides (Jing et al. 2024b), ATLAS (Vander Meersche et al. 2024) and Fast-Folding (Lindorff-Larsen et al. 2011).

• Tetrapeptide distributions approach MD replicate quality. On Tetrapeptides, EPO-List reaches Jensen-Shannon divergence values of 0.125 (backbone torsions), 0.093 (sidechain torsions), 0.107 (all torsions), 0.226 (TICA-0) and 0.311 (TICA-0,1), compared with MDGen at 0.130, 0.093, 0.109, 0.230 and 0.316, and a reference replicate-simulation baseline at 0.103, 0.055, 0.076, 0.201 and 0.268. The paper states EPO closely matches the accuracy of replicate 100-ns simulations.

• EPO recovers missing metastable states. Free energy surface analysis shows EPO can induce novel metastable states absent from the pretrained baseline (sequence ASRE) and correct high-energy biases in the pretrained distribution (sequence KSIY).

• ATLAS gains are positive but described as modest. On the 82 ATLAS test proteins, EPO-List improves pairwise RMSD correlation to 0.51/0.51 versus MDGen's 0.48/0.42, lowers MD PCA Wasserstein distance to 1.77/2.35 versus 1.89/2.43, and raises the fraction of PCs with similarity above 0.5 to 13 versus 10. The authors attribute the modest size of these gains to the very small SDE score norm (0.0001) required for long protein chains, which limited stochastic exploration, and note that larger score norms caused sampling instability and structural collapse.

• Zero-shot transfer to Fast-Folding works well. Applying the ATLAS-optimized checkpoint to Fast-Folding with no fine-tuning, EPO-List sets four new state-of-the-art results and ranks second-best on the remaining metrics — for example, JSD of 0.28/0.26 versus ConfDiff-Force at 0.29/0.27, RMSF of 0.32/0.30 versus 0.38/0.38, and Rg of 2.36/2.28 versus 3.25/3.38.

• Listwise beats pairwise in practice. Despite a theoretical equivalence that would require evaluating all (K choose 2) comparisons per batch, EPO-List outperforms EPO-Pair; the paper attributes this to lower gradient variance and less information loss under the full listwise objective.

• SDE sampling beats ODE sampling. On Tetrapeptides, EPO-SDE improves every reported JSD over EPO-ODE (0.125 vs 0.167 backbone, 0.093 vs 0.107 sidechain, 0.107 vs 0.132 all, 0.226 vs 0.247 TICA-0, 0.311 vs 0.337 TICA-0,1). On the SPFH sequence, the SDE sampler identifies distant torsion-angle peaks corresponding to metastable states separated by high energy barriers, whereas the ODE sampler converges to limited local minima.

• Two hyperparameters are sensitive. The temperature parameter β is context-dependent and must be tuned carefully (set to 1 for Tetrapeptides and 250 for ATLAS); too small a value is counterproductive. Denoising steps also matter: too few cause unstable training and mode collapse, while too many give diminishing returns at higher computational cost. Both were set to 50 in the reported runs.

Methodology in Plain English

The authors start from MDGen, a pretrained flow-matching model that generates molecular dynamics trajectories, and treat it as a competent but misaligned sampler. Rather than collecting new MD data, they let the model generate candidate conformations online and score them with Madrax, a differentiable empirical force field implemented in PyTorch. Those energy scores define a ranking: lower energy means better.

The training loop works as follows. The first frame of a protein (a single static conformation) is given to the pretrained generator. Instead of the deterministic ODE sampling the model was trained with, the authors switch to a reverse-time SDE formulation that preserves the same marginal distribution but injects controllable noise, which lets the model escape local energy minima and explore. The energies of the sampled conformations are then used as rewards in a listwise preference objective — a Plackett–Luce-style log-sum-exp over the batch — so the model is nudged toward lower-energy regions while higher-energy samples still contribute gradient.

Because the exact trajectory probability of a flow-matching model is not computable in closed form, the authors derive an upper bound on this listwise loss. They exploit the convexity of the Plackett–Luce function via Jensen's inequality, substitute the intractable reverse transition probability with the forward process probability, and simplify the remaining likelihood ratio into a difference of per-timestep mean squared errors between the reference and optimized velocity networks. That final expression is cheap to compute and fits directly onto existing pretrained generators. Parameter updates use Low-Rank Adaptation (LoRA) for efficiency.

Experiments cover three datasets. Tetrapeptides contains all-atom MD trajectories for 3,000 training, 100 validation and 100 test tetrapeptides, each simulated for 100 ns. ATLAS provides explicit-solvent, all-atom MD with three independent 100 ns trajectories for each of 1,390 structurally diverse protein entries. Fast-Folding includes MD simulation data for 12 small proteins exhibiting rapid folding–unfolding dynamics. Standard splits follow MDGen and AlphaFlow. Baselines include EigenFold, MDGen, Str2Str and ConfDiff. Training used four A100 40GB GPUs on Linux 5.4.0 with PyTorch 1.12 and CUDA 11.3; it took 30 hours on Tetrapeptides and 72 hours on ATLAS. All reported metrics come from a single experimental run, a protocol consistent with the baselines compared against.

Why This Matters

Impact on research. The paper offers an alternative to the two dominant strategies for obtaining protein ensembles: long MD simulations, which are expensive and barrier-trapped, and supervised fine-tuning on MD trajectories, which the authors argue is itself misaligned because the trajectories may undersample the true ensemble and because the learned model still collapses onto visited regions. By showing that energy-only preference signals suffice, it connects the RLHF-style post-training toolkit to continuous-time generative models of molecular structure, and it reframes DPO-style alignment as a tool for matching a physical distribution rather than optimizing a single output.

Real-world applications include:

• Structure-based drug discovery, where binding and allostery depend on conformations that are rare or transient in a single static structure.
• Protein and enzyme engineering, where identifying alternative functional states can guide mutations.
• Interpretation of experimental structural data such as NMR and Cryo-EM, which report on ensembles rather than single conformations.
• Cheap generation of ensemble hypotheses for proteins where running extensive MD is impractical.

Industry relevance. The method targets pharmaceutical and biotechnology workflows that currently depend on expensive simulation infrastructure. Because EPO requires no additional MD trajectories and uses LoRA for parameter-efficient fine-tuning, it is positioned as a refinement layer that can be bolted onto existing pretrained generators rather than a new model that must be trained from scratch. The paper's acknowledgements note support from the Administrative Committee of Zhongguancun Science City and the Beijing Natural Science Foundation (L243006).

Future Directions

• Reducing the cost of the online refinement loop. The authors identify the computationally intensive online loop as a primary limitation and suggest coarse-grained models as a way to lower cost while maintaining fidelity.

• Establishing convergence guarantees. The listwise preference objective is empirically strong but lacks theoretical guarantees of convergence to the target Boltzmann distribution, which motivates investigating optimization frameworks with provable convergence.

• Scaling to larger and more expressive models. The authors plan to scale EPO to more expressive generative models for large proteins and to adopt engineering techniques from recent RLHF-style frameworks.

• Improving SDE exploration on long protein chains. The ATLAS results were limited by the small score norm needed to avoid instability and structural collapse; finding a way to retain stochastic exploration for longer chains is left open.

Target Audience

This paper suits machine learning researchers working on generative models for scientific data, computational chemists and biophysicists interested in protein conformational ensembles, and practitioners in structure-based drug discovery or protein engineering who want to know whether preference-based post-training can substitute for long MD simulations. Readers without a background in diffusion or flow-matching models will find the empirical results accessible, but the derivations of the tractable upper bound assume comfort with continuous-time generative modeling and preference-optimization objectives.

Authors’ abstract

Accurate exploration of protein conformational ensembles is essential for uncovering function but remains hard because molecular-dynamics (MD) simulations suffer from high computational costs and energy-barrier trapping. This paper presents Energy Preference Optimization (EPO), an online refinement algorithm that turns a pretrained protein ensemble generator into an energy-aware sampler without extra MD trajectories. Specifically, EPO leverages stochastic differential equation sampling to explore the conformational landscape and incorporates a novel energy-ranking mechanism based on list-wise preference optimization. Crucially, EPO introduces a practical upper bound to efficiently approximate the intractable probability of long sampling trajectories in continuous-time generative models, making it easily adaptable to existing pretrained generators. On Tetrapeptides, ATLAS, and Fast-Folding benchmarks, EPO successfully generates diverse and physically realistic ensembles, establishing a new state-of-the-art in nine evaluation metrics. These results demonstrate that energy-only preference signals can efficiently steer generative models toward thermodynamically consistent conformational ensembles, providing an alternative to long MD simulations and widening the applicability of learned potentials in structural biology and drug discovery.