Aligning Protein Conformation Ensemble Generation with Physical Feedback
Authors: Jiarui Lu
Xiaoyin Chen
Stephen Zhewen Lu
Aurelie Lozano
Vijil Chenthamarakshan
Payel Das
Jian Tang
Venue: ICML 2025 poster (2025)
Keywords: Protein
generative models
molecular dynamics
conformation generation
alignments
Primary area: applications->health_medicine
TL;DR: incorporating physical energy feedback into diffusion model for better protein conformation ensemble generation

Abstract: Protein dynamics play a crucial role in protein biological functions and properties, and their traditional study typically relies on time-consuming molecular dynamics (MD) simulations conducted in silico. Recent advances in generative modeling, particularly denoising diffusion models, have enabled efficient accurate protein structure prediction and conformation sampling by learning distributions over crystallographic structures. However, effectively integrating physical supervision into these data-driven approaches remains challenging, as standard energy-based objectives often lead to intractable optimization. In this paper, we introduce Energy-based Alignment (EBA), a method that aligns generative models with feedback from physical models, efficiently calibrating them to appropriately balance conformational states based on their energy differences. Experimental results on the MD ensemble benchmark demonstrate that EBA achieves state-of-the-art performance in generating high-quality protein ensembles. By improving the physical plausibility of generated structures, our approach enhances model predictions and holds promise for applications in structural biology and drug discovery.

