Aligning Protein Conformation Ensemble Generation with Physical Feedback

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.

1. 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:

We shall add these results to the paper after revision, thanks!

2. 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.
3. Questions. We kindly appreciate each of the reviewer's insightful questions, and we answer each as follows:

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.

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:

As 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.

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.

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.

We 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.

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.

1. 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.
2. Training details and discussion. Thanks for pointing out the suggestion for clarification.

[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;

[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.

We 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.

3. 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.
4. 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:

• 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.
• we will clarify that the connection to DPO is novel, as previous works have not formulated preference optimization within a Boltzmann distribution motivated framework.
• 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.

In conclusion, we will re-structure the methodology section to better highlight these contributions and make clear which points are inspired by the pioneering works. 5. Responses to the specific questions:

[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.

[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.

[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.

We 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.

We sincerely appreciate the detailed and constructive feedback from the reviewer BE4T! We will address each question/concern as follows:

1. 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.
   Furthermore, 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.

2. 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.

3. 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: |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 | |---|---|---|---|---|---|---|---|---|---|---|---|---|---| |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|

4. 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.

5. Treatments of potential energy E(x). We appreciate the insightful questions which we are happy to discuss as follows:

(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.

(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.

[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.

(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).

6. 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.

7. 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.

8. 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.

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.

1. 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.
2. 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.
3. 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: Change 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. Clearly define $E_θ = -\alpha \log(p_θ/p_{ref})$ section 3.3 to 1. Simply the notation and 2. reinforce the connection to DPO modeling.
4. 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.
5. Presentation and minor corrections (comments). We also appreciate the reviewer’s attention to detail and address the specific presentation issues. We have revised accordingly:

• Correcting the equation in Line 090 to show the exact transition notation.
• Clarifying the connection between $x_θ$ and $ϵ_θ$ in the background section.
• Fixing the notation below equation (14) to ensure $x̂_0$ is properly represented.
• Adding axis labels and annotations for Figures 3 and 4.

6. 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.

Finally, 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.