Solving Inverse Problems in Protein Space Using Diffusion-Based Priors

Authors: Axel Levy, Eric Ryan Chan, Sara Fridovich-Keil, FREDERIC P POITEVIN, Ellen D Zhong, Gordon Wetzstein
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: machine learning for structural biology, protein structure determination, diffusion models, inverse problems

Abstract:
The interaction of a protein with its environment can be understood and controlled via its 3D structure. Experimental methods for protein structure determination, such as X-ray crystallography or cryogenic electron microscopy, shed light on biological processes but introduce challenging inverse problems. Learning-based approaches have emerged as accurate and efficient methods to solve these inverse problems for 3D structure determination, but are specialized for a predefined type of measurement. Here, we introduce a versatile framework to turn biophysical measurements, such as cryo-EM density maps, into 3D atomic models. Our method combines a physics-based forward model of the measurement process with a pretrained generative model providing a task-agnostic, data-driven prior. Our method outperforms posterior sampling baselines on linear and non-linear inverse problems. In particular, it is the first diffusion-based method for refining atomic models from cryo-EM maps and building atomic models from sparse distance matrices.

Decision: reject
Reviewer scores: 5, 3, 8

Reviews and discussion:
<thread>
creation date:
1730407601270

thread_type:
official_review

<thread_object>
creator: reviewer_favv

<content>
summary:
- The authors propose to use a plug-and-play framework for diffusion models to inverse problems for determining protein structure
- This allows the authors to use off-the-shelf diffusion models such as Chroma to regularise these inverse problems
- The upshot is that this line of work is task-agnostic
- The authors showcase how their method yields promising results for three inverse problems: structure estimation from an incomplete structure, structure estimation from Cryo-EM 3D maps, and structure estimation from incomplete sparse pairwise distance data

strengths:
- It is very impressive to see the results for all three inverse problems. Some aspects that are highlights to me:
    - In the structure completion it is great that the authors also test this for different levels of incompleteness. For these methods, it is good to showcase limitations to the reader
    - I am very impressed by outperforming ModelAngelo, which is as the authors mention (and to the best of my knowledge) the state-of-the-art in model building. Being able to refine the predictions to the extent in the paper is very impressive

weaknesses:
- The writing can be improved upon. In particular, the introduction and the related work could really benefit from restructuring and adding additional material. This would benefit the overall flow of the paper and provide the reader with a clearer overview of the problems being tackled in the paper. Detailed suggestions below:
    - Regarding the introduction:
        - Please allow me to paraphrase:
            - Paragraph 1: Proteins can be inferred in different ways through solving an inverse problem
            - Paragraph 2: Deep learning has been key take the next step, but are limited by their way of training (new problem setting, train again to the new setting)
            - Paragraph 3: Inverse problems are dominated by learned priors
            - Paragraph 4: Diffusion models are used in the field of biology
            - Paragraph 5: Our contribution to overcome the hurdle from the second paragrap is to use diffusion models
        - I would suggest a rewrite that would look as follows (high-level)
            - The order Paragraph 1, Paragraph 3, Paragraph 2, Paragraph 5 would be more natural to me. 
            - Paragraph 4 does not really serve any purpose and for me disrupts the flow of the introduction. So I would recommend removing it from the introduction and maybe move it to the background 
    - Regarding related work: 
        - I would expect that there would be something on finding structures from partial measurements and sparse distance matrices as well. Now there is just a lot on model building methods from Cryo-EM density maps
        - I would suggest to talk less about protein diffusion models, which seems more like something for the the background (I feel that there is not really a reason of spending a lot of time on it here)

questions:
Questions for clarification 
- Is the idea of using the score as a denoiser yours or is this something that is commonly used in the inverse problems field? You don’t seem to cite the works referred to in paragraph 3 of the introduction.
- On that note, what confuses me is the sentence “In particular, it is the first diffusion-based method for refining atomic models from simulated cryo-EM maps.” at the end of the abstract. Has it been used for the other two inverse problems? It would be great if the authors can be more precise here. (I assume that Cryo-EM is mentioned as it is potentially the most high-impact application, but it would be good to be clear about the other problems as well)