Aligning Protein Conformation Ensemble Generation with Physical Feedback
Jiarui Lu 12 Xiaoyin Chen 12 Stephen Z. Lu
Aurélie Lozano Vijil Chenthamarakshan Payel Das Jian Tang 156
Abstract
of biological molecules. These simulations evolve Newto- nian equations of motion for an entire system of particles, with accelerations determined by pre-defined force fields (physical energy functions). However, capturing biologi- cally relevant transitions, such as folding and unfolding, often requires simulations spanning micro- to millisecond timescales (Lindorff-Larse et al., 2011), I), which is computa- tionally prohibitive, often requiring hundreds to thousands of GPU days depending on system size.
Protein dynamics play crucial role in protein biological functions and properties, and their tra- ditional study typically relies time-consuming molecular dynamics (MD) simulations conducted in silico. Recent advances in generative model- ing, particularly denoising diffusion models, have enabled efficient accurate protein structure predic- tion and conformation sampling by learning distri- butions over crystallographic structures. However, effectively integrating physical supervision into these data-driven approaches remains challeng- ing, as standard energy-based objectives often lead to intractable optimization. In this paper. we introduce Energy-based Alignment (EBA), a method that aligns generative models with feed- back from physical models, efficiently calibrat- ing them to appropriately balance conformational states based on their energy differences. Exper- imental results on the MD ensemble benchmark demonstrate that EBA achieves state-of-the-ar performance in generating high-quality protein ensembles. By improving the physical plausibil- ity of generated structures, our approach enhances model predictions and holds promise for applica- tions in structural biology and drug discovery.
Recently, researchers have turned to deep generative models, particularly denoising diffusion models, to reframe this prob- lem as a conditional generation task (Arts et al., 2023; Lu et al., 2024b; Jing et al., 2024a; Zheng et al., 2024). These models learn from structure databases such as the Protein Data Bank (PDB) and predict plausible conformations con- ditioned on specific inputs, such as amino acid sequences. While these data-driven approaches generate structurally valid candidates, they do not explicitly model thermody- namic properties. A more proper formulation of the problem is equilibrium sampling (Noé et al., 2019), which aims to sample conformation ensembles from the Boltzmann distri- bution over states. This approach is crucial for accurately modeling protein ensembles and capturing thermodynamic stability. However, it remains highly intractable, as it re- quires the generative models to not only produce plausible candidates but also match the underlying energy landscape. Existing amortized sampling methods (Bengio et al., 2021; Zhang & Chen, 2021; Richter & Berner, 2023; Lahlou et al., 2023) also struggle to scale the general protein structures that typically consist of thousands of atoms.
1. Introduction
Understanding protein dynamics is a critical yet complex challenge in the study of protein functionality and regula- tion. Protein structures transition between multiple confor- mational states across varying spatial and temporal scales, influencing their biological roles. Traditionally, molecu- lar dynamics (MD) simulations have been the predominant computational tool for investigating the dynamic behavior
To address these challenges. we introduce Energy-base -basd Alignment (EBA), a framework that integrates diffusion models with physical feedback for improving protein en- semble generation. Our approach bridges the gap between purely data-driven conditional generation and physics-based simulations by incorporating fine-grained force field feed- back via scalable learning objective. Specifically, we fine-tune a pretrained all-atom denoising diffusion model to align generated conformational states with underlying physical energy landscapes. Within this alignment frame- work, the model interacts with force fields as an external environment, receiving feedback to refine its generative pro- cess. Through the EBA alignment, the diffusion model
'Equal contribution 'Mila Québec Al Institute Univerrsité de Montréal "MeGill University 'IBM Research 'HEC Montréal "CIFAR AI Chair. Correspondence to: Jiarui Lu <jiarui @eciilll@mma.cc mulaeq, Jian Tang Kjian.tang@@e @hec.ca>.
Proceedings of the 4rnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
learns balance different conformational states, resulting in more physically consistent protein conformation ensem- bles. We validate our approach on the ATLAS MD ensemble dataset (Vander Meersche et al., 2024) and demonstrate that the EBB-alligedd diffusion model achieves ftae-offtthe-ar performance compared to previous generative models by incorporating physical feedback. The proposed method pro- vides a novel pathway for integrating generative modeling with traditional simulations, presenting promising opportu- nities for protein dynamics research and drug discovery.
Reinforcement learning from human feedback. Rein- forcement Learning from Human Feedback (RLHF) repre- sents a canonical example of the alignment problem. The goal of RLHF is to align a generative model p=(x)c) to maximize the reward ro(x, c) which is provided by a model parameterized by 0, while being constrained by a reference distribution Pter(x/c) (a > 0 is the regularization factor);
max Bxo |ra (x, c)] aDDL (po(x(c) [X(r() (2) Pe
Prior work (Peters & Schaal, 2007; Peng et al., 2019;
Rafailov et al., 2024) has shown that this objective has a where 7 Ex Prer( (xc)) rix.el a
2. Preliminaries
closed-form solution:
Notations. A protein consisting of residues is charac- terized by its amino acid sequence c = (C1; (2,.....00 Ct.) E [V/4, where V denotes the vocabulary of 20 standard amino acids and L is the number of residues. The protein structure is described by the 3D positions of its constituent atoms, x3 encompassing all the heavy atoms (excluding hydrogen) of backbone and side-chains.
pe(x(c) - ()e((
the partition function.
Accordingly, the reward function can be written as:
x (z1, Z....., .N) E IR-
Pe (x(c) (x]c)
ro c) = a log
=alog Z.
(4)
Direct Preference Optimization (DPO) (Rafailov et al., 2024) has demonstrated that when the reward model is opti- mized using the Bradley-Terry (BT) model p(x" > x'|c) o (ro(x"", c) re(x',, c)), the RL problem in Eq. (2) is sim- plified to supervised learning on pairwise preference data:
Molecular dynamics. Molecular dynamics (MD) simu- lations operate by evolving an entire particle system over a simulation time T > O. governed by the physical dynamics dx; Pa/m; dr, where p E R Nx3 m, E R represents the momentum and the mass for each atom in the system. Given a force field (energy function) E : Nx3 + R, the momenta field can be updated according to the Newtonian equation of motion: dp; = Vx, Exx) dr. Using an in- tegrator like Verlet (Verlet, 1967) or Langevin dynamies (with thermostat), the position and momenta are iteratively updated each time step, producing simulation trajectory that converges to the Boltzmann distribution.
LDDO(O) - -B(c,x"xx'>>
log ao po Prot (x" a log po po (c)
prer(x'|c)
(5)
where af.) is the sigmoid function.
3. Method
Diffusion generative models. Diffusion generative mod- els (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020) aim to model data distributions using a diffusion pro- cess described by a (Gaussian) Markov chain of forward distribution: qxx Xo, c) = [(x]|axx) of I). where t E [0, 7) is the time step and ot; ag > O are the noise schedul- ing functions. The reverse (generative) process also assumes Markov structure (x(x)))) = I[ (x |xx)) and the estimated posterior can be written as (Rombach et al., of 2022), Po Po (xx) = -N(xx-1lpoxx, (), of ffj--1 of T), Gelt-19? of
In this section, we introduce Energy-based Aligmment (EBA), novel physics-informed learning objective for im- proving conformation ensemble generation. EBA leverages the underlying Boltzmann distribution by aligning a model's ss learned distribution with energy-induce probability ratios, while circumventing the need to explicitly compute the in- tractable partition function. Instead, it employs an energy- weighted classification-lis objective over stochastic mini- batch combination of conformational states, ensuring that physically meaningful energy terms directly guide the align- ment process. We further demonstrate the connection be- tween EBA and DPO, providing an alternative view of the alignment problem our method optimizes.
where the parameterized mean no (x(x)))
+
' oe 1o xx(x,, t), oi oî o² and azja a for s < l. Following Ho et al. (2020), the denoising objec-
of
tive can be expressed as:
3.1. Energy-based Alignment
C = 100
[w())] c0 (x,, 1
Given an amino acid sequence (or multiple sequence align- ment, MSA) as condition c, our goal is to conditionally sample the protein conformation ensemble {x') that be- longs to e from the target Boltzmann distribution induced
(1)
where E ~ N(,, A & of is a sggaaltt-noise ratio (SNR) and w > is the loss reweighting function.
2

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
p(x)
-
p(x,)
Intractable partition function
Energy-based Alignment e-BE(x))
Boltzmann factor
Is it good?
Potential Energy
-B
Potential Energy
E
p(x)
p(x,
AE
A2g
112
X1
X2
X3
Conformation space
Conformation space
Figure 1. Motivation behind the proposed Energy-based Alignment (EBA). To align the generative model to adhere to the Boltzmann factor across different conformational states, we employ a stochastic approximation the partition function which simplifies the intractable over all possible states to a finite summation, enabling efficient learning from physical energy while maintaining training scalability.
by some potential energy function E(x; c) E IR, that is, (x(c)c) = e ((xc)) 17, where 3 _ 1/kBT T is the tem- perature factor and Z Exe -BE(xe))) However, di- rectly sampling from such distribution is intractable due to the non-trivial partition function (or normalizing constant) Z A Ex e -9E(xx) which requires the sum (or integration) over all possible conformation states in the high-dimension space. An relevant concept is the Boltzmann factor, which describes the ratio of probabilities of two states x² and x as a function of energy difference: Pa(x'le) = -B
tropy of the target distribution which is constant w.r.t. O. Direct optimization of Eq. (6) is intractable due to the sum- mation over all possible conformational states. We approx- imate it by considering a stochastic finite subset of K rep- resentative states {x* )
K J==1 from some proposal distribution
p", such reference MD simulations or generative models. Then we rewrite the first term in Eq. (6) as:
LEBA (0) Fic,(x>D)pp
KK
e -BE(x))))
é (x(xxe))
pn(x c)
log
where AEj- Exx;; c) E(xt; c). This dependence on AEG ensures that the Boltzmann factor is invariant to shifts in absolute energy values, which is particularly important for generative model training since energy scales can vary significantly with the number of atoms in the protein. Lever- aging A Eij enables us to construct learning objectives that effectively incorporate energy feedback, capturing the rela- tive stability of states within the conformation ensemble.
2j=1 (E
2==1 C a
which we refer to as the Energy-based Alignment (EBA) ob- jective. EBA enables stochastic optimization via mini-batch sampling, where the finite-state approximation in Eq. (7) preserves the Boltzmann, mollzmannfactto factor within each mini-batch. This ensures that the probability ratio between any two states x' and xi, remains proportional to their energy difference, i.e., pa(x'e) Pa (x e -a Consequently, EBA approximately aligns the model's learned distribution with the target Boltz- mann distribution while preserving the probability ratio.
Suppose that (=(x(c) = (xE((xx))) (a > 0) is learn- able probabilistic model parameterized by Eg. then we can optimize the model by minimizing the Ballback-Leible (KL) divergence against the target Boltzmann distribution via the cross-entropy:
3.2. EBA for Diffusion Models
De (pD (x(c)((xc)))))) -(p0))=-pssh Epa(xlc) (x((x(c))
Eq. (7) presents a general form where Eg (x; c) is not specif- ically defined. To adapt this objective for diffusion models, we stipulate Ex(x; c) in terms of the probability distribu- tion (x(c).) Suppose that the diffusion model is specified with forward g(x1:7(x)) and parameterized po (Xo.T) We here explore an efficient form of objective for diffusion models under the EBA framework with, we consider the basic form of energy function by using the negative log-likelihood: Ea = - -og pa (x c), where the log Z is omitted because it is constant w.r.t. x. For diffusion models, there exist a chain of latents x1. and we alter the energy to be on the whole chain: Epo(Xirlo) log (c). Then we can factorize the
po(x)(c)
Σ
- Σ log
- H(pm:c)
follows. To begin
e 33E(x'x))
-Eixx log e (
+ Const,
Σ, a
(6)
where H(pB;C) =
is the en-
3

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
Nolsy positions
LSDED Input sequence
C(0) EBA Loss
Input repr.
Denoising steps
Single repr.
Embedding Module
Conditioning
All-Atom Diffusion Module
E 3.4
[8g x
Pair repr.
Aligned distance
-511
Add noise
Multiple sequence alignment (MSA)
X All-atom positions
E--4.0
Residue-level representation
Pre-computed once
Figure 2. Illustration of model architecture and EBA pipeline for energy-weighted preference fine-tuning. Starting from an input sequence and its MSA, residue-level embeddings pre-computed and conditioned guide the all-atom diffusion. At each step. K reference structures (c.g., K = 3) with the same sequence are sampled from the datasct, and their energies E are evaluated. The denoising error is measured by aligned distance in Eq. (14), and combined with E into the EBA-diffusion loss.
joint distribution with the reverse decomposition for pa. i.e. -Po(x)) - Ei=1 Bpo(x.xx-xxxx.e) log po xx, c). Since the prior Pa (XT) is typically chosen as stan- dard Gaussian, the first term becomes constant. How- ever, to optimize the EBA objective containing the above decomposition, we have to sample the latents ac- cording to pe (x(, X1-1|x0, c). which in practice is in- tractable because of the expensive simulation of the re- verse process. Therefore, we substitute the reverse pro- cess of po(XI:T| Xn) with the forward kernel (Wallace et al., 2024), and the energy function now becomes - - E Bo 1|xa.e) log (x 1x,,, c). Due to the introduction of forward q(-), we amend with an en- tropy term and have the final definition of energy function:
(())= S R D.
e -B((x))
o 1o.. x|.e]]
log
2 (BE(()))))
Because both q(x,-1 and po (x c) can be pa- rameterized by gaussian distribution, we thus substitute the KL divergence with the denoising term, setting the weight- ing function (Xe) constant, obtain EBA objective for diffusion as follows:
LEBA-Diftusion = -- -B --
Ea(x;) c) &
K -BE(x':o)
- (x0x
(x0)
+ H.(
j== 2 (Exx)))
log
Σ -x
(10)
(8)
3.3. DPO as a Special Case of EBA
where (g)
denotes the (self) entropy of posterior defined by forward kernel. Intu- itively, the energy function defined by the KL-divergence terms according to the decomposition of the reverse process.
Note that the derivation of the EBA objective in Eq. (7) is independent of the form of energy function Ex(x; c). A notable alternative parametrization the negative log- likelihood ratio between the training model px(x/c) and reference model (x (e.g., frozen checkpoint of the pre-trained model): Ex(x;;) log this parametrization, DPO emerges as a special case of EBA when K 2 and low temperature (i.c., 3 + oo).
By plugging Eq. (8) back into Eq. (7), we further leverage the fact that the log-sum-exp (LSE) function is convex and therefore apply Jensen's 's inequality to push out the expecta- tion (summation), leading to the following upper bound after some algebra (see Appendix A.1 for detailed derivation):
po(xle) Using par(xle)
Without loss of generality, consider a pair of samples

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
X1, X2 s.l. E(x1, c) C)
3.4. All-atom Diffusion with AlphaFold
We base our conditional generative model on the Al- phaFold3 architecture (Abramson et al., 2024), which en- codes sequence information through its MSA Module and PairFormer into conditioning embeddings. followed by an all-atom diffusion module to iteratively refine atomic coor- dinates from a noise distribution into clean conformations. AlphaFold3 explicitly models atom coordinates in the pro- tein structure, rather than relying on coarse- soause-grained rigid frames for backbone atoms and internal torsion angles for side chain atoms which can be hindered to capture subtle structural variations and alternative conformations. Adopt- ing all-atom modeling architecture enables more direct representation of conformational degrees of freedom and facilitates accurate integration of physical feedback.
LEBA)DPO(O)
log polx
2 0-E((x'00)
-B x²)
log
==1 e -8E(xx())
2j=11 a log poixlle] (e)
a log solk
=-Beex) x2) log
a log (e) (e) Px 16) + Marl (c)
log Poloed
(3 + 0c)
-B(e x²) log a aoo po(x' [c] Prer(x!|c)
ar log po 2x c)
Pret
(c)
(11)
which is equivalent to Eq. (5). This derivation demonstrates that EBA reduces to DPO objective when sampling two states and ignoring energy differences between states. Fol- lowing Wallace et al. (2024), we also derive the DPO varia- tion of EBA for diffusion models in Appendix A.2.
3.5. Structural Alignment- Objective
From the derivation of objective in Eq. (10), the stan- dard mean squared error (MSE) or L2 loss L(0) = ||e co (xxx, 1) is commonly used minimize the difference between predictions and ground truth. However, this MSE term may be suboptimal for protein conformation genera- tion. The input condition (amino acid sequences) is SE(3)- invariant' meaning it does not differentiate between roto- translational variants of the target protein, which can nega- tively affect the training efficiency. Therefore, we consider the following SE(3)-invariaa) losses for EBA fine-tuning:
Algorithm 1 Fine-tuning Diffusion Model with EBA
Require: Pre-trained denoising network Xg, conformation dataset D = {(c, x with sequence conditions e and corresponding structures Xo, energy function E'(; c), inverse temperature factors a, B, the number of time steps T, learning rate 7.
1: while not converged do
2: Sample (c, Xo) from D
Rigid aligned MSE. To ensure that the loss remains in- variant global rotational and translational differences, first align the predicted atomic coordinates to the ground truth structure using the Kabsch's 's algorithm. This align- ment minimizes the root-mean-squee deviation (RMSD) by applying an optimal rotation and translation. The MSE loss is then computed on the aligned coordinates of all atoms, allowing the model focuses on meaningful conformational changes rather than arbitrary positional offsets:
3: Let x) f Xo
4: Randomly retrieve K 1 samples (ct, xx) D
s.t. c' c (i= 2, .,...KK)
5: Calculate energy E(x];c) for each I,..., K
6: Calculate mini-batch Boltzmann weights:
e -BB(x)))
wo xxx6)
Σラ11 e (F(((()))
7:
for each candidate = 1,..., K do
N Σ
8: Sample timestep t ~ Uniform T) Add noise x ((x:1x5)
L Llligned MSE
|Align(Xx,, xo): - Xo. (12)
9:
10:
Forward denoising network x (x{
where Align(Xo, xo). aligns the predicted coordinates &, to the ground truth coordinates Xy E NRN and N is the number of atoms in the structure (I = 1, .).
11:
Calculate the aligned loss Li in Eq. (14)
12:
end for
13:
Calculate EBA loss according to Eq. (10)
While alignment- MSE ensures structural consistency in Cartesian space, it does not ex- plicitly capture inter-atomic relationships that are important for protein geometry. To address this, we additional intro- duce an auxiliary loss term based on the pairwise distances.
Smooth LDDT.
Ciu
LERA-Diflusin
Σ w (xx) log
k
_j=1
14:
Update parameters: 0 +- 0 V end while
15:
A function f is said to be SE(3)-invarina) if its output remains unchanged under arbitrary rigid transformations (i.e., rotations and translations): f(Rx + t) = fxx), VR E SO(3), E R².
16:
Return fine-tuned denosing network Xg.
5

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
Table Statistical metrics MD ensemble henchmark of ATL.AS test (N 250 following Jing (2024a)) where the median
across all test targets is reported. The runtime is reported GPU second required per sample averaged on all test targets. The best result highlighted bold, while the second-best result is underlined.
AlphaFLow-Ny MD
MSA subsampling
Ours
Metrics/ / Methods
Full
Distilled
256
MDGen
Pre-train
EBA-DPO
EBA
Pairwise RMSD 1
0.48
0.48
0.03
0.12
0.22
0.15
0.48
0.43
0.59
0.62
Predicting flexibility
Global RMSF r
0.60
0.54
0.13
0.23
0.29
0.26
0,50
0.50
0.69
0.71
Per-target RMSF
0.85
0.81
0.51
0.52
0.51
0.55
0.71
0.72
0.90
0.90
Root mean disd
2.61
3.70
6.15
5.32
4.28
3.62
2.69
3.22
2.43
2.43
7 Trans, contrib. +
2.28
3.10
5.22
3.92
3.33
2.87
2,47
2.05
2.03
Distributional accuracy
Var. contrib. 4
1.30
1.52
3.55
2.49
2.24
2.24
1.89
1.20
1.20
MD PCA W2-dist. +
1.52
1.73
2.44
2.30
2.23
1.88
1.89
1.78
1.20
1.19
Joint PCA Wa-dist. 4
2.25
3.05
5.51
4.51
3.57
3.02
2.47
2.04
PC--sim > 0.5
44
34
15
18
21
21
28
38
44
Weak contacts
0.62
0.52
0.40
0.40
0.37
0.30
0.51
0.23
0.63
0.65
Ensemble observables
Transient contacts J
0.41
0.28
0.23
0.26
0.27
0.27
0.25
0.38
0.41
Exposed residue J
0.50
0.48
0.34
0.37
0.37
0.33
0.29
0.29
0.68
0.70
Exposed MI matrix t
0.25
0.14
0.14
0.11
0.10
0.06
0.01
0.35
0.36
Runtime
GPU see, per sample
70
7
4
0.2
0.9
Specifically, for each structure, we apply the Smooth LDDT loss (Abramson et al., 2024) which defined as:
on the ATLAS dataset (Vander Meersche et al., 2024) follow- ing the benchmark in Jing et al. (2024a). For baselines, we mainly compare our methods against: AlphaFold2 (Jumper et al., 2021) with MSA subsampling (Del Alamo al., 2022), AlphaFlow (Jing et al., 2024a) and MDGen (Jing et al., 2024b). We follow the evaluation metries introduced in Jing et al. (2024a) and classify them as below:
Smooth LDDT
meaat Eisr o - A <
meant m 1 < 15 A5
(13)
where ADi A abs (D D E NNN, D Dist (xp) NNN and De Dist(xn) E RRNN represent the pairwise distance matrices of predicted and ground truth structures, and (-)) indicates the sigmoid function. This term penalizes deviations in pairwise distances, encouraging the preservation of local and global structural features.
Flexibility correlation (t): the pearson correlation r of the pairwise RMSD, global root-mean-square- fluctutation (RMSF) and per-target RMSF.
Distributional accuracy: Root mean of Wasserstein distance (W2--dist and its translation and variance con- tribution (L), MD PCA Wn-dist (4), joint PCA Wi-dist (4); the percentage of samples with PC-sim > 0.5 (f).
The final denoising training loss a weighted combination of the two components above:
Ensemble observables (f): the Jaccard similarity J of the weak contacts, transient contacts, and exposed residue as well the Spearman correlation p of the exposed mutual information (MI) matrix.
Luotal Amse Aligned MSE + Apadi LSmoouh LDDT-
(14)
Amse. Audit 2 0 are hyperparameters that control the relative importance of the aligned MSE and smooth LDDT terms. Putting it together, we compute the MSE term in Eq. (10) by Eq. (14), which ensures both structural alignment and geometric accuracy and enables the model to better cap- ture biologically relevant conformations. EBA fine-tuning procedure for diffusion models is detailed in Algorithm 1.
4.2. Training Pipeline
We adopt Protenix (Chen et al., 2025) as AlphaFold3 archi- tecture implementation and initialize the model with the pre- trained parameters from the released weights. No template input is used across our study. The ATLAS trajectories pre- pre-processed by parsing, following Jing al. (2024a), and by applying the training data pipeline in the open-sourced Protenix codebase. We fine-tune the pre- pre-trained DIFFUSION MODULE two stages including supervised fine-tuning and physical alignment, with the parameters of other modules frozen.
4. Experiments
4.1. Setup
To demonstrate the effectiveness of the proposed fine-tuning teninu pipeline, we evaluate the protein ensemble generation task
6

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
Stage 1: Supervised Fine-tuning. In the first stage, the pre-trained model is adapted via fine-tuning to the ATLAS simulated trajectories (Vander Meersche et al., 2024). AT- LAS provides diverse ensembles of protein conformations generated via all-atom MD) simulations. During this phase, the model is trained to minimize the vanilla diffusion loss. This fine-tuning stage enables the diffusion to coarsely adapt to the data distribution over conformational space sampled by the MD) simulators across different targets.
and plot the Ca-RMSF against the residue index in Figure 3.
Table 2. Ablation results MD ensemble benchmark of ATLAS test set (Jing et al., 2024a) for EBA with different number of candidate samples K in mini-batch during alignment training.
Metrics
K =
K =
K
Pairwise RMSD r t
0.62
0.61
0.62
Global RMSF r
0.71
0.71
0.72
Per-target RMSF 1
0.90
0.90
0.89
Stage 2: Physical Alignment. The second stage employs the EBA alignment to align the diffusion with physical en- ergy feedback. We use (1) Direct Preference Optimization (EBA-DPO) that leverages energy-agnostic preference pairs (binary win-lose pairs) to align the model's predictions, emphasizing conformational states with favorable energy profiles; (2) the proposed EBA objective to account for the energy difference A between multiple conformational states weighted by physical energies. To make training more efficient, we collect the ATLAS training set and annotate their potential energy off-policy with local minimization. The detailed protocol can be found in Appendix C.
Root mean W2-dist, +
2.43
2.42
2.40
4
Trans. contrib. J
2.03
2.02
2.05
4
Var. contrib. +
1.20
1.18
1.25
MD PCA W2-dist.
1.19
1.18
1.16
Joint PCA W2-dist. 4
2.04
2.04
2.15
% PC-ssim > 0.5 f
44
39
43
Weak contacts J f
0.65
0.65
0.65
Transient contacts J
0.41
0.41
0.39
Exposed residue J 1
0.70
0.69
0.70
Exposed MI matrix p
0.36
0.37
0.34
Iteration per step (s)
4.3
5.4
7.8
Avg. GPU memory (GB)
12.0
13.9
16.3
In practice, we also observed that Ex) varies significantly with protein size (the number of residues or atoms), making the energy-informed training unstable due to large variance in the objective. Inspired by Naganathan & Muñoz (2005), which shows that the folding time scales with the number of residues (with an exponent of 0.5), we introduce a sample- specific regularization factor 7,0.5 normalize the energy- Concretely, we rescale by multiplying 3 with 1.5 for each sample, where L denotes the number of residues in x.
4.4. Ablation Study
To assess the impact of the size K of mini-batch on the performance of EBA during fine-tuning, we further conduct an ablation study on the MD ensemble benchmark of the ATLAS test set. Table 2 presents the results for varying values of K, demonstrating the robustness of our mini- EBA objective Eq. (10). Note that large K is at the cost of increased computational overhead, as the model's forward pass scales linearly with K. Also, we investigate the effect of x and r), scaling factors for noise and backward step size respectively, on the quality of generated ensemble in Fig. 4.
4.3. Benchmark on Molecular Dynamics Ensembles
We evaluate the models to emulate the conformation ensem- bles of ATLAS MD simulation dataset (Vander Meersche et al., 2024), which includes three replicas of 100ns produc- tion trajectories of 1390 protein targets in total. We strictly follow the experimental settings as well as the data split in Jing et al. (2024a) and sample 250 predictions per test target using different models. Baseline evaluations are taken from the original tabular results in Jing et al. (2024a) and Jing et al. (2024b). As shown in Table 1, our methods con- sistently outperform baseline models as well as seeing an improvement than EBA-DPO (an implementation of DPO within our framework), where we set K- for EBA model to make fair comparison with DPO. The EBA model shows significant improvement in exposed residue J and the ex- posed MI matrix p. along with high RMSF correlation This suggests that our model captures collective long-range dynamics by exposing buried residues to solvent, as a phys- ically plausible behavior where other baselines perform poorly. As an illustration, we visualize the sampled ensem- bles for the test target 6uof following Jing et al. (2024a)
Additionally, we conduct simple profiling different K(= 2, 3, 5) regarding the training speed and GPU memory utilization. The DDP parallelism and 4x A100 40GB GPUs are used for this benchmarking. As shown in Table 2, the increase of K does not introduce heavy memory overhead, while the iteration step grows in quasi-linear trend.
4.5. Discussion
Our approach outperform existing methods in MD ensemble generation based off AlphaFold3's s direct atomic interaction modeling, which reduces the need for coarse- searse-graining or in- ternal coordinates while allowing for more precise capture of fine- -ne-ggri conformational changes. Beyond data-driven training, the model is fine-tuned with physical feedback, ensuring that learned structures are not only statistically plausible but also physically consistent. Especially, through EBA, we align the model to (approximately) emulate the Boltzmann distribution, enhancing its capability to generate
7

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
soammodddnamis-ccosissteee ensembles. Our adaptation ef- fectively bridges neural and physical methods a promising framework for protein conformation ensemble generation.
MD
Pretrain mo.43
EBA (ours) mo.92
Importantly, although the derivation sums over only K mini- batch samples, EBA is not designed to approximate the intractable partition function Z Eq. 3 (which beyond the scope of this paper), since doing so introduces non- negligible truncation error. Rather, the EBA objective lever- ages the Boltzmann factors, i..,, the relative weights among conformations within each sampled mini-batch, which re- main invariant to the batch size and adhere to partial knowl- edge derived from the Boltzmann distribution.
ATLAS HD
EBA
Reddue index
Figure 3. (Top) Structure ensembles for the target 6uof. in AT- LAS test set with RMSF correlation r labeled. (Bottom) Ca- RMSF versus the residue index (N + C. terminus from left right).
5. Related work
Protein conformation generation. Unlike structure pre- diction (Jumper et al., 2021) aiming identify a single, most-likely folded structure, protein conformation genera- tion focuses on sampling an ensemble of physically plau- sible states that capture the underlying energy landscape. Boltzmann generator (Noé et al., 2019) leverages normaliz- ing flows to approximate the Boltzmann distribution by train- ing on simulation data. Arts et al. (2023) applies the diffu- sion model to capture such distribution over coarse- grasse-graindd protein conformations. EigenFold (Jing et al., 2023) adopts generative perspective on structure prediction, enabling the generation of multiple structures given an input sequence. Str2Str (Lu et al., 2024b) introduces score-based sampler trained exclusively on PDB data, framing conformation gen- concttur-to-strucctte
igates these issues by directly optimizing for the optimal policy via re-parameterization of an implicit reward model. This approach has been extended beyond language model- ing: Diffusion-DPO (Wallace et al., 2024) for text-to-image generation, ABDPO (Zhou et al., 2024) for antibody design using Rosetta energy (Alford et al., 2017), and ALIDIFF (Gu et al., 2024) and DECOMPDPO (Cheng et al., 2024) for molecular optimization in structure- based drug design.
Remarks: Our method differs from existing approuches above by adopting more general-form objective, being grounded in physically meaningful motivations, addressing different task and demonstrating superior performance.
paradigm. DiG (Zheng et al., 2024) trains conditional diffusion model both PDB and in-house simulation data. ConfDiff (Wang al., 2024) incorporates the energy- and force-guidance during the reverse process of diffusion to enhance the accuracy of conformation generation. AlphaFlow (Jing et al., 2024a) repurposes the AlphaFold2 model into a denoising network via flow matching. ESMDiff (Lu et al., 2024a) fine-tunes the protein language model ESM3 using discrete diffusion to produce protein conformations. Finally, MDGen (Jing et al., 2024b) attempts direct generation of MD trajectories by modeling them as time-series of protein structures.
eration in
6. Conclusion and Limitations
Conclusion. In this study, we introduced physics- inspired alignment framework, namely EBA, for protein ensemble generation. EBA leverages energy feedback to refine the pre-trained diffusion models based on the Boltz- mann factor between different conformational states. Our approach effectively bridges structure data with physical signal, enabling scalable and physically grounded model alignment. Built upon AlphaFold3, the fine-tuned diffu- sion model demonstrates the effectiveness of the proposed EBA by benchmarking the ATLAS MD dataset. Results underscore the potential of incorporating physical energy supervision into the data-drive models, advancing more accurate and thermodynamically consistent ensemble gen- eration. Our method opens new avenues for applications in modern drug discovery and biomolecular simulations.
Alignment methods for generative models. Aligning generative models with desired objectives is becoming in- creasingly important. The Reinforcement Learning from Human Feedback (RLHF) framework optimizes models via RL using human preference rewards and has been widely ap- plied in tasks like machine translation (Kreutzer et al., 2018), summarization (Stiennon al., 2020), and instruction fol- lowing (Ouyang et al., 2022). RLHF has also been applied for alignment of text-to-image diffusion models (Black et al., 2023; Fan et al., 2024). However, RL-based fine-tuning faces significant challenges in stability and scalability. Di- rect Preference Optimization (Rafailov et al., 2024) mit-
Despite its advantages, approach has several limitations. First, since the base model, AlphaFold3, is originally designed for folding, its fine-tuned variant may not be readily well-suited for modeling tingtiime dy- namics (Lindorff- et al., 2011). Second, the accu- racy of the employed energy (force fields) needs further
Limitations.
8

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
Structire SMSOO
Structure SMIS
Impact Statement
Poluabe teap
Ratrabo R80003 Parntte RNGD
Patradte Bñ007
This study advances the intersection of generative model- ing and computational biology, enhancing protein structure prediction and conformation sampling with potential appli- cations in drug discovery and biotechnology. While these advancements hold promise for scientific innovation, we acknowledge potential risks associated with Al-driven bi- ological molecular design, including the misuse of genera- tive models for unintended purposes. Responsible use and oversight in applying these models to sensitive biological domains are important considerations. However, we do not elaborate further specific ethical concerns, the studied problem and related techniques are well established in the field of machine learning.
Laniids
Sta
1.5
Revidan Picelaiity
1001
2.23
Grabhel aHfr
130
100
Distributional
2.5
Dy Tr-diaance
24
References
1.00
125
Lambdda
3.55
3.00 Eia
Abraham, M. J., Murtola, T., Schulz, R., Páll, S., Smith, J. C., Hess, B., and Lindahl, E. Gromacs: High perfor- mance molecular simulations through multi-level paral- lelism from laptops to supercomputers. SoftwareX, 1: 19-25 15-25, 2015.
Figure 4. Evaluation performances versus different values of infer- ence hyperparameters including scaling factors for noise À and and backward step size n during diffusion sampling (N 250).
improvement, as their precision falls short of quantum-level single-point energy calculations. Third, our current study is restricted generating single-chain protein ensemble. Finally, have solely implemented and evaluated EBA within the diffusion framework, leaving open the exploration of alternative generative models. Future work will focus on addressing these limitations by extending to broader biomolecular systems, incorporating quantum-level energy calculation, and exploring more alignment implementations.
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T., Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J.. Bambrick, J., et al. Accurate structure prediction of biomolccular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Alford, R. F., Leaver-Fay, A.. Jeliazkov, J. R., O O'Meara, M. DiMaio, F. P., Park, H., Shapovalov, M. V., Renfrew, P. D., Mulligan, V. K., Kappel, K., et al. The rosetta all- atom energy function for macromolecular modeling and design. Journal of chemical theory and computation, 13 (6):3031-3048, 2017.
Acknowledgements
We thank Meihua Dang, Qijia He, Dinghuai Zhang, Zuobai Zhang as well as anonymous ICML reviewers for their help- ful feedback and discussions. We also acknowledge the authors of Protenix for open-sourcing the code of training pipeline and model implementation. The authors acknowl- edge funding from the Canada CIFAR AI Chair Program and the Intel-Mila partnership program. The computation resource of this project is supported by Mila and the Digital Research Alliance of Canada.
Arts, M., Satorras, V. G., Huang, C.-W., Zuegner, D., Fed- erici, M., Clementi, C., Noé, F., Pinsler, R., and Berg, R. v. d. Two for one: Diffusion models and force fields for coarse-grained molecular dynamics. arXiv preprint arXiv:2302.00600 2023.
Bannwarth, C., Ehlert, S., and Grimme, S. Gfn2-xtb an accurate and broadly parametrized self-consistent tight- binding quantum chemical method with multipole elec- trostatics and deeeedet dispersion contributions. Journal of chemical theory and computation, 15(3):1652- 1671, 2019.
Bengio, E., Jain, M., Korablyov, M., Precup, D., and Ben- gio, Y. Flow network based generative models for non- iterative diverse candidate generation. Advances in Neu- ral Information Processing Systems, 14:27381-27394, 2021.

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
Best, R. B., Zhu, X., Shim, J., Lopes, P. E., Mittal, J., Feig, M., and MacKerell Jr, A. D. Optimization of the additive charmm all-atom protein force field targeting improved sampling of the backbone o, s0 and side-chain x1 and x2 dihedral angles. Journal of chemical theory and computation, 8(9):3257-3273, 2012.
Jing, B., Berger, B., and Jaakkola, T. Alphafold meets flow matching for generating protein ensembles, 2024a. URL https: www.wawwwwww..///////... 240..////.22....55..
Jing, B., Stärk, H., Jaakkola, T., and Berger, B. Genera-
tive modeling of molecular dynamics trajectories. arXiv preprint arXiv:2409. 17808, 2024b.
Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13300 2023.
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zídek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-58) 2021.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint 1rXiv: 12.60080 2014.
Chen, X., Zhang, Y., Lu, C.,, Ma, W., Guan, J., Gong, C., Yang, J., Zhang, H., Zhang, K., Wu, S., et al. Protenix- advancing structure prediction through comprehensive alphafold3 reproduction. bioRxiv, pp. 2025 01, 2025.
Kreutzer, J., Uyheng, J., and Riezler, S. Reliability and learnability of human bandit feedback for sequence-to sequence reinforcement learning. Annual Meeting of the Association for Computational Linguistics, 2018. doi: 111863//11P11 1165.
Cheng, X., Zhou, X., Yang. Y., Bao, Y., and Gu, Q. Decom- posed direct preference optimization for sauucture-base drug design. arXiv preprint arXiv:2407. 1177.11390 13981, 2024.
Del Alamo, D., Sala, D., Mchaourab, H. S., and Meiler, J. Sampling alternative conformational states of transporters and receptors with alphafold2. Elife, 11:e75751, 2022.
Lahlou, S., Deleu, T., Lemos, P., Zhang, D., Volokhova, A., Hernández Gareia, A., Ezzine, L. N., Bengio, Y., and Malkin, N. A theory of continuous generative flow works. In International Conference Machine Learning, pp. 11826-18888 18300. PMLR, 2023.
Eastman, P., Galvelis, R., Peláez, R. P., Abreu, C. R., Farr, S. E., Gallicchio, E., Gorenko, A., Henry, M. M., Hu, F., Huang, J., et al. Openmm 8: Molecular dynamies simulation with machine learning potentials. The Journal of Physical Chemistry B, 2023.
Lewis, S., Hempel, T., Jiménez-Luna, J., Gastegger, M., Xie, Y., Foong, A. Y., Satorras, G., Abdin, O., Veeling, S., Zaporozhets, I., et al. Scalable emulation of pro- tein equilibrium ensembles with generative deep learning, bioRxiv, pp. 2024-12, 2024.
Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, K. Rein- forcement learning for fine-tuning text-to-image diffusion models. Advances in Neural Information Processing Sys- tems, 36, 2024.
Lindorff-Larsen. K., Piana, S., Dror, R. O., and Shaw, D. E. How fast-folding proteins fold. Science, 334(6055):517- 520, 2011.
Gu, S., Xu, M., Powers, A., Nie, W., Geffner, T., Kreis, K., Leskovec, J., Vahdat, A., and Ermon, S. Aligning target-aware molecule diffusion models with exact energy optimization. arXiv preprint 2rXiv:2407.01644, 2024.
Lu, J., Chen, X., Lu, S. Z., Shi, C., Guo, H., Bengio, Y., and Tang, Structure language models for protein con- formation generation. arXiv preprint arXiv:2410. 1..24400.44000 18403, 2024a.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in Neural Information Process- ing Systems, 33:6840-6855 2020.
Lu, J., Zhong, B., Zhang, Z., and Tang, J. Str2str: A score- based framework for zero-shot protein conformation sam- pling. In The Twelfth International Conference on Learn- ing Representations, 2024b.
Huang, J., Rauscher, S., Nawrocki, G., Ran, T., Feig, M., De Groot, B. L., Grubmüller, H., and MacKerell Jr, A. D. Charmm36m: an improved force field for folded and intrinsically disordered proteins. Nature methods, 14(1): 71-73, 2017.
Naganathan, A. N. and Muñoz, V. Scaling of folding times with protein size. Journal of the American Chemical Society, 127(2):480-48) 2005.
Jing, B., Erives, E.,, Pao-Huang, P., Corso, G., Berger, B., and Joakkola, T. Eigenfold: Generative protein struc- ture prediction with diffusion models. arXiv preprint arXiv:2304.00 02/ 198, 2023.
Nguyen, H.. Roe, D. R., and Simmerling, C. Improved generalized born solvent model parameters for protein simulations. Journal of chemical theory and computation, 9(4::2200-200 2013.
10

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltzmann gener- ators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457):eaawl 147, 2019.
Verlet, L. Computer" experiments" classical fluids. i. thermodynamical properties of lennard-jones molecules. Physical review, 159(1):98, 1967.
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C. L., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., Schulman, J., Hilton, J., Kelton, F., Miller, L. E., Simens, M., Askell, A., Welinder, P., Christiano, P., Leike, J. and Lowe, R. J. Training language models to follow instructions with human feedback. Neural Information Processing Systems, 2022.
Wallace, B., Dang, M., Rafailov, R., Zhou, L.,, Lou, A., Pu- rushwalkam, S., Ermon, S., Xiong, C., Joty, S., and Naik, N. Diffusion model alignment using direct preference op- timization. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), PP. 8228-8238, June 2024.
Wang, Y., Wang, L.,, Shen, Y., Wang, Y., Yuan, H., Wu, Y., and Gu, Q. Protein conformation generation via force- guided se(3) diffusion models, 2024. URL https: axxiv.org/abe htt///bb/223..10000
Peng. X. B., Kumar, A., Zhang, G., and Levine, S. Davantage-weighee regression: Simple and scalable off-policy reinforcement learning, arXiv preprint arXiv: 1910.00177. 2019.
Zhang, Q. and Chen, Y. Path integral sampler: stochas- tic control approach for sampling. arXiv preprint arXiv:2111. 1..1:/1, 2021.
Peters, J. and Schaal, S. Reinforcement learning by weiiht-weighted regression for operational space con- trol. Proceedings of the 24th International Confer- ence on Machine Learning, ICML '07, pp. 745-750, New York, NY, USA, 2007. Association for Comput- ing Machinery. ISBN 9781595937933. doi: 10. 1145/1273496 12733900 URL setps://doi...
Zheng, S., He, J., Liu, C., Shi, Y., Lu, Z.. Feng. W., Ju. F., Wang, J., Zhu, J., Min, Y., al. Predicting equilibrium distributions for molecular systems with deep learning. Nature Machine Intelligence, pp. 1-10, 2024.
Zhou, X., Xue, D., Chen, R., Zheng, Z.. Wang, L., and Gu, Q. Antigen-specifee antibody design via direct energy-based preference optimization. arXiv preprint 10Xhv:2403.16576, 2024.
Rafailov, R., Sharma, A., Mitchell, E., Manning. D., Er- mon, S., and Finn, C. Direct preference optimization: Your language model is secretly reward model. Ad- vances in Neural Information Processing Systems, 36, 2024.
Richter, L. and Berner, J. Improved sampling via learned diffusions. arXiv preprint arXiv:2307.01 198, 2023.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffiusinn models. In Proceedings of the IEEE/CVF con- ference on computer vision and pattern recognition, pp. 10684-10695, 2022.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International Conference on Machine Learning, pp. 2256-2265. PMLR, 2015.
Song. J., Meng, C., and Ermon, S. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Stiennon, N., Ouyang, L., Wu, J., Ziegler, D. M., Lowe, J., Voss, C., Radford, A., Amodei, D., and Christiano, P. Learning to summarize from human feedback. Newral Information Processing Systems, 2020.
Vander Meersche, Y., Cretin, G., Gheeraert, A., Gelly, J.-C., and Galochkina, T. Atlas: protein flexibility description from atomistic molecular dynamics simulations. Nucleic acids research, 52(D1):D384-D333 2024.

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
A. Theoretical Results
A.1. Proof of Equation 9
For simplicity, we denote the decomposed negative KL-divergence term (xx 11x (x/-11xX,, c)] as F(x),
where t is the index of time step and is the sample index. Let w(x;) - - weight. We aim to establish that:
Σ5 oe(wie) (i 1...., K) be the Boltzmann
K w ((+)) log
- Pap F(x',))
Σ
- -(,
K E B/
u//o (ng) P(a)a))
Eg
1=1
C
T
w(x;) log e F(alit)
s {-
xx (x
(15)
Σg cl(x),))
1=1 1=
Firstly, we rewrite the lent-hand side by expanding the log: The left-hand side is:
K
Ie Fepaeii F(a
Σ (x;) log
(x) JND
Ex e
E, E -e-oea//in F(z0,t)
K
K
E x mat F(a),))
-B((,(())))
Σ wxxx)
Σ Exi Bxt F(x², - log
j=1
= (o(0
Σ (xx;)
K
T
Ex xxx/ (x))))) F(xi, + log
K
e
Ee
2f ind) F(at.))
j=1
For the denominator term:
K
Σ B (g P F(22.)
log
j=l
we notice that this the log-sum-exp (LSH)) function. Apply the Jensen's 's inequality for with convexity of the LSE function and we can pull the summation and expectation outside:
T
log
e
E. F 10/000//1b (F(2,00
Σ
E t=1
~ (xe)
log
P(*),))
Thus, the left-hand side satisfies:
E, B A (F('
((a')))
log Σ, Σ, Axgaann xg) [[(0),0]
B (x//////) log E, cf((),))
Because the Boltzmann weights are strictly positive, ΣK w(x,) is convex linear combination, following the inequality above and thus we have:
K w(xx log
2 Pafarcttt ((aia))
T K
e((a',))
-B(o,()))) JDD Σ
E, E. F(o) F(w),)) S --B (
DD
E xx
w(x,) log
ΣS K (F(c)
ΣØ c
The inequality is now proved.
12

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
A.2. Derivation for EBB-DPO-Diftusion
To address the intractability of p=(x)c for diffusion models, Wallace et al. (2024) has derived ELBO objective for diffusion DPO. We provide brief summary of their derivation here and refer the readers to Wallace al. (2024) for the complete details.
The diffusion ELBO for the objective in Eq. 11 is:
log log (alog ( pe po(x}-11x3: c) Pron [2] c) (16)
LEBA-DUE DPAAA-DDD-iiiisioonn (0) - -
The reverse process po xx:T|x)) can be further approximated by the forward process q(x1:1|x0):
CEBA-DPO-Difussoon (0)
log a-aa((
DKL DxLla(x;-1) [2 Dxx [ x (x) x7 c)] Di 1x
Dxi
|xx
(17)
With some algebra, the loss simplifies to:
LEBA-DPO-Diffusionn (0) - logo--o
||e - co t, c)ll - ||e* (x) t, c)ll - ||e² tt -x,((x,
(18)
In practice, we combine a and T into a single hyperparameter, denoted as T. Similar to EBA-Diffusion in Algorithm 1, the MSE terms are computed by Eq. 14.
A.3. Further Discussion on the Learning Objective
Here we show the optimization of DxL(PF) lig regarding parameters 0 in Eq. 6 will finally converge to a minimizer of
DEL. (pal (pallpa) in the following sense. Suppose the a global minimizer of Dxx (1 w.r.t. V( [ ~ PB. Let
K = 2, Vx x. we have p K which implies that (a, 6 > 0):
exp(-bE(x")) (xp())/xxppp((((((p = exp(=aEg(z')) exp(-axxppppleeeee
(19)
Then we see -b[B(x*) - B(r)) = -a -a(Fa - -E( }, or equivalently
Eo(x*) = b E(x²) [Eo())) b
E(r')),
a
a
Note that by marginalizing.j j. we that
Ex(x²)- = E(x') + Const, Vi
Finally plug in, Var.
b po(x*) = exp(- = exp( aa( =((x') + Co =xx =pn(x*).
or (z) = (x) This indicates that model alignment according to Boltzmann factor can also provide useful guidance towards the underlying Boltzmann distribution: when the equality Eq. 19 holds for any pair æ , then we immediately have Po = PB-
13

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
B. Additional Experimental Results
In this section, we additionally incorporate more reports as comparison to give interested readers better overview of our methods as shown in Table SI. Results of ConfDiff (Wang et al., 2024) and BioEmu (Lewis al., 2024) are obtained by performing inference pipeline using code and checkpoints from their official repositories. For the ConfDiff, we use the checkpoint named OpenF OpenFold-r3e 31d-r3-MD which was finetuned on the ATLAS dataset (Wang al., 2024). Notably, the proposed EBA consistently outperformed the pre-trained and the Stage2-SFT checkpoints (after two stages of supervised fine-tuning, see Section 4.2), which demonstrates the proposed EBA objective provides additional knowledge from the energy landscape beyond the maximizing likelihood training.
Table S1. Supplementary comparison MD ensemble benchmark of ATLAS test following Jing al. (2024a) where the median
across all test targets is reported.
ESMFLOW-MD
Ours
Metrics/ Methods
Full
Distilled
AlphaFold
ConfDDff
BioEmu
Pre-train
Stage2-SFT
EBA
Pairwise RMSD t
0.19
0.19
0.10
0.59
0.46
0.43
0.57
0.62
Predicting flexibility
Global RMSF r
0.31
0.33
0.21
0.67
0.57
0,50
0.69
0.71
Per-target RMSF
0.76
0.74
0.52
0.85
0.71
0.72
0.89
0.90
Root mean Wisdist.
3.60
4.23
3.58
2.76
4.32
3.22
2.58
2.43
+ Trans, contrib. +
3.13
3.75
2.86
2.23
4.04
2.47
2.15
2.03
Distributional accuracy
+ Var. contrib. /
1.74
1.90
2.27
1,40
1.77
1.89
1.28
1.20
MD PCA Wy-dist. J
1.51
1.87
1.99
1.44
1.97
1.78
1.29
1.19
Joint PCA Wi-dist.
3.19
3.79
1.58
2.25
3.98
2,47
2.13
2.04
PC-sim > 0.5 f
26
33
23
28
42
44
Weak contacts J
0.55
0.48
0.23
0.59
0.33
0.23
0.63
0.65
Ensemble observables
Transient contacts J f
0.34
0.30
0.28
0.36
0.25
0.43
0.41
Exposed residue J
0.49
0.43
0.32
0,50
0,29
0.67
0.70
Exposed MI matrix f
0.20
0.16
0.02
0.24
0.07
0.01
0.35
0.36
C. Implementation Details
Rigid align. Algorithm 3 describes the rigid alignment procedure used for atomic protein structures during loss calculation, which follows the Kabsch algorithm with reflection consideration. Notably, the structural alignment is in practice performed by transforming the ground truth structure to match the predicted structure, allowing gradients to propagate correctly during optimization, i.e., re x
Diffusion inference. During inference, Algorithm 2 outlines the diffusion module sampling procedure, which iteratively refines an initial conformation drawn from a Gaussian distribution. At each step, a randomly sample rigid transformation applied, followed by the standard denoising step using a learned diffusion module. In Algorithm 2, / is the Gaussian distribution; the Uso(3) indicates the uniformly distribution rotations in three dimensions and implemented using the scipy. spatissasemm... i al rans form. Rotatittiioooo iintatiio.. random. We use "o 0.8, Tmin 1.0, Natep 20, 1.75, 1.25 as the default hyperparameters for sample diffusion across all finetuned and aligned models; for pre-traines model, keep the default À= = 1.5, 1.5 since we find that the change of scaling factors will yield significantly worse results for diffusion module without fine-tuning. The inference noise scheduler also has the same configuration as Abramson et al.
(2024), i.e., Smaz = 160.0, 8min x 10-1 = 7 and data 16.
Training data. For training, we follow the dataset splitting of Jing et al. (2024a) for ATLAS. Specifically, we download the ATLAS MD trajectories, which comprises 1.390 proteins selected for structural diversity based on ECOD domain classification. This results in train validation test splits of 1,266 39 82 MD ensembles, with the rest excluded due excessive sequence length (Jing al., 2024a) The conformation stratified sampled from the production trajectories with an interval as Iss (per 100 frames in the original saving frequency). We combine the three replicas as the final ensemble with 300 structures per target. The multiple sequence alignments (MSA) for each target are downloaded from the pre-computed deposit in Protenix training data (Chen et al., 2018). . For the energy-based alignment fine-tuning, we employ an off- off-policy strategy where energy labels are assigned by scanning the entire training dataset using a force field prior to loading the data for model training. This allows us to efficiently decouple data processing from model updates to accelerate training.
14

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
Algorithm 2 Inference of Diffusion Module (Algo. 18 in Abramson et al. (2024))
Require: Features f embeddings jiputs s , , 2imi noise schedule [co, CI..... .r.
Require: Hyperparameter: = 0.8, "miin 1.0, noise scale X = 1.5, step scale 1/ = 1.25
1: X ~ co (N Ia) // Initialization
2: for & E Ci,..... cr do
3: XI f- x Center coordinates
4: R ~ ~'s00(3) ~ (0, Iz)
5:
R. x: + t, vl // Apply roto-translate
6:
7 To if c Ymin else 0
7:
[ t G--1 + +)
8:
82 11 . e + \v
9:
moisy
t
10:
x dennised t DiffusionModull((/)) moisy f, imputs stthh zamb f*)
6, + (x!tity muisy x xdenoised) //
dl C. -i
hoisy Xi + 1 dt 6
13:
14: end for
15: Return X]
Algorithm 3 Structure Rigid Align (Kabsch-Umeyaa Algorithm)
Require: X1, x
1: μ. pf + Σ, x1:
E1 x
2: X1, vet
xi + X1 - 11, tel
nd
3: U,V +- -trrch.ss svd ) x 8 x
4: B-+UV UV
5: torcch
linalg........ det (R) then
1
6:
F+
0
0
0
-1
7: UFVV
8: end if
9: x align
Rxi + 1
10: Return x
align
15

---

Aligning Protein Conformation Ensemble Generation with Physical Feedback
Alternatively, an on-policy strategy could be used, where conformations are generated dynamically during training, with energy calculated on the fly. Such approach would necessitate maintaining replay buffer store and reuse past samples, enabling more adaptive training based on real-time model predictions. We leave this exploration for future work.
Backhone model. We adopt PROTENIX (Chen et al., 2025), a PyTorch implementation of AlphaFold3, as the backbone architecture. For all ATLAS training targets, we pre-computed the PairFormer embeddings (a collection of s s_input, using Neycle = 10. During training and inference, we activate only the input embedder and diffusion module while leaving other components, such as the MSA module, PairFormer, Distogram head and confidence heads frozen. Additionally, no template module or template-base -batee input is used in our study. During training, the atom coordinates are properly permuted to match the best alignment with ground truth structure (Chen al., 2025).
8 s.trunk, z.trunk,
Supervised fine-tuning (SFT). Before performing EBA alignment, the pre-trained diffusion model is first fine-tuned using the standard score matching objective (Ho et al., 2020). The parameter optimization is performed with the Adam optimizer (Kingma& & Ba, 2014), using learning rate of 0.001, / values of (0.9, 0.95), and a weight decay of 1 x 10-8. The learning rate follows an ExponentialLR schedule with a warm-up phase of 200 steps and a decay factor y of 0.95 applied every 50k optimizer steps. We set AMSE = 1.0 and ALDD = 1.0 in Eq. (14). During training the noise level is sampled from datu 1.2+1.5- as the default setting with data = 16. In each optimizer step, we clip the gradient norm by 10. The SFT process consists of two stages: in the first stage, input structures with more than 384 residues are randomly cropped to a fixed size of 384. For cropping, half of the time, contiguous (on sequence) cropping used, while the other half employs the spatial cropping (Abramson et al., 2024). Random rigid augmentation is applied during diffusion training with an internal diffusion batch size of 32. In the second stage, the cropping size is increased to 768, and random rigid augmentation is applied with a reduced internal batch size of 16. In both stages, conformation samples are uniformly drawn from the training dataset without any sample weighting. The training was conducted with NVIDIA A A100 GPUs.
Alignment fine-tuning, As an alignment baseline, the EBA-DPO is implemented within the EBA framework by setting K = and replacing the softmax with sigmoid function as in Wallace et al. (2024). The binary preference label is annotated by selecting the sample with smaller energy in the mini-batch as "win" while the other as "lose". The reference model is selected be the SFT model with all parameters frozen. For EBA, we follow the same optimizer and scheduler as SFT stage but use smaller base learning rate of 1.0 10 T. To reduce the variance of gradient during training, we accumulate the gradient per 16 steps and also clip the norm by 10. We set the energy temperature factor B 1/5 where L is the protein length of the current mini-batch of samples, and the combined model temperature factor of = 50 The (internal) diffusion batch size is set to be 8 during alignment. To obtain a batch of sequence-coupled samples, we first sample single structure (c, x) uniformly from the alignment dataset, then we retrieve using the corresponding sequence e against the dataset for additional - 1 sample with the same sequence e (uniformly sample from the dataset). Similar to SFT, we set Amse = 1.0 and A 1.0 as the structural loss weights. Similarly, experiments were run on NVIDIA A100 GPUs.
Energy annotation. To evaluate the potential energy E of protein structures for preference fine-tuning, we adopted the OpenMM suite (Eastman et al., 2023) and CHARMM36 (Best al., 2012) force field with GBn2 model (Nguyen et al., 2013) for protein solvation. Each protein structure was fixed by the PDBFixer add hydrogen atoms and neutralized with Na Na+/CI- ions at a concentration of I 150mM. We then perform energy minimization with tolerance as 10 kJ /(mol nm) until converge. The energy (unit in kJ/ / mol) after minimization was assigned to each structure. To calculate the Boltzmann weight, we adopted kB- = 8.314 x 10-3 //mom - K) and set the temperature to be 300K (room temperature). Note that our energy evaluation does not perfectly reproduce the ATLAS simulation protocol, which uses GROMACS v2019.4 (Abraham et al., 2015) with the CHARMM36m force field (Huang et al., 2017) in an explicit TIP3P water environment. Since protein-solvent interactions can significantly contribute to the overall potential energy, using explieit water for energy annotation can introduce large variance in energy estimates and is far less efficiint in computation. Future work may consider quantum-leve level mniitt energy evaluation using density functional theory (DFT) methods, semi semi-empirical methods (such as GFN2-xTB (Bannwarth et al., 2019)) or even ab initio computational methods suitable for protein structure, which we believe is more promising since they are more accurate and less affected by solvent.
Runtime profiling, The MD ensemble generation runtime in Table I for basclines models (Jing et al., 2024a;b) and EBA are benchmarked on NVIDIA A100 GPUs The training iteration per step and memory profiling in Table 2 are calculated based on 4x A100 GPUs as wall clock time and per-device memeory consumption, respectively.
16


Reviews:
[{"review_id": "oeKrINzoba", "reviewer": "Reviewer_Gaxj", "summary": "This paper tackles the task of conformational ensemble generation in protein structure prediction. While folding models like AlphaFold predict individual states given a sequence, the task here is to be able to sample from the entire Boltzmann distribution instead. To this end, the authors propose an alignment scheme that leverages physical energies from molecular force fields to align the diffusion model component of modern folding models like AlphaFold3 to sample from the Boltzmann distribution. In particular, the authors introduce a scheme that effectively weighs different samples from the distribution according to their Boltzmann factors when fine-tuning the model, related to reward-based fine-tuning methods. The model is trained on molecular dynamics data that offers approximate samples from the Boltzmann distribution for training and compared to recent works such as AlphaFlow and MDGen, achieving favourable performance.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer Gaxj's detailed feedback and constructive criticism. Below, we address the main concerns raised in the review and clarify points of potential confusions.\n1. Clarification of the learning objective. The reviewer raises a point regarding the “approximation of the partition function Z” using a small set of samples (eg., K=2) and the potential biases introduced therein. There have been works exploring effective estimation of intractable partition function and the induced unbiased sampling; however, we do not claim such. We acknowledge that approximating Z with limited samples introduces large variance, but we emphasize that our method is motivated by/focuses on aligning with the relative Boltzmann factor (BF, p(xi)/p(xj)) between the off-policy conformations as the learning objective. Indeed, EBA can be viewed as an alignment “towards” the Boltzmann distribution while using partial information from mini-batched samples. Our finetuning objective is \"emulating\" the Boltzmann-weighted distributions rather than precisely sampling from the Boltzmann distribution, to (ideally) make the model situate between the pre-training and Boltzmann weight-aware. To further clarify this, we will include extra discussion in the revised manuscript regarding the impact of K, along with a more in-depth explanation of how our approximation aligns with other energy-based modeling approaches. \n2. Discussion on the sensitivity to K. The reviewer points out that the increased K which better approximates Z does not “consequently improve performance”. Our hypothesis is that the effectiveness of our method depends more on the distribution/quality of the off-policy structures rather than the size of the mini-batch. Firstly, K=2 does not imply there are totally 2 conformations sampled from the same target; different batches will roll over the acquired K samples and thus theoretically have access to each pair of p(xi)/p(xj). Secondly, we were adopting a simplified off-policy dataset from 100ns MD trajectories, where the dataloaders do not always yield a distinct variety of modes, since the structural fluctuations can be small for 100ns simulations. This may explain why performance does not always improve. To make readers aware of this concern, we shall expand our discussion on the interplay between K and performance, and such challenges probably anticipated in the on-policy setting.\n3. Clarification of $\\alpha$ notation in Derivations. We appreciate the feedback on the possible confusion of using α and provide an improving suggestion. Indeed, $\\alpha$ appears in multiple contexts (in background we simply borrow the notation from previous literature), leading to current ambiguity. To improve clarity, we will modify the notation as suggested:\nChange away the present α introduced above equation (6), making it explicit that the regularization $\\alpha$ in section 3.3 is the same as the one in background RLHF.\nClearly define $E_θ = -\\alpha \\log(p_θ/p_{ref})$ section 3.3 to 1. Simply the notation and 2. reinforce the connection to DPO modeling.\n4. Incremental gains of EBA Over EBA-DPO. The reviewer notes that EBA achieves on-par performance with EBA-DPO in table 1. We shall argue that EBA is designed to provide a more structured and interpretable framework rather than just maximizing out the performance in the given setting (Atlas MD data). We also discovered that small energy difference will lead to relatively large weight difference due to the exponential exp, thus making the weighted objective (EBA) resemble the EBA-DPO (binary, win-lose) in this case. More ideal tasks to show advantages of EBA can contain simpler systems with well-balanced energy differences. On the other hand, adapting DPO for protein conformation generation is also novel, and EBA provides a general framework and theoretical justification for what DPO optimizes in this context.\n5. Presentation and minor corrections (comments). We also appreciate the reviewer’s attention to detail and address the specific presentation issues. We have revised accordingly:\n- Correcting the equation in Line 090 to show the exact transition notation.\n- Clarifying the connection between $x_θ$ and $ϵ_θ$ in the background section.\n- Fixing the notation below equation (14) to ensure $x̂_0$ is properly represented.\n- Adding axis labels and annotations for Figures 3 and 4.\n6. Additional literature on energy-based models. We appreciate the suggestion to reference broader energy-based modeling literature, such as the mentioned (Song and Kingma, 2021). We will incorporate relevant citations and compare our approach with established energy-based training techniques, including those that require MCMC-based partition function estimation.\n\n\nFinally, we kindly thank the reviewer Gaxj for their insightful feedback, which has helped us refine both methodology and presentation. We will appreciate any further suggestions of revisions in our following discussion.", "claims_and_evidence": "The main claims made by the submission are generally supported through empirical evidence presented in the experiments.", "methods_and_evaluation_criteria": "The methods and evaluation criteria are generally appropriate and make sense for the problem at hand. I do have some concerns regarding the presentation and derivation of the newly proposed method, though. See below.", "theoretical_claims": "Yes, I checked the correctness of the proofs in the appendix. Everything looks correct to me.", "experimental_designs_or_analyses": "Experimental design and analyses overall seem appropriate and seem sound an valid. I do have some questions regarding the experimental results in an ablation experiments, though. See below.", "relation_to_broader_scientific_literature": "The key contributions of the paper are positioned appropriately with respect to the existing literature. In particular, the paper appropriately references and discusses the most important protein folding models, methods to perform ensemble generation instead, as well as the most important alignment and fine-tuning methods.", "essential_references_not_discussed": "I think all most essential papers are cited.\n\nHowever, when deriving its alignment scheme, the paper employs an energy-based formulation, essentially starting its derivation from a forward KL minimization (equation (6)). The method then approximates the partition function through a small set of empirical samples, if I am understanding the paper correctly (in practice, it is merely 2 samples in the main experiments). This raises questions, which I will discuss below. With that in mind, it would be appropriate here to cite works from the broader literature on energy-based models and position the method appropriately, contrasting it to other approaches. See, for instance, Song and Kingma, *How to Train Your Energy-Based Models* https://arxiv.org/abs/2101.03288, 2021, for an overview and further citations therein to more methods (contrastive divergence, noise contrastive estimation, Stein discrepancy, etc.). Among other things, these methods usually require a careful handling of the partition function, which in practice can require the need for MCMC and similar methods, but this is not the case in the submission here.", "other_strengths_and_weaknesses": "**Strengths:**\n- The empirical results in the main table, although somewhat incremental, show the proposed method's advantage.\n- Generally, the approach to incorporate explicit physical energies based on atomistic force fields when doing generative modeling of conformational ensembles makes sense to me, so I think this is a worthwhile direction.\n\n**Weaknesses:**\n- In their derivations, the authors simply approximate the partition function through a small set of samples from the training data (a mini batch, or just two samples, i.e. K=2). I think this is a very, very coarse approximation of the true partition function. However, I think this not only introduces variance in the partition function estimate, but also a strong bias relative to what one would obtain with the ground-truth partition function. Usually, when training energy-based models, it is this partition function that leads to many complications. This approximation and the potential bias is not discussed at all and completely ignored. I would like the authors to comment on that.\n- One would think that using more samples, i.e. a higher K, to approximate the partition function would be more accurate and consequently lead to better results. But this does not seem to be the case. The ablation study seems to indicate that the method is not very sensitive to K and that for some metrics the results actually get worse. Again, this raises questions, but is not discussed at all.\n- The way the method is derived is confusing: In the background section, in equation (2), the authors first introduce $\\alpha$ as the regularization strength in RLHF. Then above equation (6), $\\alpha$ is used again, but this time as a scaling of the energy, which is different from the previous $\\alpha$ and in fact would not be necessary at all, as it could just be absorbed into the energy function definition. Moreover, after equation (8), the $\\alpha$ effectively cancels out and there is no $\\alpha$ anymore in equation (12). Then, in section 3.3 the $\\alpha$ re-appears, and now it is reinterpreted as the regularization-$\\alpha$ from RLHF. This approach is not incorrect, but it is confusing. It would make a lot more sense to not have the second $\\alpha$ introduced above equation (6) at all, and instead in section 3.3 when making the connection to DPO to define $E_\\theta=-\\log (\\frac{p_\\theta}{p_{ref}})^\\alpha$. This makes intuitively much more sense, because now it is clear that this is again the regularization $\\alpha$ that modulates the fine-tuned $p_\\theta$ with respect to $p_{ref}$. Overall, I thought the derivation, although mathematically correct, was quite confusing.\n- In the main experiment table, it seems that EBA-DPO performs almost on-par with EBA. Hence, the advantage of the full EBA framework over regular DPO is very incremental.\n\nThe above concerns do impact the paper's clarity and significance, unfortunately.", "other_comments_or_suggestions": "- Line 090: The equation seems incorrect. For $q$, the authors have the transition from $x_{t-1}$ to $x_t$, but then they write the Normal kernel that goes from clean data all the way to $x_t$.\n- For consistency, the authors should make the connection between $x_\\theta$ and $\\epsilon_\\theta$ in the background section.\n- Below equation (14), I assume it should be *predicted coordinates* $\\hat{x}_0$ with the hat?\n- Figure 3 is missing the x axis description and annotation.\n- Figure 4, last 2 plots, are missing the y axis description.", "questions_for_authors": "I do not have any further questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nn4QHc41Dr", "reviewer": "Reviewer_BE4T", "summary": "This work concerns the problem of improving diffusion models for protein conformation generation using the information from a physical model. It proposes a new fine-tuning loss EBA for diffusion model training, based on the principle of offline RL using energy labels as the a negative reward feedback. \n\nBy balancing the diffusion loss for different conformations of a protein during mini-batch training, it aims to train the model to better reflect the thermodynamics of the protein. In addition, the authors show that the popular DPO objective is a special case of EBA, by using a mini-batch of 2 and setting an infinitely low temperature. Through the evaluation on the standard benchmark of Atlas, they shows that models fine-tuned using EBA objective achieve better performance in capturing the flexibility of the protein and predicting the dynamics-related MD observables.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the detailed and constructive feedback from the reviewer BE4T! We will address each question/concern as follows:\n\n1. Target distribution and bias away. Indeed, given sufficient unbiased MD data, directly maximizing the likelihood (MLE) is ready to converge to the Boltzmann distribution (Noé 2019, Arts 2023). However, we note that the ATLAS MD dataset, containing \"collections\" of 3 parallel trajectories per target, may introduce bias when using MLE alone. In this context, EBA adjusts for potential bias by incorporating partial information from the Boltzmann factor.  \nFurthermore, our objective function is unbiased towards $p(x)$ from the following sense: one can prove if there exists a distribution $q_\\theta(x) = \\exp(−β’ E_\\theta (x)) / Z’$ such that $q(x_i)/q(x_j) = p(x_i)/p(x_j) \\forall i, j$, then $p(x)$ equals to $q(x)$ for all x. This implies that optimizing w.r.t. the Boltzmann factor converges towards the target Boltzmann distribution. We will add this proof to the next revision to inspire better insights.\n\n2. Eq. (11, 12) notation of T. It is a result of taking expectation of $t \\sim U(0, T)$. During derivation, we make the original sum over diffusion time step $\\sum_{t=0}^T$ in ELBO as $T * [(1/T) * \\sum_{t=0}^T]$ and extract the second part [**] out using the Jensen Inequality to be there. \n\n3. Base performance before EBA. This is indeed an important ablation study. Below is the evaluation results of the SFT checkpoint before EBA. We will add this into the next revision:\n|Name| Pairwise RMSD r | Global RMSF r | Per target RMSF r | RMWD | RMWD trans | RMWD var | MD PCA W2 | Joint PCA W2 | PC sim > 0.5 % | Weak contacts J | Transient contacts J | Exposed residue J | Exposed MI rho |\n|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n|Stage 2-SFT| 0.57|0.69|0.89|2.58|2.15|1.28|1.29|2.13|42|0.63|0.43|0.67|0.35|\n\n4. Selecting K value. We note that using K samples in the EBA objective is not intended to estimate the partition function and likely cannot do so accurately. Instead, the EBA objective aligns with Boltzmann factors, i.e., the relative weights between conformations. Since these relative weights are invariant to the number of samples used, increasing K does not necessarily improve performance. \n\n5. Treatments of potential energy E(x). We appreciate the insightful questions which we are happy to discuss as follows: \n\n(1) This is due to the change of scale of energy. Traditional molecular mechanism force field (MMFF) is defined by enumerating the interactions in the molecule. The increase of #atom (or protein length) will quadratically increase the pairwise terms. In large proteins, the scale of E(x) is much larger, such that the energy difference [E(xi) - E(xj)] is linearly larger, which will make the learning difficult for large systems since the energy softmax can spike instead of being balanced. \n\n(2) [i] Correct, following the point above, the introduction of length-factor implicitly changes the temperature of different proteins, thus calibrating the coefficient. The potential issue can be a push-away from the target relative weight between conformations yet the order of them is kept therein. Thus, EBA is inducing a much “flatter” distribution for large proteins. \n\n[ii] the folding time can be hypothesized as a result of free energy barriers. Inspired by their empirical findings (Naganathan, 2005), we thus applied the L^0.5 and found it empirically worked out fine. We cite this to make an analogy instead of being a rigorous scaling rule.\n\n(3) In fact, it is the problem of “energy variance” that motivates us to use the off-policy ensembles where the conformations embody smaller variance (due to being sampled from MD). As a result, the range of energy difference is small and the softmax weights stay within (1e-2, 1). In early experiments, we found that the softmax weights will spike (as small as 1e-7 vs 1.0) if we do not introduce the length factor, and even worse when we used on-policy data (sampled using pre-trained AF3).\n\n6. Merged hyperparameters. We are sorry if this introduces ambiguity. In fact, this trick is inspired by the diffusion-DPO where the authors did the same combination. In practice, we only tune $\\alpha * T$ as an independent variable, keeping the sampling diffusion step as T=20 for all configs.\n\n7. Structural aligned loss. The rationale of using these losses come from the original training of AF3 ’s diffusion module, and we mostly follow their loss definition. The consequence is that the space of protein structure is not a quotient space R^{3N} / SE(3) according to AlphaFlow paper section 3.2 (Jing, 2024). We have not yet abated on the loss choices and will clarify this in the paper.\n\n8. Additional related works, notation/typo, and formatting issues (comments): We are grateful for pointing out these valuable suggestions. We appreciate the attention to detail and will dedicate the discussion and the corrected issues into our next revision.", "claims_and_evidence": "The main claims made by the authors includes: \n\n1. “Our approach bridges the gap between purely data-driven conditional generation and physics-based simulations by incorporating fine-grained force field feedback via a scalable learning objective”\n2. “Through the EBA alignment, the diffusion model learns to balance different conformational states, resulting in more physically consistent protein conformation ensembles”\n3. “the EBA-aligned diffusion model achieves state-of-the-art performance compare to previous generative models by incorporating physical feedback”\n4. “the proposed method provides a novel pathway for integrating generative modeling with transitional simulations”\n\nWhile claim 1 and 4 are supported by the construction of the problem, questions remain for claim 2 on how does the learned diffusion model align with the target Boltzmann distribution (Q1); for claim 3 on the practical improvement of using EBA over SFT-fine-tuned model (Q3).", "methods_and_evaluation_criteria": "See Q1, Q2, Q5 regarding the questions on the proposed methods.\n\nThe evaluation criteria is standard in the field and is appropriate for the proposed problem.", "theoretical_claims": "I read through the derivation and proofs of the paper. Most of the derivations are standard and resembles those in DPO and Diffusion-DPO (Wallace et al). See Q2 for a specific question.", "experimental_designs_or_analyses": "As proposed method is a new objectives for diffusion fine-tuning, one key comparison should be proposed method vs SFT, which is missing from the main result (see Q3).", "relation_to_broader_scientific_literature": "The general problem of this work is to improve a pre-trained diffusion model with an explicit reward feedback (i.e., energy), with a specific focus on distributional matching. This general problem was discussed in a recent paper Lu et al (ICML 2023, https://arxiv.org/abs/2304.12824) with an extension to protein conformation (Wang et al, ICML 2024, https://arxiv.org/abs/2403.14088). However, above work approaches the problem from the perspective of guidance while this work attacks problem by forming a new fine-tuning objective.", "essential_references_not_discussed": "Recent efforts have attempted to incorporate physical information during the training and sampling of diffusion models to improve the protein conformation generation. For example: \n\nConfDiff (ICML 2024, https://arxiv.org/abs/2403.14088) trains plug-in guidance models to integrate energy and force information during the reverse sampling. It also targets the energy-tilted distribution (eq 3 in their paper) and connects adjusted reverse diffusion with the potential energy and force of conformations (eq 4, eq 5 in their paper).\n\nBioEmu (preprint, https://www.biorxiv.org/content/10.1101/2024.12.05.626885v1) proposes to fine-tune pre-trained diffusion models with the property prediction fine-tuning, integrating the energy difference of pairs of conformation states (Section S3.6 in their paper).\n\nThese works should be discussed in order to better understand the current approach.", "other_strengths_and_weaknesses": "- Strength: this work shows a new training objective to integrate protein conformation energy for diffusion model training, using AlphaFold3 as a modern all-atom framework, and empirical results show strong performance on Atlas benchmark.\n- Weakness: my main concerns originate from some specific details of the model (see questions). In addition, there is an excessive mixed use of notations (see Other Comments)", "other_comments_or_suggestions": "There are several notation issues and typos in the current manuscript:\n\n- Equation (1): what is $\\lambda_t$  in $\\omega(\\lambda_t)$?\n\n- In Equation (2), using $\\theta$ for both $r_\\theta$ and $p_\\theta$ can be misleading that i) two models share parameters and ii) reward model is also optimized in equation (2)\n\n- Equation (6):\n  - There is no superscript \"i\" in first summation and there is a missing “i” in $\\log(p_\\theta(x|c))$\n  - The summation is a Monte Carlo estimation of the KL divergence, yet a equal-by-define $\\triangleq$ instead of approximation $\\approx$\n  - In Line 162: used summation over $x^i$\n\n- Typo: Line 246 right column: “predicted coordinates $\\mathbf{x}_0$”\n- Typo: Line 608: mixed use of index i and j\n- Equation 17, additional comma ‘,’ in the denominator\n- In Section A.1.: are $\\forall j$ and $\\forall i$ necessary in some of equations?", "questions_for_authors": "1. Can you clarify which target distribution the model learns to match for the general EBA objective (Equation 7) and for the model trained in Table 1? Based on the model, it should match to a weighted distribution of data $\\propto p_\\text{data}(x)\\exp(-\\beta E(x))$, however, the data from MD simulation already follows the Boltzmann distribution of the potential energy $\\mathbf{x}_\\text{data} \\propto \\exp(-\\beta E(x))$. Imposing this additional energy weights seems to bias the target distribution away from the Boltzmann distribution. Can you provide further clarification?\n2. In equation 11 and 12, why is the total number of diffusion steps (T) in the exponential part of the denominator instead of before the summation? \n3. What is the base performance for SFT model, after the training stage 1?\n4. At first, It appears that a large K value is required for an accurate Monte Carlo estimation of the partition function in Equation 7; however, the ablation study in Table 2 suggests that the results are not sensitive to the choice of K. Can you explain the principle behind selecting an appropriate K value?\n5. I have several questions regarding the treatment of potential energy E(x):\n    1. The authors mentioned that E(x) varies significantly with protein size, introducing large variance in the objective. Can you clarify why this is an issue, given that all data are normalized by mini-batch Boltzmann weights (Algorithm 1, step 5) and are from the same protein (i.e., same protein size)? \n    2. By using length dependent normalization, the proposed method appears to impose different temperatures for different proteins. Can you clarify if this is the case, and if so, what potential issues may arise from this treatment? Additionally, can you explain why “the scaling of folding time with respect to the length” translates into the rule of scaling the potential energy with the length? \n    3. Large energy variance can still exist within different conformations of the same protein and cause the softmax operation in Algorithm 1, step 5 to collapse  — only one or a few data have probability ~1 and others to be 0. Did you check the mini-batch Boltzmann weights across different proteins? \n6. In Appendix A.2, the authors “combine alpha and T” into a single hyperparameter, presumedly the “model factor T” mentioned in Appendix B (Alignment fine-tuning). This combination is somewhat confusing. For exmple, when varying the “model factor T”, did the authors also change the number of diffusion steps (the original T)?\n7. While using structural losses specific to protein (aligned MSE and smooth LDDT) can be a reasonable alternative for training in practice, the reasoning is not entirely convincing: the use of MSE error of predicted noise is derived from the KL divergence of stepwise reverse Gaussian steps whereas the structural losses may not. Is there any approximation introduced by replacing with proposed structural losses? Did you conduct empirical studies on choosing different losses?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "59QCmX0Lwc", "reviewer": "Reviewer_yVMH", "summary": "This work presents Energy-based Alignment (EBA), a method to fine-tune a pretrained diffusion model to sample mini-batches of protein structures that match the underlying Boltzmann distribution.  This alignment is achieved by minimizing the cross entropy between the ground truth Boltzmann distribution, and the distribution approximated by the model, evaluated on the mini batches of size K. Authors provide a theoretical derivation of the exact form of the cross entropy for the diffusion model. Additionally, authors show that EBA with K=2, in the low temperature limit, can be formulated as a Direct Preference Optimisation problem (EBA-DPO). \n\nAuthors fine-tune AlphaFold 3 (Protenix) on the ATLAS dataset. To compare against baselines, authors use metrics from Jing et al. 2024. Both EBA and EBA-DPO outperform baselines across almost all metrics.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We kindly appreciate the reviewer yVMH’s supportive and insightful feedback. Below, we address each of the key questions raised and provide clarifications where necessary.\n1. Clarifications on theoretical derivation. We acknowledge that the substitution of the denoising loss into the KL divergence (Line 185) was not made fully explicit. In the revised manuscript, we will provide a more detailed explanation of this substitution and its connection to ELBO. Specifically, we will highlight the derivation steps that lead to the energy-based alignment (EBA) objective and its direct connection to narrow the gap towards the Boltzmann distribution.\n2. Training details and discussion. Thanks for pointing out the suggestion for clarification. \n\n[Training details] In the current paper version, we actually postponed the implementation details in the appendix to save space in the main text. To clarify: we used Adam as the optimizer with a learning rate of $10^{-3}$ during SFT and $10^{-7}$ during EBA, $\\beta$ values of (0.9, 0.95), and a weight decay of $1×10^{−8}$. The learning rate follows an ExponentialLR schedule with a warm-up phase of 200 steps and a decay factor $\\gamma$ of 0.95 applied every 50k optimizer steps. Gradient norm clipping was applied to each step by 10; \n\n[Line 338] The training instabilities were primarily observed early that when we did not impose the length factor $1/L^{0.5}$, the softmax of energy would spike. This is due to the change of scale of energy. Traditional molecular mechanism force field (MMFF) is defined by enumerating the interactions existing in the molecule. The increase of atom number (or protein length L) will (quadratically) also increase the pairwise terms. One can imagine that in large proteins, the energy scale E(x) is much larger, such that the energy difference [E(xi) - E(xj)] is thus linearly larger, which makes the learning difficult for large systems since the energy softmax can spike instead of being balanced between conformational states. \n\nWe will expand on these points and organically incorporate them into the main text in the revised version to provide more transparency on the training process.\n\n3. Discussion on prior works and clarification on our novel contributions. Thanks for the good suggestions. We acknowledge that Boltzmann generators (BG, Noé et al., 2019) is a very pioneering work in this research direction. While our approach differs here in using diffusion models and model alignment for fine-tuning, we agree that the core principle of learning towards the Boltzmann weighted ensembles is shared. We were so inspired by many prior works including BG, RLHF, DPO, AlphaFlow to name a few, and will dedicate a discussion paragraph in the introduction section to motivate the EBA more explicitly. \n4. Tentative summary of novel contributions. We acknowledge that the methodology section does not clearly separate our novel contributions from the prior works. As for our contributions, we showcase an successful application of fine-tuning AF3 for MD ensemble generation. Also: \n- we shall explicitly state that EBA introduces a new physical-inspired fine-tuning objective based on energies for diffusion models, aligning generated distributions with the Boltzmann distribution.\n- we will clarify that the connection to DPO is novel, as previous works have not formulated preference optimization within a Boltzmann distribution motivated framework.\n- the use of mini-batch Boltzmann weighting for preference optimization is also a novel aspect of our work to enable scalable training of models such as AF3. \n\nIn conclusion, we will re-structure the methodology section to better highlight these contributions and make clear which points are inspired by the pioneering works.\n5. Responses to the specific questions:\n\n[**Q1**] Smooth LDDT Loss: The smooth LDDT loss is inspired from structural alignment objectives in training the diffusion module of AlphaFold3, and is so used for diffusion fine-tuning in our work. We will clarify this in the manuscript. \n\n[**Q2**] Retrieving K-1 Samples: K-1 indicates the extra samples additional to the original training sample i. In this sense, K=2 means we would retrieve 1 extra related sample to compute the energy weight; Yes, this notation simply denotes enumeration over samples. We will revise this for clarity.\n\n[**Q3**]Learning rate eta and loss weight lambda: We are sorry for abusing the usage of these letters, they do have different meanings: in the Figure 4, these symbols are introduced from the original AF3; while in the algorithm 1 they are learning rate and loss weight. Subsequently, we will rename them and clarify this explicitly in the text.\n\nWe thank the reviewer yVMH for their positive evaluation of our work and for recognizing the contributions of EBA. We will incorporate these suggestions to improve clarity, expand discussion on related work, and provide additional experimental details.", "claims_and_evidence": "Authors claim in the abstract that “Experimental results on the MD ensemble benchmark demonstrate that EBA achieves state-of-the-art performance in generating high-quality protein ensembles”, and that indeed seems to be supported by the experimental evaluation. AlphaFlow-MD, MSA subsampling and MDGen are strong baselines that EBA outperforms.", "methods_and_evaluation_criteria": "ATLAS dataset is a commonly used dataset for this task, and the chosen baselines make sense. The generation success is assessed using metrics commonly used in similar works.", "theoretical_claims": "I went through derivations in the main body, but didn’t check the Appendix. Authors optimise ELBO using previously derived results, however the substitution of the denoising loss into KL divergence (line 185) should be made more explicit.", "experimental_designs_or_analyses": "Using AlphaFold 3 (via Protenix) is a good choice of a model for the task at hand. It is not clear what was the training (fine-tuning) procedure (for example, optimizers, or expanded discussion of training instabilities from line 338).", "relation_to_broader_scientific_literature": "The problem that this work is addressing, namely sampling from the ensemble that obeys the ground truth Boltzmann distribution, is a known and important problem in structural biology, see for example a review by Henin et al. 2022 (https://arxiv.org/pdf/2202.04164)", "essential_references_not_discussed": "While the presented methodology is original, it is based on previous findings that are only shortly discussed. Boltzmann generators from Noe et al. 2019 (https://arxiv.org/abs/1812.01729) are based on the similar insights about matching equilibrium distributions, and if they are the inspiration for this work, they should receive more attention.", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper is generally well written and has a solid evaluation. Including the pre-trained version of the model, but before EBA fine-tuning, shows evidence that EBA indeed results in improved performance.\n\nWeaknesses:\n\nAuthors don’t discuss properly in the Methodology Section what is their own novel idea, and what is inspired by other works.", "other_comments_or_suggestions": "In Section 2, where DPO is introduced for the first time, ${x^w}$ and ${x^l}$ are not explained. It would improve reading of the work if more intuition about BT model was given.\n\nIn Algorithm 1, step 10, the comment is maybe a bit close to the text and looks like a division.\n\nSmall typo: line 109, ‘that is ,’ instead of ‘that is, ‘\n\nMain body, rather than the Appendix, could mention how the energy function ${E}$ is chosen.", "questions_for_authors": "1.\tIs the smooth LDDT loss an entirely novel thing introduced here, or it is taken from other works?\n\n2.\tIn Algorithm 1, in step 3., why $K-1$ samples are retrieved? Does $c^j = c(j=…)$ simply mean that enumeration?\n\n3.\tFrom reading of the manuscript it seems that eta is the learning rate, and lambda is the loss weight. Are those the same eta and lambda as in the Figure 4?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qypE0bv4cx", "reviewer": "Reviewer_PSvo", "summary": "his paper introduces Energy-based Alignment (EBA), a novel approach that enhances generative models by incorporating feedback from physical models. EBA efficiently calibrates these models to balance conformational states based on energy differences, overcoming intractable optimization issues.\n\nTested on the MD ensemble benchmark, EBA achieves state-of-the-art performance in generating high-quality protein ensembles. By enhancing the physical plausibility of predicted structures, EBA improves model accuracy and shows potential for advancing structural biology and drug discovery applications.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the thoughtful feedback and helpful suggestions from the reviewer PSvo. We here address each point raised and clarify aspects of our work where necessary.\n1. Additional baselines. We acknowledge the importance of comprehensive baseline comparisons. In our study, we primarily focused on methods that directly optimize for MD ensembles with specific parameters (temperature, force field, etc.). Though the mentioned methods are all pioneering and relevant works, most of them are not trained on the Atlas data so we think it is not very suitable to compare with. Yet we are happy to include these baselines (with necessary MD fine-tuning) and set up a benchmark in the near future. We here complement the most recent two, Confdiff and BioEmu to the comparison:\n\n|Baseline| Pairwise RMSD r | Global RMSF r | Per target RMSF r | RMWD | RMWD trans | RMWD var | MD PCA W2 | Joint PCA W2 | PC sim > 0.5 \\% | Weak contacts J | Transient contacts J | Exposed residue J | Exposed MI matrix rho |\n|----|----|----|----|----|----|----|----|----|----|----|----|----|----|\n| BioEmu | 0.46 |0.57 | 0.71  | 4.32 | 4.04 | 1.77 | 1.97 | 3.98 | 51 | 0.33 | - | - | 0.07 |\n|ConfDiff (OpenFold-r3-MD)| 0.59 | 0.67 |0.85|2.76|2.23|1.40|1.44|2.25| 35 |0.59|0.36| 0.5 | 0.24 |\n|EBA(Ours)| 0.62|0.71|0.90|2.43|2.03|1.20|1.19|2.04|44|0.65|0.41|0.70|0.36|\n\nWe shall add these results to the paper after revision, thanks!\n\n2. Related work section on *Generative AI for Protein Conformation Generation*. We are sorry if we don’t make it clear. The current manuscript does contain a paragraph for protein conformation generation in Section 5-Part 1, which includes several related works but may not sufficient. To provide a better context, we will make a more dedicated discussion on this topic in the revision, especially with the aforementioned baselines.\n3. **Questions**. We kindly appreciate each of the reviewer's insightful questions, and we answer each as follows:\n\n**Q1**: By scaling up to larger systems such as membrane proteins or IDP, the network may suffer from capacity issues and lack of accuracy. 1. The training step (of AF2, AF3-like) is based on a cropped region of original structure (eg. 384, 768) which may be difficult in accurately predicting the distribution of larger protein. 2. AF3 is pretrained on structural proteins mostly and thus the sampled ensemble has limited diversity; the distribution of IDP is more flat in 3D space so the model should be finetuned to be able to cover such “free” ensembles.\n\n**Q2**: Sure. We conduct a simple benchmarking on different K (2,3,5) regarding training speed and GPU memory utilization. We use 4xA100 for profiling:\n\n| K | Iteration per step (sec) | Avg memory consumed in MiB (DDP, per device) |\n|--------|----------------------------------|---------------------------------|\n| 2  | 4.3                          | 12288.0                              |\n| 3  | 5.4                          | 14254.1                              |\n| 5  | 7.8                          | 16711.7                              |\n\nAs we adopt the backbone of Protenix, the training of larger system (eg. >10k atoms or ~1k tokens) will be trained with a random cropping at a feasible training size (Eg. 384) on common A100 GPUs. In our opinion, it is doable in scaling up, while we acknowledge the challenge in accuracy for longer protein may exist due to the cropping, which is the shared issue for all structure prediction models.\n\n**Q3**: Yes. For the quantum-level, we plan to first experiment with the small molecule/peptide dataset. For example, we plan to curate a molecular conformation dataset annotated by the gfn2-xtb, a well-known semi-empirical method. For the protein, the quantum calculation is not computationally feasible; we may consider the estimated potential such as the one in the AI2BMD paper. While compared with quantum-level methods would provide additional insights, such approaches are computationally expensive and are generally limited to small peptides rather than protein targets from ATLAS at similar time scales. We leave for future exploration.\n\n**Q4**: Yes, in fact extending to multi-chain (complex) is right on our way. We are actively experimenting to extend the EBA to the protein-ligand simulation data (specifically, the MISATO dataset (Siebenmorgen et al., 2024)). Since the AF3/Protenix architecture by design can handle complex input, we think it is ready w.r.t. model architecture. We also anticipate that this will require modifications to the energy function to account for the inter-chain interactions.\n\nWe sincerely appreciate the insightful comments and questions, which have helped us identify areas for clarification and improvement. We hope that these revisions will address your concerns and further strengthen our manuscript.", "claims_and_evidence": "Introduction of Energy-based Alignment (EBA):\n\nClaim: The paper introduces the EBA method to enhance protein conformation ensemble generation by integrating physical energy feedback with generative models.  \nEvidence: The authors compare their method with traditional molecular dynamics (MD) simulations and generative models, demonstrating that EBA improves the physical plausibility of generated structures.\nValidation of EBA's Performance:\n\nClaim: The EBA-aligned model outperforms existing generative models in terms of generating high-quality protein ensembles.  \nEvidence: Experimental results on the ATLAS MD ensemble benchmark show superior performance in metrics like RMSD, RMSF, and exposed residue similarity compared to other methods (e.g., AlphaFold, AlphaFlow, and MDGen). The paper provides a table with detailed statistical metrics to support these claims​\n.\nTheoretical Justification of EBA:\n\nClaim: EBA effectively aligns the model’s learned distribution with the Boltzmann distribution, ensuring better alignment with physical laws governing protein stability.  \nEvidence: The derivation of the EBA objective and its mathematical formulation (e.g., Equation 6 and 7) demonstrate how the method avoids the intractable partition function while still maintaining energy consistency across conformations​\n.\nAblation Study on Hyperparameters:\n\nClaim: The model's performance is robust across different values of the mini-batch size (K) during fine-tuning.  \nEvidence: The paper presents an ablation study showing consistent results for various K values, with only a slight increase in computational overhead for larger values​\n.\nLimitations:\n\nClaim: The method has limitations, such as challenges with modeling long-time scale dynamics and the accuracy of the energy functions.  \nEvidence: The authors acknowledge that while the method works well for single-chain protein ensembles, further improvements are needed, especially in terms of force field precision and extending the method to more complex biomolecular systems​", "methods_and_evaluation_criteria": "Yes, the methods and evaluation criteria makes sense to me.", "theoretical_claims": "Yes, the proofs are mainly correct.", "experimental_designs_or_analyses": "Comparison with Baseline Methods:\nExperimental Design: The authors compare their EBA-aligned model to existing protein structure generation methods, including AlphaFold2, AlphaFlow, and MDGen. These comparisons are made using the ATLAS MD ensemble benchmark, which involves generating protein ensembles and assessing their physical plausibility and other properties like RMSD, RMSF, and Jaccard similarity.  \nValidity Check: The use of established baselines and the clear presentation of comparative metrics (e.g., RMSD, RMSF, and other structural observables) helps establish the validity of their approach. The experimental results are robust, with the EBA-aligned model showing consistent improvements in key metrics like RMSF correlation and exposure of residues.  \nPotential Issue: The performance of baseline methods like AlphaFold2 and MDGen on different datasets or settings could impact the generalizability of the results. For example, if these baselines were trained or optimized differently, comparisons may not fully reflect the potential of the proposed method. This could be a source of bias.\n\nDataset and Training Pipeline:\nExperimental Design: The training pipeline involves fine-tuning a pre-trained denoising network on protein ensembles from the ATLAS dataset. This ensures that the model learns from simulation data that captures diverse conformational states.  \nValidity Check: Using the ATLAS dataset, which consists of MD simulation data with atomistic resolution, is a good way to validate the approach against realistic, high-quality data. The dataset also covers a range of protein types, which improves the robustness of the results.  \nPotential Issue: The authors acknowledge that the training data is limited to single-chain proteins. Extending the analysis to multi-chain or larger protein systems could reveal the limitations of the model. Furthermore, while ATLAS is a useful dataset, the findings may not fully translate to other datasets without similar structural diversity or complexity.", "relation_to_broader_scientific_literature": "Yes, key contributions of the paper rare elated to the broader scientific literature", "essential_references_not_discussed": "Yes, some critical related works are not discussed.\n\nPlease see: Protein Conformation Generation via Force-Guided SE(3) Diffusion Models (ICML24) and the related works part. \n\nIn the submission, the authors did not have a paragraph for related works on protein conformation generation. Moreover, the baselines are not compared.", "other_strengths_and_weaknesses": "A series of baseline methods are not compared, including but not limited toEigenFold, Str2Str, ConfDiff, BioEmu, and AI2BMD.", "other_comments_or_suggestions": "Please add one paragraph in the related work to discuss generative AI for conformation generation.", "questions_for_authors": "The evaluation was conducted on the ATLAS MD ensemble dataset. Have you tested the model on other protein datasets with different properties or larger systems, such as those involving membrane proteins or disordered proteins? If not, do you anticipate any challenges in transferring the model to these types of datasets?\n\nIn your ablation study, the results indicate that increasing the mini-batch size (K) does not drastically affect the performance but increases computational overhead. Can you quantify the computational trade-offs in terms of training time and hardware requirements for different K values? How scalable is the method for much larger datasets or systems with tens of thousands of atoms?\n\nWhile the paper shows promising results against empirical MD simulation data, how do you plan to validate your approach against quantum-level calculations, especially for proteins with highly intricate folding mechanisms? Are you planning any future benchmarks that involve comparing your method to quantum chemistry methods, such as those based on density functional theory (DFT)?\n\nYour current work is focused on generating ensembles for single-chain proteins. How do you plan to extend the model to handle multi-chain or large protein complexes? Are there any expected challenges or changes in the model architecture when scaling to multi-chain systems?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces Energy-Based Alignment (EBA), a method for fine-tuning diffusion models to generate protein conformational ensembles. By incorporating energy-based feedback from molecular simulations into the generative training process, the authors propose a way to bridge data-driven learning and classical physical modeling. Empirical evaluation shows improved performance over prior methods, such as AlphaFlow and MDGen.

The paper addresses an important and timely problem in the intersection of structural biology and generative modeling. The proposed method is well-motivated and represents a promising direction for incorporating physical priors into deep generative models. Several aspects of the paper, including the application to protein ensembles and the connection to preference optimization (EBA-DPO), were recognized as meaningful contributions.

On the flip side, the reviewers raised concerns about the clarity of the presentation. Specifically, the derivation of the EBA objective and the approximations involved in aligning with the Boltzmann distribution were found to be difficult to follow. Questions remained about the rigor of the theoretical framework, particularly regarding the approximation of the partition function and its implications for the learned distribution. Additionally, some reviewers felt that the empirical gains over standard fine-tuning were modest, and requested clearer baselines and more thorough comparisons.

Despite these reservations, the overall direction of the work is promising. The paper introduces a solid contribution to an important area and lays the groundwork for further advances. With additional work to improve clarity, expand theoretical justifications, and benchmark against more recent baselines, this line of research has the potential to make significant impact. On balance, I recommend acceptance.
Decision: Accept (poster)