Solving Inverse Problems in Protein Space Using Diffusion-Based Priors

We would like to thank our reviewer for emphasizing the quality of our results, as well as the relevance of our experiments, especially our study on the influence of the completeness level in the model refinement part. We provide an updated version of the manuscript where modifications are shown in red.

Introduction

We thank our reviewer for thoroughly reading our introduction and suggesting ways to clarify and improve it. We agree that Paragraph 4 (development of diffusion models in the space of protein atomic models) breaks the flow: it has been shortened and merged to the end of the previous paragraph. Instead of swapping Paragraphs 2 and 3, we propose to split the third paragraph in two. The first part focuses on drawing the parallel between the paths taken by protein structure determination methods and by computational imaging methods. The second part (now Paragraph 4) focuses on the possibility of using generative models as priors in inverse problems, an idea mostly introduced and developed in the field of imaging. Our introduction can now be summarized as follows:

• Paragraph 1: Protein structure determination was first formulated as a MAP problem. However, this approach requires hand-crafted priors.
• Paragraph 2: Fully supervised approaches have been developed and some of them redefined the state of the art. However, they need paired data and do not handle distribution shifts.
• Paragraph 3: The field of imaging has followed the same path: MAP with heuristic priors followed by fully-supervised approaches.
• Paragraph 4: Generative models can serve as implicit priors, which has recently been used extensively in the field of imaging. Diffusion models were developed in protein space, but have not yet been leveraged as priors for structure determination
• Paragraph 5: ADP-3D uses a MAP framework and a prior implicitly defined by a diffusion model. It circumvents the need for heuristic priors, as well as paired data.

We kindly refer to the updated version of the manuscript, where our introduction has been updated.

Related Work

The “distances-to-structure” task is a strong simplification of the reconstruction problem from NMR restraints while the cryo-EM task is more realistic. For this reason, we decided to focus on cryo-EM methods more than on NMR methods in the Related Work section. We currently provide a reference to RosettaNMR, one of the existing methods for structure determination from NMR data, in Section 5. We added references to Xplor-NIH, CS23D and CS-Rosetta.

We made an effort to shorten the first paragraph of the Related Work section. Thank you for the suggestion!

Questions

• The idea of using the denoiser of a pretrained diffusion model in a plug-n-play framework was introduced in DiffPIR (Zhu et al., 2023), a method for solving inverse problems in image space. We mention this at the end of Section 3. To clarify this point, we added a reference to DiffPIR in the last paragraph of the Introduction.
• ADP-3D is the first diffusion-based method applied to cryo-EM model refinement, and the first method applied to the “distances-to-structure” task. For the structure completion task, the conditioning framework introduced in the Chroma publication can be considered as an existing baseline, which we compare our results to in Figure 2. We modified the abstract to clarify this point.

Additional feedback

• We thank our reviewer for their suggestion regarding the first sentence of the manuscript. We rephrased it and swapped it with the next sentence.
• Thank you for pointing out the lack of theory on convergence of existing MLE and MAP-based reconstruction methods. We rephrased the last sentence of the first paragraph.
• We de-emphasized the focus on ModelAngelo in the second paragraph of the introduction.
• We corrected the second contribution: “We outperform an existing posterior sampling method at reconstructing full protein structures from partial structures”.
• Thank you for your suggestion regarding Section 4.3. Although it would be possible to have an additional section to describe the three forward models, the first and third tasks can be compactly described with one equation. In the interest of conciseness, we chose to briefly describe the forward model of each task at the beginning of each paragraph of Section 5. For the model refinement task, the forward model is more complex and we describe the technical contributions that are specific to this task in Section 4.3 (e.g., MC sampling of the side-chain angles).
• Thank you for the suggestion, we added an intermediate step in Equation 13.
• For the model refinement task, we moved our results to the supplements and replaced them with results obtained on real data. We removed the unclear remark highlighted by our reviewer.