Additional feedback
- p.1, l.28: I would not start this paragraph with this sentence. The second sentence starting with “Experimental methods in structural biology…” would be a better choice to give the reader a better idea of what you are actually addressing in this paragraph. The first sentence can still be used later on in the paragraph to motivate the experiments.
- p.1. l.37: I would be reluctant to use the sentence “These approaches can estimate the uncertainty of their answers and provide theoretical guarantees of correctness, …” as there is little underlying theory that even guarantees that the iterates converge (the loss might converge at best, for which there is some theory in the case of RELION)
- In section 2, the authors mainly focus on ModelAngelo. I would recommend to discuss the broader problem so that the reader gets a better feeling for the broader picture
- p.2 l.90: In the second bullet of the contributions the authors claim to compare to sampling methods. In the experiment only one method is discussed. Please change this to singular (method)
- Section 4.3: It feels a bit odd that there is a whole section for the model building part from Cryo-EM maps, but that the other inverse problems are defined in the experiments. I would suggest to create a new section in between what is now section 4 and 5 and discuss all three inverse problems there. This will be useful and more straightforward for the reader who is interested in just one of the applications.
- p.7 l.360: I understand that in equation 13 several terms that are being conditioned on can be removed (\bar{x} in the last term), but for the reader this might be a big step. I would suggest to write out the steps explicitly to get from the left hand side to the right hand side
- p.8 l.409. typo: evalute -> evaluate
- p.9 l.458. What do you mean with the sentence “Note that we do not chose a structure that was resolved with cryo-EM because most cryo-EM-resolved models are incomplete.”
    - Also typo: chose -> choose
</content>
</thread_object>

</thread>
<thread>
creation date:
1730653500869

thread_type:
official_review

<thread_object>
creator: reviewer_1gsm

<content>
summary:
The paper introduces ADP-3D (Atomic Denoising Prior for 3D reconstruction), an algorithm to combine a pretrained protein diffusion model with additional data. The method uses a plug-n-play approach that splits the log posterior into a prior and likelihood term, duplicates the variables and couples them via a quadratic term. This allows the authors to avoid incorporating the likelihood directly into the diffusion process. Instead only a simple quadratic term needs to be added to the likelihood, and a diffusion step is used to incorporate the prior. No retraining of the diffusion model is necessary. The authors focus on the Chroma model (Ingraham et al., 2023) for protein structures and combine the model with various likelihoods (structure completion from sub-sampled backbone positions, model refinement based on a density map, sequence information, incomplete distance matrix). Using specific examples, the authors find that the incorporation of the diffusion process improves the quality of the reconstructed structures.

strengths:
* A general approach to combining pretrained diffusion models with data based on a variable-splitting framework.
* The resulting algorithm is quite simple and allows for combining pretrained diffusion models with new data (thereby avoiding a retraining step).
* The approach is illustrated for various protein structure modeling tasks.

weaknesses:
* The general approach (algorithm 1) seems to be a minor modification of existing work.
* The method is illustrated only for simulated data.

__Recommendation__

I recommend to reject the article. ADP-3D is a straight forward application of the plug-n-play framework to protein diffusion models. The idea of using a variable splitting approach to combine diffusion models with additional data has already been proposed by Zhu et al. (2023) in the context of image restoration. So the major novelty is in the application of the method to protein structure modeling. However, the authors only use simulated data in their numerical experiments.

