EBA (Energy-based Alignment)

Proteins are not single structures. They interconvert between conformational states, and which states are populated — and in what proportion — is what determines binding, allostery, and cryptic pocket formation. Molecular dynamics simulation answers this question from physics but needs micro- to millisecond trajectories, often hundreds of GPU-days per system. Generative models trained on crystallographic structures answer it far faster, but they learn the distribution of deposited structures, not the Boltzmann distribution over states, and so have no principled reason to get the relative populations right.

Energy-based Alignment (EBA) closes that gap by fine-tuning a pretrained all-atom diffusion model with feedback from a physical force field. The method treats a mini-batch of K candidate conformations as a finite-state approximation of the Boltzmann distribution and minimizes the KL divergence between that energy-weighted target and the model's own distribution over the same candidates. The formulation is a strict generalization of Direct Preference Optimization: DPO is recovered as the special case K = 2 with the temperature taken to infinity, so instead of binary win/lose pairs, EBA weights every candidate by its actual energy gap. The work is by researchers at Mila, Université de Montréal, McGill University, IBM Research, and HEC Montréal, and was published at ICML 2025.

EBA sits alongside AlphaFlow, which fine-tunes AlphaFold 2 with flow matching, and BioEmu-1, which emulates equilibrium ensembles directly. Its distinguishing move is that the supervision signal is the force field itself rather than more simulation data.

#Key Features

• Energy-weighted alignment objective: Candidate conformations are scored by a physical potential and the model is trained to reproduce their Boltzmann weights, rather than to imitate simulation frames or satisfy binary preferences.
• DPO as a special case: The paper shows analytically that Direct Preference Optimization falls out of the objective at K = 2 and infinite temperature, connecting physics-based calibration to the alignment literature for language models.
• Diffusion-module-only fine-tuning: Only the all-atom diffusion module is updated; the trunk and embedding modules stay frozen, which keeps the cost of alignment far below retraining a structure predictor.
• Length-normalized energies: Because potential energy scales with system size, energies are rescaled by the inverse square root of residue count, which stabilizes training across targets of very different lengths.

#Technical Details

The generative backbone is the AlphaFold 3 architecture as implemented in the open-source Protenix codebase, initialized from Protenix's released weights, with no template inputs. Training proceeds in two stages. Stage one is supervised fine-tuning of the diffusion module on all-atom MD trajectories from the ATLAS dataset — three 100 ns replicas across 1,390 protein targets — using the standard diffusion loss. Stage two applies the EBA objective, with potential energies annotated off-policy by local energy minimization of collected training structures.

Evaluation follows the AlphaFlow protocol on the 250-target ATLAS test split, reporting medians across targets. EBA reaches a pairwise-RMSD correlation of 0.62 and a global RMSF correlation of 0.71, against 0.48 and 0.60 for AlphaFlow-MD and 0.48 and 0.50 for MDGen. On distributional accuracy it reports a root-mean Wasserstein-2 distance of 2.43 (AlphaFlow-MD: 2.61) and an MD PCA Wasserstein-2 distance of 1.19 (AlphaFlow-MD: 1.52). An ablation over the mini-batch size K shows the energy-weighted objective outperforming its DPO reduction on most metrics.

#Applications

Physically calibrated ensembles matter wherever the relative population of states, not just their existence, drives the answer: estimating conformational entropy and flexibility, identifying cryptic and transient pockets for structure-based drug design, interpreting NMR and hydrogen-exchange observables, and generating starting structures for further simulation. Because generation is a diffusion sampling pass rather than a trajectory, ensembles for many targets can be produced in the time a single MD run would take, making the model a practical screening layer ahead of selective, expensive simulation.

#Impact

EBA offers a reusable recipe — treat a force field as a reward model and align a structural generative model against it — that is not specific to conformational ensembles and can in principle be applied to any generative model whose outputs a physical potential can score. The authors are explicit about the limits: the backbone was designed for folding, so its aligned variant inherits folding-oriented biases; energies come from a classical force field and inherit its inaccuracies; and evaluation is confined to ATLAS, an in-silico benchmark, with no experimental validation. Code and the ~3.1 GB aligned checkpoint are released, with Protenix-derived components under Apache 2.0 and the remaining code under MIT.