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

TL;DR

EPO tackles the difficulty of sampling full Boltzmann protein ensembles by online refinement of pretrained generators using energy-based, listwise preferences. It integrates SDE-based sampling to explore rugged landscapes and adopts a listwise energy ranking with a practical upper bound to align generated ensembles with thermodynamic realism without additional MD trajectories. The approach yields state-of-the-art results across Tetrapeptides, ATLAS, and Fast-Folding benchmarks and demonstrates that energy-only signals can effectively steer generative models toward diverse, physically plausible conformations. This online, MD-free framework potentially broadens the applicability of learned potentials in structural biology and drug discovery by enabling efficient exploration of conformational diversity.

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.

Figures (6)

• Figure 1: EPO aligns the model's energy landscape with the ground truth by leveraging energy rankings from a physics-based molecular force field (PMFF), guiding the model to assign lower energy to more favorable conformations.
• Figure 2: Overview of EPO. (a) Given the first frame (static conformation) of a protein, we leverage a pretrained ensemble generator and introduce an ODE-to-SDE strategy to enable stochastic sampling for online optimization. Energies of the online samples are used as rewards to update the model with LoRA. (b) SDE-based sampling effectively overcomes the energy barriers inherent in the original model (sequence: ASRE). Consequently, the optimized EPO model generates diverse and physically realistic protein ensembles, bypassing the need for any post-hoc MD simulations.
• Figure 3: Comparison of torsion angle distributions and free energy surfaces (FES) for two tetrapeptide sequences. EPO successfully captures crucial metastable states absent in pretrained model outputs and effectively corrects high-energy biases observed in pretrained distributions (highlighted in dotted boxes, respectively), demonstrating strong alignment with the ground-truth energy landscapes.
• Figure 4: Sample distributions projected onto the first two time-lagged independent components for four proteins from the Fast-Folding dataset. EPO-List demonstrates improved diversity by exploring a broader conformational landscape.
• Figure 5: Analysis of key components of EPO. (A) Impact of the temperature hyperparameter $\beta$. (B) Effect of the number of denoising steps in SDE sampling. An insufficient number of steps causes training instability, whereas an excessive number provides no significant improvement. (C) Torsion angle distributions for SDE vs. ODE strategies on the SPFH sequence. The SDE-based model successfully identifies distant peaks corresponding to metastable states separated by high energy barriers.