questions:
* The abstract states that this is "the first diffusion-based method for refining atomic models from simulated cryo-EM maps." Why did you not apply the method to real cryo-EM maps? The paper would be much stronger, if you showed applications to real data. Likewise, you could use NMR distance bounds as experimental distance information rather than simulated distances. 
* How do you choose the parameter $\rho$? Are the results sensitive to the choice of $\rho$? 
* How do you define __completeness__? You state that completeness is the "RMSD of the predicted structure for the X% most well-resolved alpha carbons (compared to the deposited structure)." How do you assess if an alpha carbon is well-resolved? Please note that the resolution of your model (in my understanding something like its precision) and its accuracy are not the same thing. 
* Figures 3, 4, 5 report the resolution of the deposited structure. What is the purpose of showing this number? The resolution of the PDB structure does not tell us anything about the difficulty of the structure calculation task.
* Figure 5, right panel: Why does the RMSD increase with increasing number of distances (#pd > 4k) when not using a diffusion model? It should go down and should also approach zero for as many distances as 8 k pairwise distances for a protein of length 127 residues.

__Additional feedback__

* The selection of distances is unrealistic in the sense that, for example, NMR experiments such as NOESY only probe short distances up to 6 \AA. Since there are many more long-range distances, your way of selecting distances randomly will tend to pick distances that are not experimentally accessible, which makes the selection process unrealistic. 
* The abstract states that "raw biophysical measurements such as cryo-EM density maps" can be used by your method. To categorize cryo-EM maps as raw data is misleading. A density map is the result of a complex sequence of processing steps starting with motion correction, particle picking, etc. The raw measurements in cryo-EM are the micrographs. 
* Figure 1 does not add much information beyond what is already said in the text. You could remove the figure and use the space to show more details about your numerical experiments such as your tests with RFdiffusion (Watson et al., 2023). 
* Likewise, figure 2 is a specific result for the correlation matrix used by Chroma. It could be moved to the supplementary information. 
* On page 8 you state that you avoid testing on structures that were part of the Chroma training set by using only structures that were published after the release of Chroma. However, this does not guarantee that the Chroma training set and your test structure do not overlap. First, many examples of structures of one and the same protein (or a protein with a highly similar sequence) can be found in the PDB. Second, protein structures are much more persistent than sequences. So even proteins with a low sequence similarity can have similar structures. 
* In the captions of figure 3 and 5, you report the lowest RMSD obtained with 8 runs. Please also report the maximum RMSD (you could also report the mean and the standard deviation) to give an indication of how much your models vary. 
* There is a slight inconsistency in your whole approach: You claim to do MAP estimation, but then you use diffusion-based sampling.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701033060

thread_type:
official_review

<thread_object>
creator: reviewer_gxxw

<content>
summary:
The paper introduced a framework, ADP-3D, which uses a diffusion model as a prior to solve the inverse problem in structural biology. The authors leverage the pretrained diffusion models in the atomic model space like Chroma and RFDiffusion as their prior. The downstream tasks include structure completion (as a toy problem), atomic model refinement (model completion given a simulated density map, an incomplete atomic model, and the sequence), and solving structure from given pairwise distances.

strengths:
As far as my knowledge goes, this work is one of the first to leverage a pretrained diffusion model of protein atomic structures as a Bayes prior, to solve the inverse problems which are very common in structural biology. The connection between the generative model and the experimental observation is very important in expanding the scope of the AI for science field.

weaknesses:
1. Since the paper uses pretrained diffusion models from Chroma and RFDiffusion, and the measurement models for the tasks are apparent, from a methodological standpoint, my understanding is that the main contribution of this paper is the MAP estimation method given a diffusion prior. As ADP-3D seems like a generic algorithm that is not heavily tailored for structural biology, there are many similar methods in the field of diffusion posterior sampling for inverse problems, e.g. DPS, $\Pi$GDM, as surveyed in [1] and also mentioned in the related work section by the authors. The authors should compare the similarities and differences between the propose ADP-3D method and other similar algorithms. The authors should also use at least one of these methods as the baseline and compare the results to show the advantage of the proposed ADP-3D method.

2. I understand that real data is not considered in the scope of this paper, but this makes the downstream tasks a bit far from the real world implementation. For example, the model refinement task uses a small, simulated density as a given condition, while ModelAngelo itself was trained on real experimental density maps. In reality, it is possible that ModelAngelo does not even output a severely incomplete model, making the whole setting of this task futile.

3. The model refinement task seems to only contain one example (7pzt).


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

questions:
1. How does ADP-3D related and differ from other diffusion posterior sampling methods?

2. In the model refinement task, what happens if using an experimental density map as the condition? The authors claimed that "most cryo-EM-resolved models are incomplete", but in my opinion the authors should at least have one more realistic case to show the setting in this task is not a castle in the air.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732760964864

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Shared Response

comment:
We sincerely thank all our reviewers for their constructive feedback. We value their appreciation of the novelty of our approach **[gXXw]**, the quality of our results **[fAVV]** and our illustration of the method on diverse protein structure determination tasks **[1GSm]**. We are glad they emphasized the significance of this work, establishing new connections between generative modeling and scientific disciplines **[gXXw]**. We appreciate their positive comments on the relevance of our experiments **[fAVV]**.

We provide an updated version of the manuscript where modifications are shown in red. Our main modifications include:
* The addition of results on experimental cryo-EM structures for the model refinement task instead of only synthetic data.
* A comparison to the Diffusion Posterior Sampling (DPS) method.
* A re-wording in the Introduction and Related Work sections.

In order to clarify our contribution, we would like to emphasize that ADP-3D demonstrates, for the first time, the possibility of using a pretrained diffusion model as a Bayesian prior for protein structure determination **[gXXw]**. While our approach builds on recent works solving inverse problems in image space (e.g., inpainting, deblurring, super-resolution), we believe the application to protein structure modeling is a substantial contribution in itself. To this end, ADP-3D introduces several domain-specific technical innovations that are critical for its success and go beyond a direct translation of existing methods. The paper demonstrates the versatility of the method by tackling three different inverse problems, and the possibility to leverage any user-chosen diffusion model as a prior.

In order to enable the scientific community to use and build on ADP-3D, we have recently made our code publicly available and fully open source on github: https://github.com/qt7391/adp-3d-anonymous
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors address inverse problems in structural biology by using pre-trained diffusion modeling. The paper received three reviews with two arguing for rejection and one for acceptance. In general, the reviewers appreciated the application of SOTA AI models to fundamental scientific problems. However, concerns were raised that it represents a straightforward modification of existing work and only considers simulated data. Two reviewers agreed that, in their opinions, the distances chosen to evaluate the method are not appropriate for the NMF experiment. Overall, the sense was that the technical AI/ML contribution of this work did not rise to what would be expected for ICLR.