3D Structure Prediction of Atomic Systems with Flow-based Direct Preference Optimization

We thank the reviewer for the valuable comments, and answer the questions as follows.

W1: The significance of the work.

Thank you for your kindly comment. Our paper, while applying DPO to atomic systems, introduces significant innovations tailored for this domain:

• Extension of Gaussian Paths: We expand the typical diffusion paths from previous works[14,20] to a broader set of Gaussian paths, and ensure these paths still meet the symmetry requirements essential for accurate structural predictions, enhancing exploration of conformational space.

• Unified DPO Objective: Our theoretical analysis ensures the universality of the DPO objective that is applicable across all Gaussian paths, not limited to the traditional diffusion paths [30].

• Automated Preference Dataset Generation: We introduced a method to automatically generate preference datasets from the training data, reducing the need for manual annotations and lowering resource demands.

• Empirical Validation: Our method has been validated on challenging antibody and crystal datasets, demonstrating its effectiveness and practical utility in real-world 3D structure prediction.

These contributions address specific challenges in 3D structure prediction for atomic systems, offering substantial advancements over existing methods.

W2: The general presentation could be improved.

Thank you for your constructive feedback.

We acknowledge your concerns regarding the general presentation and the intergration of Section 3.3. The primary intention of our paper is to demonstrate the effectiveness of DPO across various flow models and Gaussian paths for atomic structure prediction tasks. In Section 3.3, we aimed to establish the generality of DPO by ensuring the training objective is universal for arbitrary Gaussian paths, as introduced in Section 3.1. This theoretical foundation is crucial for supporting the feasibility of the experiments detailed in Section 4, bridging the gap between theoretical derivations and practical applications.

Moreover, the impact of DPO has been visualized in Figure 4 in our paper. Specifically, Figure 4 illustrates the distribution of a set of samples generated by our model, where the x-axis represents the RMSD to the ground truth and the y-axis represents the probability density. The blue curve represents the model's performance before DPO optimization, and the red curve represents performance after DPO optimization. Notably, the blue curve often exhibits a bimodal distribution, with the first peak at a lower RMSD indicating higher quality conformations, and the second peak at a higher RMSD representing conformations that deviate more from the ground truth. These higher RMSD samples often include physically implausible conformations. After applying DPO, we observe a suppression of the erroneous peak, indicating an overall improvement in the quality of the generated structures. This visualization clearly demonstrates DPO's role in enhancing model performance by effectively reducing the generation of less accurate structures. We provide additional visualizations in Figure S1 of general response, which also support the above observations.

For the overfitting problem and the bad cases, one worth-noticed phenomenan is that the learning rate for DPO training affects the model performance. For instance, a high learning rate can cause the model to deviate from the initial state, and the original distribution is forgetten. We provide an example in Table S2 of the general response for the validation performance on MPTS-52 crystal structure prediction with OT-OT path as follows, where high learning rate hinders the model performance.

Q1: How efficient is the proposed approach w.r.t. the competitors?

In assessing the efficiency of our proposed approach relative to competitors, we provide the inference time (minutes) for each method when applied to MP-20 in Table S3 of the general response. Note that DPO does not change the inference process of the flow models. Consequently, models optimized with DPO maintain comparable generation efficiency to their original versions.

Q2: The role of symmetries.

Nice Question! Symmetries indeed play a crucial role in the field of atomic systems. In our study, we have carefully addressed this issue by ensuring that all proposed flow paths rigorously maintain the symmetries specified in Eq. (6) and Eq. (12). As a result, our model inherently respects these symmetries. Specifically, the predicted antibody CDRs exhibit equivariance relative to the provided context, and the modeled crystal distributions are invariant under E(3) transformations. Given this inherent compliance with symmetry requirements, we do not need to employ additional data augmentation strategies to artificially introduce symmetries.

Dear Reviewer jaPH,

Thanks for your valuable questions regarding the symmetries and evaluation methods. We hope all your concerns have been addressed, but please let us know if there is anything further you would like to discuss.

Best,

Authors

We thank the reviewer for the valuable comments, and answer the reviewer’s questions as follows.

W1: Missing motivation for a generative approach to 3D structure prediction of atomic systems. Q1.1: why would you use a generative approach?

Thank you for your insightful question. The motivation for using generative model first arises from the scarcity of experimental data, which typically provides only a single near-stable structure. Generative models can learn and represent a distribution of potential stable structures, by maximizing the likelihood of the observed samples within the modeled distribution, filling gaps in the data. Their success in works like DiffAb, DiffCSP, DiffDock, and AlphaFold3 further supports this direction.

W2: Unclear/imprecise use of the term "hallucination". Q2: What do you understand as hallucinated samples?

Thank you for your feedback. We recognize that the term "hallucination" was used without a clear definition, leading to potential confusion. In the context of our work, "hallucination" refers to instances where the generative model produces conformations that are not only incorrect but also highly improbable from a physical view.