We thank our reviewer for their detailed comments and feedback. We provide an updated version of the manuscript where modifications are shown in red. In this revised manuscript, we have now added results on experimental cryo-EM density maps for the model refinement task. We respectfully disagree with the characterization of ADP-3D as a straightforward adaptation of prior diffusion modeling approaches. While our work builds on the general plug-and-play framework, we believe the application to protein structure modeling is a substantial contribution in itself. Furthermore, ADP-3D introduces several domain-specific technical innovations that are critical for its success and go beyond a direct translation of existing methods.

Main contribution

As highlighted by our reviewer, the main contribution of this paper is the introduction of a method leveraging any pretrained diffusion model for solving a large spectrum of inverse problems in protein space. The method is inspired by DiffPIR (Zhu et al., 2023), a plug-n-play method designed for image inpainting, deblurring and super-resolution. ADP-3D operates in protein space and is evaluated on three different tasks using two different diffusion models of protein atomic models. We argue that ADP-3D is not a straightforward translation of previous work. Several technical contributions are specific to the fact that we are operating in the space of atomic models. For example, the preconditioning strategy (Section 4.2 and Appendix A) used in the structure completion and model refinement tasks is key for convergence and represents one of the methodological contributions of this paper (see, for example, our ablation study in Figure S4). For the model refinement task, our coarse-to-fine (i.e., frequency marching) strategy is also necessary for convergence and only applies to this task.

We would like to emphasize the fact that the paper demonstrates the versatility of the method and the possibility of using it as a true “plug-n-play” method. ADP-3D can be used to leverage any pretrained diffusion model as a prior, which we demonstrate by using both Chroma and RFdiffusion. ADP-3D is compatible with any inverse problem for which the likelihood of the measurements is differentiable with respect to the atomic model, making it a very general tool, which we demonstrate by tackling three different problems: structure completion, model refinement and distances-to-structure.

Results with experimental data

In order to demonstrate the applicability of ADP-3D to real scenarios, we provide additional results using experimental density maps. Our results are described in Section 5 (Figure 3) and replace our previous results on synthetic data (moved to the Supplementary Materials, G.6). We use three publicly available density maps of single-chain structures: EMD-33209 (map for PDB:7xkw), EMD-32760 (map for PDB:7wsm) and EMD-26482 (map for PDB:7ug0). We show that, by leveraging the pretrained model Chroma as a prior, ADP-3D can refine incomplete models provided by ModelAngelo. These results validate the applicability of the method on experimental data, thereby proving its relevance in real scenarios. We kindly refer to the updated Supplementary Materials for further details.

Influence of $\rho$

We realize that our notation is overloaded and apologize for the confusion. $\rho$ represents both the penalty parameter of a proximal operator and the “momentum” of a gradient step. We assume that our reviewer’s question concerns the former.

In ADP-3D, we do not perform a full maximization on the proximal operator. Instead, we perform a single gradient step (with momentum) from the iterate $\tilde{z}_0$ (L259, page 5). In $\tilde{z}_0$, the gradient of the penalty term $(\rho/2)\Vert z-\tilde{z}_0\Vert_2^2$ with respect to $z$ is $0$ in $\tilde{z}_0$. Therefore, $\rho$ has no influence on the results. We now explicitly mention this at the end of Section 4.1 and use a different notation ($\gamma$) for the penalty parameter to avoid the potential confusion with the momentum parameter. We thank our reviewer for raising this point.

Clarification on “completeness” vs. RMSD

We define the “completeness” of an incomplete model as the ratio between the number of modeled residues over the total number of residues. This term is sometimes replaced with “completion”, like in the ModelAngelo publication. When plotting completeness vs. RMSD, we first align the fully reconstructed atomic model with the deposited model. We then compute the Euclidean distances between the two models, for each alpha-carbon. The alpha-carbons are ranked by increasing deviation to the “ground truth” location. We then compute the RMSD between the two models, for the $N$ first alpha-carbons only, where $N$ ranges from 1 to the total number of residues. This gives $N$ points in the “completeness vs. RMSD” plot.

We thank our reviewer for highlighting the novelty of our method leveraging a pretrained diffusion model of protein atomic structures as a Bayesian prior, and its role in establishing new connections between generative modeling and scientific disciplines. We also thank the reviewer for their feedback and suggestions. We provide an updated version of the manuscript where modifications are shown in red. Our main additions are:

• Additional experiments with real data for the model refinement task.
• A comparison to the DPS method for the structure completion task.

Additional results with experimental data

In order to demonstrate the applicability of ADP-3D to real scenarios, we provide additional results using experimental density maps. Our results are described in Section 5 (Figure 3) and replace our previous results on synthetic data (moved to the Supplementary Materials, G.6). We use three publicly available density maps of single-chain structures: EMD-33209 (map for PDB:7xkw), EMD-32760 (map for PDB:7wsm) and EMD-26482 (map for PDB:7ug0). We show that, by leveraging the pretrained model Chroma as a prior, ADP-3D can refine incomplete models provided by ModelAngelo. These results validate the applicability of the method to experimental data, thereby proving its relevance in real scenarios. We kindly refer to the updated manuscript and Supplementary Materials for further details.

Comparison to other diffusion-based posterior sampling methods

Thank you for raising this point. As mentioned by our reviewer, there exist many different methods to perform posterior sampling using diffusion models in image space (DPS [1], PiGDM [2], etc…). We also thank our reviewer for pointing us to the survey [3], which we added to the second paragraph of the Related Work section in addition to other related works.

While designing the method, we experimented with a few of these techniques, in particular with DPS, ScoreALD [4] and DiffPIR [5]. DiffPIR is generally more accurate than scoreALD because gradients of the log-likelihood are computed on the denoised, protein-like structure. DPS uses the same idea, but gradients are computed through the denoiser, making the whole computation significantly slower.

Following our reviewer’s suggestion, we now include a comparison to DPS in the Supplementary Materials (Appendix G.4 and Figure S3). We use the structure completion task on PDB:8ok3 with a subsampling factor of 4. For DPS, we run a sweep over the parameter $\zeta^\prime$. Both methods run over 1,000 steps on 64 replicas. We find that ADP-3D leads to more accurate results and is, on average, about six times faster than DPS because it does not need to compute gradients through the denoiser. We kindly refer to the updated Supplementary Materials for further details.

[1] Chung, Hyungjin, et al. "Diffusion posterior sampling for general noisy inverse problems." arXiv preprint arXiv:2209.14687 (2022).

[2] Song, Jiaming, et al. "Pseudoinverse-guided diffusion models for inverse problems." International Conference on Learning Representations. 2023.

[3] Daras, Giannis, et al. "A survey on diffusion models for inverse problems." arXiv preprint arXiv:2410.00083 (2024)

[4] Kawar, Bahjat, et al. "Denoising diffusion restoration models." Advances in Neural Information Processing Systems 35 (2022): 23593-23606.

[5] Zhu, Yuanzhi, et al. "Denoising diffusion models for plug-and-play image restoration." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

ModelAngelo on synthetic data

We agree with our reviewer that, since ModelAngelo was trained on experimental maps but is here used on simulated maps, the input may be out-of-distribution and ModelAngelo’s output more incomplete than in a real setup. The completeness we obtained with ModelAngelo on PDB:7pzt ranges between 50% and 70% (Figure S5b). In comparison, the completeness obtained with ModelAngelo on experimental maps ranges between 16.3% and 100% (mean $\pm$ std = 80 $\pm$ 24%, see Extended Data Table 1 in [6]), These numbers show that ModelAngelo provides (sometimes significantly) incomplete models and motivate the need for a refinement method.

[6] Jamali, Kiarash, et al. "Automated model building and protein identification in cryo-EM maps." Nature 628.8007 (2024): 450-457.