As depicted in Figure 4, the generative model's output distributions are represented by two curves on where the x-axis measures the RMSD from the ground truth and the y-axis indicates the probability. The blue curve, which illustrates the model's performance prior to the application of DPO, often displays a bimodal distribution. The first peak, with lower RMSD, indicates higher quality predictions, while the second peak, with higher RMSD, corresponds to conformations that deviate significantly from the ground truth.

Upon closer examination, we found that this second peak often includes physically implausible conformations. In scenarios without DPO, the probability associated with this erroneous peak can surpass that of the more accurate peak. We have termed this phenomenon as "hallucination," indicating the unexpected high probability of low-quality, implausible generations, which DPO effectively helps to suppress.

If the term "hallucination" leads to ambiguity, we are open to describing it more directly as incorrect or low-quality generations.

W3: Self-contradictory method. Q1.2: Why do you construct preferences based on distances to a single ground truth without accounting for the possibility of multiple solutions given the conditioning?

Thank you for your insightful comments. As mentioned above, high RMSD values often indicate low-quality conformations. Our DPO approach treats low RMSD structures as preferred, guiding the model to reduce low-quality generations. This adapts DPO to a quantitative metric relevant to 3D atomic structures, ensuring objectivity and scalability without relying on costly expert evaluations.

W4: Contradictory method and evaluation. Q1.3: Why does the evaluation consider distance to a single ground truth?

Thank you for your observations. We employ RMSD as an evaluation metric, recognizing its widespread acceptance and proven utility in our domain. Experimentally measured structures, which serve as benchmarks in our evaluations, are generally near a stable state. Consequently, structures generated close to these experimental conformations are typically more reliable and indicative of stability. Moreover, this metric has been widely adopted in the field, including in seminal works such as AlphaFold3, and is well-recognized within our community.

W5: The comparison with baselines is limited to generative methods. Q3: Are there non-generative (regression) baselines and if so, how is the performance compared to them?

Thanks for your advice. We utilize the same backbone model (MEAN) directly for a regression task as an additional baseline. The results are shown in Table S1 in the general response. We observe that the generative models surpass the regressive model on 4 of the 6 CDRs, notably in the most variable and critical regions, CDR-H3 and CDR-L3. Additionally, we report not only the mean RMSD across 20 generations for each generative model but also the minimum RMSD. It can be seen that the minimum RMSDs are significantly lower, showcasing that the generative models not only provide predictions that are closer to the observed reference structure but also demonstrate the ability to generate multiple viable structures around the stable state.

Q1: Is there uncertainty in two given tasks, i.e., can there be multiple different valid solutions to predicting the 3D structure given the respective prior knowledge as conditioning?

Certainly, there is inherent uncertainty in predicting the 3D structures of molecules due to the dynamic nature of molecular conformations. These conformations are not static or fixed; instead, there are many possible structures clustering around the energy minima[a]. Therefore, a generative modeling approach is employed to capture the probability distribution of these stable structures, aiming to maximize the likelihood of observing conformations close to these energy minima. Regarding the use of RMSD as a metric, it serves to assess how closely the generated structures align with a known stable conformation. Since the plausible structures only deviate slightly from each other, structures distinct from the observed stable conformations are still less reliable. Therefore, although this approach does not account for the multiplicity of valid solutions, it provides a practical measure of deviation from a recognized ground truth, facilitating the evaluation of model performance in generating physically plausible structures.

[a] Fernández-Quintero, Monica L., et al. "CDR-H3 loop ensemble in solution–conformational selection upon antibody binding." MAbs. Vol. 11. No. 6. Taylor & Francis, 2019.

We thank the reviewer for the constructive comments and address the follow-up questions as follows:

Motivation:

Yes, your understanding is correct. We utilize flow matching models to generate candidate structures, and those with high RMSDs are considered low-quality. These low-quality predictions, referred to as "hallucinations" in the paper, are suppressed using the DPO proposed in Section 3.3.

Regarding W3:

There is indeed a similarity between the flow matching training objective and the RMSD metric. In fact, the rectification strategy suggested by the reviewer—to align the model-predicted $x_0$ at each step with the ground truth—is essentially equivalent to the original flow matching objective. Taking the OT path as an example, the flow matching objective on the vector field is to minimize $\\\|v_t-\frac{x_t-x_0}{t}\\\|_2^2$, which can be reparameterized as $\\\|x_0-(x_t-v_t\cdot t)\\\|_2^2$.

Then why does DPO outperform the original model? The key difference lies in the training objective. Unlike the original loss, which solely fits predictions to the ground truth, DPO not only guides the model towards the "win" cases but also corrects inaccuracies in the "lose" cases, steering the model away from low-quality predictions.

Moreover, regarding the potential causes of hallucinations, one possible factor is the imperfect modeling of intermediate states. In a multi-step generation process, small errors can accumulate, leading to inaccuracies in the final structure. The DPO objective mitigates this by correcting the outputs from the inaccurate ("lose") cases at each timestep, thereby improving overall performance.

Regarding W5:

The performance of the regression baseline on CDR-L1 can be attributed to the relatively low diversity of this particular CDR, where the structure patterns are similar. Hence the regression task might be better in this specific case. However, research has primarily focused on more flexible regions like CDR-H3 and CDR-L3, where structure complexity and context-dependency make prediction tasks more challenging. In these cases, generative models exhibit significantly better performance.

We appreciate the kindly discussions with the reviewer, and look forward to your further feedback.

The rebuttal addresses most of the mentioned weaknesses and questions. After carefully considering all reviews and answers by the authors, I am still advocating for accepting the paper by increasing my rating to: 6 Weak Accept (cf. edited rating)

Thanks for your constructive comments! We provide more explanations to address your concerns as follows.

W1: The motivation for symmetry preservation and transform invariance of the flow trajectory is unclear and seems not quite relevant to the main idea of using DPO to improve sample accuracy. The theoretical justification for the trajectory needing to be transform invariant is not well explained, apart from the empirical results.

Thanks for your feedback. In many atomic systems, symmetry indeed plays a fundamental role, as the physical laws governing these systems are invariant under certain transformations. More specifically, for antibody structure prediction, the E(3) equivariance of the trajectory in Eq. (7) ensures that the predicted CDRs are equivariant to the given context [20]. Similarly for crystal structure prediction, the periodic E(3) equivariance described in Eq. (13-14) ensures that the predictions retain the inherent symmetries in crystal structures [14]. As we also explore various new flow paths beyond those in previous studies [14,20], it is imperative to ensure that these newly proposed paths still adhere to the necessary symmetries, thereby validating their correctness. We will expand on these points in the revised paper to more clearly articulate the role and necessity of maintaining these symmetries.

W2: In L57, the authors claim that they are "the first to theoretically prove the compatibility of DPO with arbitrary Gaussian paths by deriving a unified objective." However, I can hardly find anything relevant to this claim, and the paper seems more focused on application rather than theory.

Thank you for your observation. We apologize for any ambiguity in our presentation. In Section 3.3, we generalize the feasibility of DPO from diffusion-based VP paths to arbitrary Gaussian paths. Specifically, as detailed in Eq. (27), we demonstrate that Gaussian paths can be universally discretized using conditional mean values parameterized by $x_i$ and $x_0$. As $x_i$ is given, the KL divergence term in Eq. (26) can be reparameterized by the MSE of the ground truth $x_0$ and $x_{0,\theta}$ predicted by the flow model. This formulation allows us to derive a unified DPO training objective applicable to all Gaussian paths, as outlined in Eq. (28). We acknowledge the need for clearer exposition on this theoretical contribution and will enhance the relevant sections in our paper to better articulate these derivations and their implications.

W3: In general, the contribution is limited, considering that diffusion DPO is a known method in [30], which is theoretically universally applicable to all diffusion flow matching models. While this work touches on the new field of atomic system prediction, using a universally applicable method in this field does not seem very novel to me.

Thanks for pointing this out. As mentioned above, while it is true that Diffusion-DPO [30] establishes the framework of DPO within the context of diffusion models like DDPM, which predominantly utilize the Variance Preserving (VP) path, our contribution extends this framework to arbitrary Gaussian paths and adapts to a broader class of flow models, and demonstrates the practical applicability on 3D structure prediction tasks.

W4: The overall structure is not focused. While the main idea of this work is the application of diffusion DPO rather than proposing a new theory, the paper spends a lot of space on preliminaries with very specific examples. The main body, especially the experiments section, is short and lacking in detail.

Thanks for your feedback. We acknowledge that the current structure of the paper may not optimally highlight the core contributions of our work. Our primary aim is to demonstrate the efficacy of DPO across a range of flow models characterized by arbitrary Gaussian paths, specifically within the context of 3D structure prediction tasks.

To establish this, we introduced multiple flow paths that extend beyond the typical diffusion-based paths previously explored in literature [14,20]. These paths are crucial for demonstrating the broad applicability of DPO in handling diverse flow models. Additionally, maintaining the necessary symmetries for each specific prediction task is essential for the design of each proposed path, which we underscore in Section 3.1. We derives the generality of DPO in Section 3.3, and demonstrate its effectiveness for multiple tasks and multiple flow paths in the experiments. We will carefully reorganize the content to ensure a more logical flow that aligns with the main contributions of our work.

W5: The subscript 'i' is a bit confusing as it sometimes denotes the index of a data sample and sometimes the index of a timestep.

Thanks for your kind advice. We would like to change the notation of timesteps into 's' to avoid confusion.

Thank you again for your valuable suggestions. We believe we have addressed your concerns and kindly hope for your feedback and reconsideration of the scores.

Dear Reviewer SYsR,

Thanks again for your insightful comments. We have addressed your concerns by clarifying the motivations and overall structure of this work in our rebuttal. Please let us know if you have any further questions.

Best regards,

Authors