Solving Inverse Problems in Protein Space Using Diffusion-Based Priors

## Abstract

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 raw biophysical measurements of varying types 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 both linear and non-linear inverse problems. In particular, it is the first diffusion-based method for refining atomic models from cryo-EM density maps.

## 1 Introduction

Our ability to understand the molecular machinery of living organisms and design therapeutic compounds depends on our capability to observe the three-dimensional structures of protein and other biomolecular complexes.
Experimental methods in structural biology such as X-ray crystallography, cryogenic electron microscopy (cryo-EM) and nuclear magnetic resonance (NMR) spectroscopy provide noisy and partial measurements from which the 3D structure of these biomolecules can be inferred. However, turning experimental observations into reliable 3D structural models is a challenging computational task. For many years, reconstruction algorithms were based on Maximum-A-Posteriori (MAP) estimation and often resorted to hand-crafted priors to compensate for the ill-posedness of the problem. State-of-the-art algorithms for cryo-EM reconstruction [, ] are instances of such “white-box” algorithms. These approaches can estimate the uncertainty of their answers and provide theoretical guarantees of correctness, but can only leverage explicitly defined regularizers and do not cope well with complex noise sources or missing data. Very recently, methods like Blush regularization in cryo-EM reconstruction [] use a data-driven prior based on the noise2noise framework [] to bypass the need for heuristic regularizers. However, the denoiser does not explicitly leverage the knowledge that proteins are made of atoms and cannot take advantage of a known amino acid sequence.

Recently, supervised-learning approaches have emerged as an alternative to the MAP framework and some of them established a new empirical state of the art for certain tasks. The algorithm ModelAngelo [], for example, can convert cryo-EM density maps into 3D atomic models with unprecedented accuracy. Typically, these supervised learning methods view the reconstruction problem as a regression task where a mapping between experimental measurements and atomic models needs to be learned. Some of these, like ModelAngelo, can even combine experimental data with sequence information by leveraging a pretrained protein language model []. However, these methods can only cope with a predefined type of input data. If additional information is available in a format that the model was not trained on (e.g., structural information about a fragment of the protein), or if the distribution of input data shifts at inference time (e.g., if the noise level changes due to modifications in the experimental protocol), a new model needs to be trained to properly cope with the new data.

In the field of imaging, scenarios where an image or a 3D model must be inferred from corrupted and partial observations are known as “inverse problems”. To overcome the ill-posedness of these problems, regularizers were heuristically defined to inject hand-crafted priors and turn Maximum Likelihood Estimation (MLE) problems into MAP problems. In a similar fashion, machine learning-based methods were recently shown to outperform hand-crafted algorithms for a wide variety of tasks: denoising [], inpainting [], super-resolution [], deblurring [], monocular depth estimation [], and camera calibration [], among others. Recently, generative models were shown to be an effective way to inject data-driven priors into MAP problems, making inverse problems well-posed while replacing heuristic priors []. Among these generative methods, diffusion models gained popularity due to their powerful capabilities in the unconditional generation of images [], videos [], and 3D assets [], and were recently leveraged to solve inverse problems in image space [, ].

The field of structural biology has also witnessed the application of diffusion models in protein structure modeling tasks [, ]. Building on recent progress in protein structure prediction [] (sequence to structure) and fixed-backbone design [] (structure to sequence), these diffusion models have opened the doors to de novo protein design [, , , ] (joint generation of structure and sequence). In particular, the recently released generative model Chroma [] stands out in part thanks to its “programmable” framework, i.e. its ability to be conditioned on external hard or soft constraints. However, despite experimentally validated state-of-the-art results at unconditional generation and its programmable framework, Chroma has not yet been applied to inverse problems like atomic model building.

Here, we introduce ADP-3D (Atomic Denoising Prior for 3D reconstruction), a framework to condition a diffusion model in protein space with any observations for which the measurement process can be physically modeled.
Instead of using unadjusted Langevin dynamics for posterior sampling, our approach performs MAP estimation and leverages the data-driven prior learned by a diffusion model using the plug-n-play framework [].
We demonstrate that our method handles a variety of external information: simulated cryo-EM density maps, amino acid sequence, partial 3D structure, and pairwise distances between amino acid residues, to refine a complete 3D atomic model of the protein. We show that our method outperforms posterior sampling baselines and, given a cryo-EM density map, can accurately refine incomplete atomic models provided by Modelangelo. In all our experiments, we use Chroma [] as our pretrained generative model and highlight the importance of the diffusion-based prior. We therefore make the following contributions:

- •
We introduce a versatile framework, inspired from plug-n-play, to solve inverse problems in protein space with a pretrained diffusion model as a learned prior;
- •
We outperform existing posterior sampling methods at reconstructing full protein structures from partial structures;
- •
We show that a protein diffusion model can be guided to perform atomic model refinement in simulated cryo-EM density maps;
- •
We show that a protein diffusion model can be conditioned on a sparse distance matrix.

## 2 Related Work

#### Protein Diffusion Models.

Considerable progress has been made in leveraging diffusion models for protein structure generation.
Lee et al. [] demonstrated de novo protein generation using diffusion models over 2D distance matrices, requiring a post-processing step to produce 3D structures.
Conditioned on secondary structure adjacency matrices, Anand and Achim [] used a 3D point cloud representation to generate new structures.
Trippe et al. [] circumvented the need for conditioning information, demonstrated unconditional generation of proteins larger than 20 amino acid residues and introduced a novel conditional sampling algorithm to generate structures respecting a target motif.
Wu et al. [] built a generative model that directly operates on backbone internal coordinates (i.e., dihedral angles), thereby ensuring $\text{SE}(3)$-equivariance.
Based on theoretical works extending diffusion models to Riemannian manifolds [, ], Yim et al. [] introduced a principled formulation of diffusion models in $\text{SE}(3)$ and represented protein backbones as elements of $\text{SE}(3)^{N}$.
In RFdiffusion, Watson et al. [] experimentally designed the generated proteins and structurally validated them with cryo-EM.
Recently, AlphaFold 3 [] showed that a diffusion model operating on raw atom coordinates could be used as a tool to improve protein structure prediction.
Diffusion models have been used to address specific tasks in protein space, like the generation of complementarity-determining region (CDR)-loops conditioned on non-CDR regions of antibody-antigen complexes [], or sequence-to-structure prediction [].

Out of the existing generative models for proteins, Chroma [] has reported state-of-the-art “designability” metrics with a model that can be conditioned to generate proteins with desired properties (e.g., substructure motifs, symmetries). Despite its modular interface in which users can define their own conditioners, Chroma has not yet been used to guide the generative model with experimental measurements and solve real-world inverse problems.
Here, we introduce a framework to efficiently condition a pretrained protein diffusion model and demonstrate the possibility of using cryo-EM maps as conditioning information.

#### Diffusion-Based Posterior Sampling in Image Space.

An inverse problem in image space can be defined by $\mathbf{y}=\Gamma(\mathbf{x})+\eta$ where $\mathbf{x}$ is an unknown image, $\mathbf{y}$ a measurement, $\Gamma$ a known operator and $\eta$ a noise vector of known distribution, potentially signal-dependent. The goal of posterior sampling is to sample $\mathbf{x}$ from the posterior $p(\mathbf{x}|\mathbf{y})$, the normalized product of the prior $p(\mathbf{x})$ and the likelihood $p(\mathbf{y}|\mathbf{x})$. Bora et al. [] showed that generative models could be leveraged to implicitly represent a data-learned prior and solve compressed sensing problems in image space. Motivated by the success of diffusion models at unconditional generation [], several works showed that score-based and denoising models could be used to solve linear inverse problems like super-resolution, deblurring, inpainting and colorization [, , , , , ], leading to results of unprecedented quality. Other methods leveraged the score learned by a diffusion model to solve inverse problems in medical imaging [, , , , ] and astronomy []. Finally, recent methods went beyond the scope of linear problems and used diffusion-based posterior sampling on nonlinear problems like JPEG restoration [], phase retrieval and non-uniform deblurring []. Taking inspiration from these methods, we propose to leverage protein diffusion models to solve nonlinear inverse problems in protein space.

#### Model Building Methods.

Cryogenic electron-microscopy (cryo-EM) provides an estimate of the 3D electron scattering potential (or density map) of a protein. The task of fitting an atomic model $\mathbf{x}$ into this 3D map $\mathbf{y}$ is called model building and can be seen as a nonlinear inverse problem in protein space (see 4.3).

Model building methods were first developed in X-Ray crystallography [] and automated methods like Rosetta de-novo [], PHENIX [, ] and MAINMAST [] were later implemented for cryo-EM data.
Although they constituted a milestone towards the automation of model building, obtained structures were often incomplete and needed refinement [].
Supervised learning techniques were built for model building, relying on CNN and U-Net-based architectures [, , ]. EMBuild [] was the first method to make use of sequence information and Modelangelo [] established a new state of the art for automated de novo model building. Trained on 3,715 experimental paired datapoints, Modelangelo uses a GNN-based architecture and processes the sequence information with a pretrained language model [].
Although fully-supervised methods outperform previous approaches, they still provide incomplete atomic models and cannot use a type of input data it was not trained with as additional information.

Here, we propose a versatile framework to solve inverse problems in protein space, including atomic model refinement. Our approach can cope with auxiliary measurements for which the measurement process is known. Our framework allows any pretrained diffusion model to be plugged-in as a prior and can therefore take advantage of future developments in generative models without any task-specific retraining step.

## 3 Background

Figure: Figure 1: Overview of ADP-3D. Our method turns partial and noisy measurements (the “conditioning information”) into a 3D structure by leveraging a pretrained diffusion model (here, Chroma []) and physics-based models of the measurement processes. Starting from a random structure $\mathbf{x}_{T}$, our method iterates between a denoising step and a data-matching step. The denoiser comes from the pretrained diffusion model. The data-matching step aims at maximizing the likelihood of the measurements.
Refer to caption: x1.png

### 3.1 Diffusion in Protein Space

We describe the atomic structure of a protein of $N$ amino acid residues by the 3D Cartesian coordinates $\mathbf{x}\in\mathbb{R}^{4N\times 3}$ of the four backbone heavy atoms (N, Cα, C, O) in each residue, the amino acid sequence $\mathbf{s}\in\{1,...,20\}^{N}$, and the side chain torsion angles for each amino acid $\mathbf{\chi}\in(-\pi,\pi]^{4N}$ (the conformation of the side chain can be factorized as up to four sequential rotations).

In Chroma [], the joint distribution over all-atom structures is factorized as

$$ $p(\mathbf{x},\mathbf{s},\mathbf{\chi})=p(\mathbf{x})p(\mathbf{s}|\mathbf{x})p( \mathbf{\chi}|\mathbf{x},\mathbf{s}).$ (1) $$

The first factor on the right hand side, $p(\mathbf{x})$, is modeled as a diffusion process operating in the space of backbone structures $\mathbf{x}$. Given a structure $\mathbf{x}$ at diffusion time $t$, Chroma models the conditional distribution of the sequence $p_{\theta}(\mathbf{s}|\mathbf{x},t)$ as a conditional random field and the conditional distribution of the side chain conformations $p_{\theta}(\mathbf{\chi}|\mathbf{x},\mathbf{s},t)$ with an autoregressive model.

Adding isotropic Gaussian noise to a backbone structure would rapidly destroy simple biophysical patterns that proteins always follow (e.g., the scaling law of the radius of gyration with the number of residues). Instead, Chroma uses a non-isotropic noising process as an inductive bias to alleviate the need for the model to learn these patterns from the data. The correlation of the noise is defined in such a way that a few structural properties are statistically preserved throughout the noising process. Specifically, the forward diffusion process is defined by the variance preserving stochastic process

$$ $\mathrm{d}\mathbf{x}=-\frac{1}{2}\mathbf{x}\beta_{t}\mathrm{d}t+\sqrt{\beta_{t }}\mathbf{R}\mathrm{d}\mathbf{w},$ (2) $$

where $\beta_{t}$ is a time-dependent noising schedule and $\mathrm{d}\mathbf{w}$ is a standard Wiener process of dimension $\mathbb{R}^{4N\times 3}$.
The matrix $\mathbf{R}\in\mathbb{R}^{4N\times 4N}$ is fixed and defined explicitly based on statistical considerations regarding the structure of proteins (see [] and supplementary).
Starting from $\mathbf{x}_{0}$ at $t=0$, a solution to this stochastic differential equation (SDE) at time $t$ is given by

$$ $\mathbf{x}_{t}\sim\mathcal{N}(\mathbf{x};\alpha_{t}\mathbf{x}_{0},\sigma_{t}^{ 2}\mathbf{R}\mathbf{R}^{T}),$ (3) $$

where $\alpha_{t}=\exp\left(-\frac{1}{2}\int_{0}^{t}\beta_{s}\mathrm{d}s\right)$ and $\sigma_{t}=\sqrt{1-\alpha_{t}^{2}}$.

New protein samples can be generated by sampling $\mathbf{x}_{T}$ from $\mathcal{N}(0,\mathbf{R}\mathbf{R}^{T})$ and integrating the following reverse-time SDE over $t\in[T,0]$ []:

$$ $\mathrm{d}\mathbf{x}=\biggl{[}-\frac{1}{2}\mathbf{x}-\mathbf{R}\mathbf{R}^{T} \nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})\biggr{]}\beta_{t}\mathrm{d}t+\sqrt{ \beta_{t}}\mathbf{R}d\bar{\mathbf{w}},$ (4) $$

where $d\bar{\mathbf{w}}$ is a reverse-time Wiener process. Following Tweedie’s formula [], the score $\nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})$ is an affine function of the time-dependent optimal denoiser, approximated by $\hat{\mathbf{x}}_{\theta}(\mathbf{x},t)$:

$$ $\nabla_{\mathbf{x}}\log p_{t}(\mathbf{x})=\frac{(\mathbf{R}\mathbf{R}^{T})^{-1 }}{1-\alpha_{t}^{2}}{(\alpha_{t}\mathbb{E}[\mathbf{x}_{0}|\mathbf{x}_{t}= \mathbf{x}]-\mathbf{x})},\quad\hat{\mathbf{x}}_{\theta}(\mathbf{x},t)\approx \mathbb{E}[\mathbf{x}_{0}|\mathbf{x}_{t}=\mathbf{x}].$ (5) $$

In practice, the SDE (4) is modified with a low-temperature sampling strategy in order to promote high-likelihoood states. The reverse-time low-temperature diffusion process is equilibrated with Langevin dynamics and solved numerically with the Euler-Maruyama method.

### 3.2 Half Quadratic Splitting and Plug-n-Play Framework

An objective function of the form $f(\mathbf{x})+g(\mathbf{x})$ can be efficiently minimized over $\mathbf{x}$ using a variable splitting algorithm like Half Quadratic Splitting (HQS) []. By introducing an auxiliary variable $\tilde{\mathbf{x}}$, the HQS method relies on iteratively solving two subproblems:

$$ $\displaystyle\tilde{\mathbf{x}}_{k}$ $\displaystyle=\text{prox}_{g,\rho}(\mathbf{x}_{k})=\operatorname*{argmin}_{ \tilde{\mathbf{x}}}~{}{g(\tilde{\mathbf{x}})+\frac{\rho}{2}\|\tilde{\mathbf{x} }-\mathbf{x}_{k}\|_{2}^{2}},$ (6) $\displaystyle\mathbf{x}_{k+1}$ $\displaystyle=\text{prox}_{f,\rho}(\tilde{\mathbf{x}}_{k})=\operatorname*{ argmin}_{\mathbf{x}}~{}{f(\mathbf{x})+\frac{\rho}{2}\|\mathbf{x}-\tilde{ \mathbf{x}}_{k}\|_{2}^{2}},$ $$

where prox are called “proximal operators” and $\rho>0$ is a user-defined proximal parameter.

If $f$ represents a negative log-likelihood over $\mathbf{x}$ and $g$ represents a negative log-prior, the above problem defines a Maximum-A-Posterior (MAP) problem. The key idea of the plug-and-play framework [] is to notice that the first minimization problem in (6) is exactly a Gaussian denoising problem at noise level $\sigma=\sqrt{1/\rho}$ with the prior $\exp(-g(\mathbf{x}))$ in $\mathbf{x}$-space. This means that any Gaussian denoiser can be used to “plug in” a prior into a MAP problem.

Once a diffusion model has been trained, it provides a deterministic Gaussian denoiser for various noise levels, as described in (5). As recently shown in Zhu et al. [], this optimal denoiser can be used in the plug-n-play framework to solve MAP problems in image space. Here, we propose to apply this idea to inverse problems in protein space, leveraging a pretrained diffusion model.

## 4 Methods

In this section, we formulate our method, ADP-3D (Atomic Denoising Prior for 3D reconstruction), as a MAP estimation method in protein space and explain how the plug-n-play framework can be used to leverage the prior learned by a pretrained diffusion model. We then introduce our preconditioning strategy in the case of linear problems. Finally, we describe and model the measurement process in cryogenic electron microscopy. ADP-3D is described with pseudo-code in Algorithm 1.

### 4.1 General Approach

Given a set of independent measurements $\mathbf{Y}=\{\mathbf{y}_{i}\}_{i=1}^{n}$ made from the same unknown protein, our goal is to find a Maximum-A-Posteriori (MAP) estimate of the backbone structure $\mathbf{x}^{*}$.
Following Bayes’ rule,

$$ $\mathbf{x}^{*}=\operatorname*{argmax}_{\mathbf{x}}~{}\bigl{\{}p_{0}(\mathbf{x} |\mathbf{Y})\bigr{\}}\\ =\operatorname*{argmin}_{\mathbf{x}}~{}\Bigl{\{}\underbrace{-\sum_{i=1}^{n} \log p_{0}(\mathbf{y}_{i}|\mathbf{x})}_{f(\mathbf{x})}\underbrace{-\log p_{0}( \mathbf{x})}_{g(\mathbf{x})}\Bigr{\}}.$ (7) $$

While most of previous works leveraging a diffusion model for inverse problems aim at sampling from the posterior distribution $p(\mathbf{x}|\mathbf{Y})$, we are interested here in scenarios where the measurements convey enough information to make the MAP estimate unique and well-defined. In the results section, we show that MAP estimation outperforms posterior sampling in such well-defined problems.

We take inspiration from the plug-and-play framework [] to efficiently solve (7). We propose to use the optimal denoiser $\hat{\mathbf{x}}_{\theta}(\mathbf{x},t)$ of a pretrained diffusion model to solve the first subproblem in (6). Framing the optimization loop in the whitened space of $\mathbf{z}=\mathbf{R}^{-1}\mathbf{x}$, which provides more stable results, our general optimization algorithm can be summarized in three steps:

$$ $\displaystyle\tilde{\mathbf{z}}_{0}$ $\displaystyle=\mathbf{R}^{-1}\hat{\mathbf{x}}_{\theta}(\mathbf{R}\mathbf{z}_{t },t)$ $\displaystyle\text{Denoise at level }t\text{,}$ $\displaystyle\hat{\mathbf{z}}_{0}$ $\displaystyle=\operatorname*{argmin}_{\mathbf{z}}\frac{\rho}{2}\|\mathbf{z}- \tilde{\mathbf{z}}_{0}\|_{2}^{2}-\sum_{i=1}^{n}\log p_{0}(\mathbf{y}_{i}| \mathbf{z})$ Maximize likelihood, $\displaystyle\mathbf{z}_{t-1}$ $\displaystyle\sim\mathcal{N}(\alpha_{t-1}\hat{\mathbf{z}}_{0},\sigma_{t-1}^{2})$ $\displaystyle\text{Add noise at level }t-1.$ $$

Here, no specific assumptions have been made on the likelihood term and this framework could hypothetically be applied on any set of measurements for which we have a physics-based model of the measurement process.
Since the second step is not tractable in most cases, we replace the explicit minimization with a gradient step with momentum from the iterate $\tilde{\mathbf{z}}_{0}$.
This step can be implemented efficiently using automatic differentiation.

Figure: Algorithm 1 ADP-3D (Atomic Denoising Prior for 3D reconstruction)

### 4.2 Preconditioning for Linear Measurements

We consider the case where the measurement process is linear:

$$ $\mathbf{y}=\mathbf{A}\mathbf{x}_{0}+\eta=\mathbf{A}\mathbf{R}\mathbf{z}_{0}+ \eta,\quad\eta\sim\mathcal{N}(0,\mathbf{\Sigma}\in\mathbb{R}^{m\times m}),$ (8) $$

with $\mathbf{y}\in\mathbb{R}^{m}$ and $\mathbf{A}\in\mathbb{R}^{m\times 4N}$ being a known measurement matrix of rank $m$. In this case, the log-likelihood term is a quadratic function:

$$ $\log p_{0}(\mathbf{y}|\mathbf{z})=-\frac{1}{2}\|\mathbf{A}\mathbf{R}\mathbf{z} -\mathbf{y}\|_{\mathbf{\Sigma^{-1}}}^{2}+C,\text{where }\|\mathbf{x}\|_{ \mathbf{\Sigma}^{-1}}^{2}=\mathbf{x}^{T}\mathbf{\Sigma}^{-1}\mathbf{x},$ (9) $$

Figure: Figure 2: Condition number of $\mathbf{R}$ vs. number of residues $N$.
Refer to caption: x2.png

and $C$ does not depend on $\mathbf{z}$. As shown in Figure 2, the condition number of $\mathbf{R}$ (i.e., the ratio between its largest and smallest singular values) grows as a power function of the number of residues. For typical proteins ($N\geq 100$), this condition number exceeds $100$, making the maximization of the above term an ill-conditioned problem. In order to make gradient-based optimization more efficient, we propose to precondition the problem by precomputing a singular value decomposition $\mathbf{A}\mathbf{R}=\mathbf{U}\mathbf{S}\mathbf{V}^{T}$ and to set $\mathbf{\Sigma}=\sigma^{2}\mathbf{U}\mathbf{S}\mathbf{S}^{T}\mathbf{U}^{T}$. Note that this is equivalent to modeling the measurement process as $\mathbf{y}=\mathbf{A}\mathbf{R}(\mathbf{z}+\tilde{\eta})$ with $\tilde{\eta}\sim\mathcal{N}(0,\sigma^{2})$. In other words, we assume that the noise $\eta$ preserves the simple patterns in proteins, which is a reasonable hypothesis if, for example, $\mathbf{y}$ is an incomplete atomic model obtained by an upstream reconstruction algorithm that leverages prior knowledge on protein structures. The log-likelihood then becomes

$$ $\log p_{0}(\mathbf{y}|\mathbf{z})=-\frac{1}{2\sigma^{2}}\bigg{\|}\begin{ pmatrix}\mathbf{I}_{m}&0\\ 0&0\end{pmatrix}\mathbf{V}^{T}\mathbf{z}-\mathbf{S}^{+}\mathbf{U}^{T}\mathbf{y }\bigg{\|}_{2}^{2}+C.$ (10) $$

The maximization of this term is a well-posed problem that gradient ascent with momentum efficiently solves (see supplementary analyses). In (10), $\mathbf{S}^{+}$ denotes the pseudo-inverse of $\mathbf{S}$.

### 4.3 Application to Atomic Model Building

Figure: Figure 3: Structure Completion. Results on the ATAD2 protein (PDB:7qum, 130 residues). (a) Qualitative results. The input structure is a subsampled version of the target structure (subsampling factor in the top row). In the “input” row, we show the target (unknown) in gray and the locations of the known alpha carbons in colors. We report the lowest RMSD over 8 runs. (b) RMSD vs. subsampling factor. Our method is compared to Chroma conditioned with the SubstructureConditioner. The importance of the diffusion-based prior is shown. We report the mean RMSD ($\pm 1$ std) over 8 runs. The experimental (deposited) resolution is indicated with a dashed line.
Refer to caption: x3.png

#### Measurement Model in Cryo-EM.

In single particle cryo-EM, a purified solution of a target protein is flash-frozen and imaged with a transmission electron microscope, providing thousands to millions of randomly oriented 2D projection images of the protein’s electron scattering potential. Reconstruction algorithms process these images and infer a 3D density map of the protein. Given a protein $(\mathbf{x},\mathbf{s},\mathbf{\chi})$, its density map is well approximated by []

$$ $\mathbf{y}=\mathcal{B}(\Gamma(\mathbf{x},\mathbf{s},\mathbf{\chi}))+\eta\in \mathbb{R}^{D\times D\times D},$ (11) $$

where $\Gamma$ is an operator that places a sum of 5 isotropic Gaussians centered on each heavy atom. The amplitudes and standard deviations of these Gaussians, known as “form factors”, are tabulated [] and depend on the chemical element they are centered on.
$\mathcal{B}$ represents the effect of “B-factors” [] and can be viewed as a spatially-dependent blurring kernel modelling molecular motions and/or signal damping by the transfer function of the electron microscope.
$\eta$ models isotropic Gaussian noise of variance $\sigma^{2}$. This measurement model leads to the following log-likelihood:

$$ $\log p_{0}(\mathbf{y}|\mathbf{x},\mathbf{s},\mathbf{\chi})=-\frac{1}{2\sigma^{ 2}}\|\mathbf{y}-\mathcal{B}(\Gamma(\mathbf{x},\mathbf{s},\mathbf{\chi}))\|_{2} ^{2}+C.$ (12) $$

#### Likelihood Terms in Model Refinement.

We consider a 3D density map $\mathbf{y}$ provided by an upstream reconstruction method and an incomplete backbone structure $\bar{\mathbf{x}}\in\mathbb{R}^{m}$ ($m\leq 4N$) provided by an upstream model building algorithm (e.g., ModelAngelo []). Sequencing a protein is now a routine process [] and we therefore consider the sequence $\mathbf{s}$ as an additional source of information. The side chain angles $\chi$ are, however, unknown.

The log-likelihood of our measurements for a given backbone structure $\mathbf{x}$ can be decomposed as

$$ $\log p_{0}(\mathbf{y},\mathbf{s},\bar{\mathbf{x}}|\mathbf{x})=\log p_{0}( \mathbf{s}|\mathbf{x})+\log p_{0}(\bar{\mathbf{x}}|\mathbf{x})+\log p_{0}( \mathbf{y}|\mathbf{x},\mathbf{s}).$ (13) $$

On the right-hand side, the first term can be approximated using the learned conditional distribution $p_{\theta}(\mathbf{s}|\mathbf{x},t=0)$. We model $\bar{\mathbf{x}}=\mathbf{M}\mathbf{x}+\eta$ so that the second term can be handled by the preconditioning procedure described in the previous section. Finally, the last term involves the marginalization of $p_{0}(\mathbf{y}|\mathbf{x},\mathbf{s},\mathbf{\chi})$ over $\chi$. This marginalization is not tractable but (12) provides a lower bound:

$$ $\log p_{0}(\mathbf{y}|\mathbf{x},\mathbf{s})\geq\mathbb{E}_{\mathbf{\chi}\sim p _{0}(\mathbf{\chi}|\mathbf{x},\mathbf{s})}\Bigl{[}\log p_{0}(\mathbf{y}| \mathbf{x},\mathbf{s},\mathbf{\chi})\Bigr{]}\approx\mathbb{E}_{\mathbf{\chi} \sim p_{\theta}(\mathbf{\chi}|\mathbf{x},\mathbf{s},t=0)}\Bigl{[}\log p_{0}( \mathbf{y}|\mathbf{x},\mathbf{s},\mathbf{\chi})\Bigr{]},$ (14) $$

using Jensen’s inequality.
The expectation is approximated by Monte Carlo sampling and gradients of $\chi$ with respect to $\mathbf{x}$ are computed by automatic differentiation through the autoregressive sampler of $\mathbf{\chi}$.

## 5 Experiments

#### Experimental Setup.

Our method uses the publicly released version of Chroma(^1^11[https://github.com/generatebio/chroma](https://github.com/generatebio/chroma)) []. We run all our experiments using structures of proteins downloaded from the Protein Data Bank (PDB) []. In order to select proteins that do not belong to the training dataset of Chroma, here we only consider structures that were released after 2022-03-20 (Chroma was trained on a filtered version of the PDB queried on that date). We only consider single-chain structures for which all the residues have been spatially resolved, so that the Root Mean Square Deviation (RMSD) of the predicted Cartesian coordinates can be computed for all the heavy atoms in the backbone. In each experiment, we run 8 replicas in parallel on a single NVIDIA A100 GPU.

#### Structure Completion.

Given an incomplete atomic model of a protein, our first task is to predict the coordinates of all heavy atoms in the backbone. The forward measurement process can be modeled as $\mathbf{y}=\mathbf{M}\mathbf{x}$ where $\mathbf{M}\in\{0,1\}^{(4N/k)\times 4N}$ is a masking matrix ($\mathbf{M}\mathbf{1}=\mathbf{1}$) and $k$ is the subsampling factor. We consider the case where, for each residue, the location of all 4 heavy atoms on the backbone (N, Cα, C, O) is either known or unknown. Residues of known locations are regularly spaced along the backbone every $k$ residues. We compare our results to the baseline Chroma conditioned with a SubstructureConditioner []. This baseline samples from the posterior probability $p(\mathbf{x}|\mathbf{y})$ using unadjusted Langevin dynamics. We use 1000 diffusion steps for our method and the baseline.

In Figure 3, we show our results on ATAD2 (PDB:7qum) [, ], a cancer-associated protein of 130 residues. The protein was resolved at a resolution of 1.5 Å using X-ray crystallography. Our method recovers the target structure without loss of information ($\text{RMSD}<1.5~{}\text{\AA}$) for subsampling factors of 2, 4 and 8. Fig. 3.b shows that our method outperforms the baseline and highlights the importance of the diffusion-based prior. When the subsampling factor is large ($\geq 32$), the reconstruction accuracy decreases but the method inpaints unknown regions with realistic secondary structures (see quantitative evaluation in the supplementary). Note that making the conditioning information sparser (increasing the subsampling factor) tends to close the gap between our method (MAP estimation) and the baseline (posterior sampling). We provide results on four other protein structures in the supplementary.

Figure: Figure 4: Atomic Model Refinement. Results on the TecA bacterial toxin (PDB:7pzt, 160 residues). (a) Qualitative results. From left to right: the input density map at 2.0 Å resolution, the incomplete model given by ModelAngelo and our refined models (1 sample and 5 samples), overlaid on the target structure in transparency. (b) RMSD of alpha carbons vs. completeness (number of predicted residues / total number of residues) with ModelAngelo (MA) and our method. We run 5 experiments and report the mean of the lowest RMSD on $\alpha$-carbons over 8 replicas ($\pm 1$ std). The experimental (deposited) resolution is indicated with a dashed line.
Refer to caption: x4.png

#### Atomic Model Refinement.

Next, we evaluate our method on the model refinement task. We use ChimeraX [, ] to simulate the 3D density maps of a single-chain protein, at different resolutions. We run the state-of-the-art model building method ModelAngelo [] on the simulated maps, using the known sequence. Then, we provide ModelAngelo’s output (an incomplete model) to our method, along with the density map and the sequence. We use ModelAngelo’s default parameters. To evaluate our method, we report the RMSD of the predicted structure for the $X$% most well-resolved alpha carbons (compared to the deposited structure), for $X\in[0,100]$ ($X$ is called the “completeness”).

In Figure 4, we show our results on a synthetic density map of the TecA bacterial toxin (PDB:7pzt).
This protein structure belongs to the CASP15 dataset, released in May 2022 to evaluate protein structure modeling methods.
Proteins belonging to this dataset were selected for not having close homologous in the PDB, which ensures that Chroma was not trained on similar structures.
The protein was resolved at 1.84 Å resolution using X-ray crystallography. Note that we do not chose a structure that was resolved with cryo-EM because most cryo-EM-resolved models are incomplete. For both input resolutions (2.0 Å and 4.0 Å), our method improves on ModelAngelo’s accuracy for a fixed completeness level (and reaches a higher completeness for the same accuracy). We explore the effect of removing some of the conditioning information in the supplementary materials.

#### Pairwise Distances to Structure

Finally, we assume we are given a set of pairwise distances between alpha carbons of a protein and we use our method to predict a full 3D structure. This task is a simplification of the reconstruction problem in paramagnetic NMR spectroscopy, where one can obtain information about the relative distances and orientations between pairs of atoms via the nuclear Overhauser effect and sparse paramagnetic restraints, and must deduce the Cartesian coordinates of every atom [, ]. Formally, our measurement model is $\mathbf{y}=\|\mathbf{D}\mathbf{x}\|_{2}\in\mathbb{R}^{m}$,
where $\mathbf{D}\in\{-1,0,1\}^{m\times 4N}$ is the distance matrix and the norm is taken row-wise (in $xyz$ space). $\mathbf{D}$ contains a single “1” and a single “-1” in each row and is not redundant (the distance between a given pair of atoms is measured at most once). $m$ corresponds to the number of measured distances.

We evaluate our method on the bromodomain-containing protein 4 (BRD4, PDB:7r5b []), a protein involved in the development of a specific type of cancer (NUT midline carinoma) [] and targeted by pharmaceutical drugs []. For a given number $m$, we randomly sample $m$ pairs of alpha carbons (without redundancy) between which we assume the distances to be known. Our results are shown in Figure 5. When 500 pairwise distances or more are known, our method recovers the structural information of the target structure ($\text{RMSD}<1.77~{}\text{\AA}$, the resolution of the deposited structure resolved with X-ray crystallography). We conduct the same experiment without the diffusion model and show a drop of accuracy, highlighting the importance of the generative prior. Note that, when the diffusion model is removed, increasing the number of measurements increases the number of local minimima in the objective function and can therefore hurt the reconstruction quality (plot in Fig. 5, orange curve, far-right part).

Figure: Figure 5: Pairwise Distances to Structures. Results on BRD4 (PDB:7r5b, 127 residues) (a) Qualitative results. The reconstructed structures are shown in colors, depending on the number of known pairwise distances. We report the lowest RMSD over 8 runs. The target structure is shown in gray along with its pairwise distance matrix. (b) RMSD vs. number of known pairwise distances. Each experiment is ran 10 times with randomly sampled distances. We report the mean of the lowest RMSD obtained over 8 replicas ($\pm 1$ std). The plot demonstrates the importance of the diffusion model. The experimental (deposited) resolution is indicated with a dashed line.
Refer to caption: x5.png

## 6 Discussion

This paper introduces ADP-3D, a method to leverage a pretrained protein diffusion model for protein structure determination.
ADP-3D is not tied to a specific diffusion model and allows for any data-driven denoisers to be plugged in as priors. Our method can therefore continually benefit from the development of more powerful or more specialized generative models.

We focused here on cases where the measurement process is simulated and perfectly known. Considering real experimental measurements (e.g., cryo-EM or NMR data) would raise complex and exciting challenges. In particular, we hope ADP-3D can serve as a precursor for reconstructing all-atom models directly from cryo-EM images, a key to turn empirical conformational distributions into thermodynamic information [].
In cases where the measurement process cannot be faithfully modeled due to complex nonidealities, or when the measurement process is not differentiable, our framework reaches its boundaries. At the detriment of versatility, exploring the possibility of finetuning a pretrained diffusion model on paired data for conditional generation constitutes another promising avenue for future work.

#### Acknowledgements.

The author(s) are pleased to acknowledge that the work reported on in this paper was substantially performed using the Princeton Research Computing resources at Princeton University which is a consortium of groups led by the Princeton Institute for Computational Science and Engineering (PICSciE) and Office of Information Technology’s Research Computing. This material is based upon work supported by the National Science Foundation under award number 2303178 to SFK. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation.

We thank Mike Dunne and Jay Shenoy for their insightful feedback on this manuscript.

## Appendix

## Appendix A Log-Likelihood Functions

We define here the log-likelihood functions, as mentioned in Algorithm 1.

#### Structure Completion.

For structure completion, we use the log-likelihood function defined in (10):

$$ $f(\mathbf{y},\mathbf{z})=-\bigg{\|}\begin{pmatrix}\mathbf{I}_{m}&0\\ 0&0\end{pmatrix}\mathbf{V}^{T}\mathbf{z}-\mathbf{S}^{+}\mathbf{U}^{T}\mathbf{y }\bigg{\|}_{2}^{2}.$ (15) $$

#### Model Refinement.

In the model refinement task, we combine three sources of information with the following three log-likelihood functions:

$$ $\displaystyle f_{m}(\bar{\mathbf{x}},\mathbf{z})$ $\displaystyle=-\bigg{\|}\begin{pmatrix}\mathbf{I}_{m}&0\\ 0&0\end{pmatrix}\mathbf{V}^{T}\mathbf{z}-\mathbf{S}^{+}\mathbf{U}^{T}\bar{ \mathbf{x}}\bigg{\|}_{2}^{2}$ (incomplete model) (16) $\displaystyle f_{s}(\mathbf{s},\mathbf{z})$ $\displaystyle=\log p_{\theta}(\mathbf{s}|\mathbf{R}\mathbf{z},t=0)$ (sequence) $\displaystyle f_{d}(\mathbf{y},\mathbf{z})$ $\displaystyle=-\|\mathbf{y}-\mathcal{B}(\Gamma(\mathbf{R}\mathbf{z},\mathbf{s} ,\mathbf{\chi}))\|_{2}^{2},\text{ where }\mathbf{\chi}\sim p_{\theta}(\mathbf{ \chi}|\mathbf{R}\mathbf{z},\mathbf{s},t=0)$ (cryo-EM density) $$

$\Gamma$ is an operator turning an all-atom model into a density map. Given $\mathbf{X}=(\mathbf{x},\mathbf{s},\mathbf{\chi})=[X_{1},\ldots,X_{A}]^{T}\in(
\mathbb{R}^{3})^{A}$, the Cartesian coordinates of $A$ heavy atoms,

$$ $\Gamma(\mathbf{X}):x\in\mathbb{R}^{3}\mapsto\sum_{i=1}^{A}\sum_{j=1}^{5}4\pi \frac{a_{i,j}}{b_{i,j}}\exp\left(-\frac{4\pi^{2}}{b_{i,j}}\|X_{i}-x\|_{2}^{2} \right)\in\mathbb{R},$ (17) $$

where $a_{i,j}$ and $b_{i,j}$ are tabulated form factors [].
The norm $\|\mathbf{y}-\Gamma(\mathbf{X})\|_{2}^{2}$ is computed directly in Fourier space using Parseval’s theorem, up to a time-dependent resolution $r_{t}$. Here, we assume the B-factors are unknown and we neglect their contribution in the log-likelihood function as we consider high-resolution density maps ($\leq 4~{}\text{\AA}$).

#### Pairwise Distances to Structure.

We use the following log-likelihood function:

$$ $f(\mathbf{y},\mathbf{z})=-\big{\|}\mathbf{y}-\|\mathbf{D}\mathbf{R}\mathbf{z} \|_{2}\big{\|}_{2}^{2}.$ (18) $$

## Appendix B Optimization parameters

#### Structure Completion.

All experiments are ran for 1,000 epochs, with a learning rate $\lambda$ of 0.3 and a momentum $\rho$ of 0.9. The time schedule is linear:

$$ $t=1-\text{epoch}/1000.$ (19) $$

#### Model Refinement.

All experiments are ran for 4,000 epochs. The learning rates are

$$ $\lambda_{m}=0.1,\quad\lambda_{s}=\begin{cases}0\text{ if epoch}<\text{3,000}\\ 1\times 10^{-5}\text{ otherwise}\end{cases},\quad\lambda_{d}=3\times 10^{-5} \cdot(r_{t}/r_{0})^{3},$ (20) $$

and all the momenta are set to 0.9. $r_{t}$ is set to 15 Å for the first 3,000 epochs and linearly decreases to 5 Å during the last 1,000 epochs. The side chain angles ($\mathbf{\chi}$) are sampled every 100 epochs. The time schedule is

$$ $t=1-\sqrt{\text{epoch}/\text{4,000}}.$ (21) $$

#### Pairwise Distances to Structure.

All experiments are ran for 1,000 epochs, with a learning rate $\lambda=200/n$, where $n$ is the number of known pairwise distances, and a momentum $\rho$ of 0.99. The time schedule is

$$ $t=1-\sqrt{\text{epoch}/\text{1,000}}.$ (22) $$

## Appendix C Correlated Diffusion

We use the “$R_{g}$-confined globular polymer” correlation matrix, as defined in []:

$$ $\mathbf{R}=a\begin{bmatrix}\nu b^{0}&&&&&\\ \nu b^{1}&b^{0}&&&&\\ \nu b^{2}&b^{1}&b^{0}&&&\\ \vdots&&\ddots&\ddots&&\\ \nu b^{N-2}&&&b^{1}&b^{0}&\\ \nu b^{N-1}&&\cdots&b^{2}&b^{1}&b^{0}\end{bmatrix},$ (23) $$

where $a$ is a free parameter, $\nu=\frac{1}{\sqrt{1-b^{2}}}$ and $b$ is chosen such that the radius of gyration $R_{g}$ scales with the number of residues as $R_{g}\sim rN^{\alpha}$ ($r\approx 2.0~{}\text{\AA}$, $\alpha\approx 0.4$ [, ]).

## Appendix D Target Structures

Information on the proteins used as target structures are given in Table 1.

**Table 1: Information on target structures.**
| PDB | Nb. of residues | Resolution | Release date | Imaging method | $\in$ CASP15 |
| --- | --- | --- | --- | --- | --- |
| 7r5b | 127 | 1.77 Å | 2023-02-08 | X-ray diffraction | No |
| 7qum | 130 | 1.50 Å | 2023-03-01 | X-ray diffraction | No |
| 1cfd | 148 | N/A | 1995-12-07 | NMR | No |
| 7pzt | 160 | 1.84 Å | 2022-11-02 | X-ray diffraction | Yes |
| 1a2f | 291 | 2.10 Å | 1998-11-25 | X-ray diffraction | No |

## Appendix E Additional Results

### E.1 Structure Completion

We provide additional results in Table 2. The Evidence Lower Bound (ELBO) is computed with Chroma and is a lower bound of the learned log-prior. Note that, as the subsampling factor increases (i.e., as the problem becomes less constrained), the gap between our approach (MAP estimation) and the baseline (low-temperature posterior sampling) decreases. For highly undersampled structures, the RMSD with the target structure is high but the generated structure remains plausible under the learned prior distribution.

**Table 2: Additional structure completion results. We compare our method to Chroma [] with a SubstructureConditioner. We report the RMSD in Å (best of 8) and the associated Evidence Lower Bound (ELBO) provided by Chroma as a proxy for realism.**
| PDB (# res.) | Metric | Method | Subsampling factor |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 2 | 4 | 8 | 16 | 32 | 64 | 128 |  |  |  |
| 7r5b (127) | RMSD ($\downarrow$) | Chroma* | 0.41 | 1.34 | 3.73 | 9.34 | 15.95 | 16.27 | – |
| ADP-3D | 0.10 | 0.47 | 2.21 | 4.47 | 9.45 | 11.54 | – |  |  |
| ELBO ($\uparrow$) | Chroma* | 7.79 | 6.72 | 7.27 | 7.56 | 7.40 | 9.10 | – |  |
| ADP-3D | 7.38 | 7.87 | 7.48 | 6.65 | 7.64 | 7.98 | – |  |  |
| 7qum (130) | RMSD ($\downarrow$) | Chroma* | 0.28 | 0.81 | 2.14 | 3.46 | 7.70 | 11.49 | 14.45 |
| ADP-3D | 0.21 | 0.27 | 1.13 | 2.68 | 5.60 | 7.76 | 13.01 |  |  |
| ELBO ($\uparrow$) | Chroma* | 5.87 | 6.08 | 7.16 | 7.13 | 8.43 | 8.89 | 9.27 |  |
| ADP-3D | 5.64 | 6.47 | 7.87 | 7.56 | 7.99 | 8.76 | 8.98 |  |  |
| 1cfd (148) | RMSD ($\downarrow$) | Chroma* | 0.42 | 1.17 | 3.06 | 5.84 | 9.44 | 14.28 | 15.87 |
| ADP-3D | 0.20 | 0.59 | 1.83 | 4.62 | 12.76 | 12.51 | 15.29 |  |  |
| ELBO ($\uparrow$) | Chroma* | 5.65 | 5.97 | 6.49 | 5.46 | 8.87 | 9.22 | 8.94 |  |
| ADP-3D | 5.22 | 5.34 | 7.60 | 8.00 | 8.91 | 8.72 | 7.82 |  |  |
| 7pzt (160) | RMSD ($\downarrow$) | Chroma* | 0.50 | 1.61 | 4.82 | 10.36 | 16.15 | 16.81 | 16.93 |
| ADP-3D | 0.23 | 0.67 | 2.64 | 9.39 | 14.26 | 15.47 | 14.07 |  |  |
| ELBO ($\uparrow$) | Chroma* | 5.99 | 5.55 | 5.39 | 7.47 | 8.52 | 8.74 | 9.04 |  |
| ADP-3D | 5.79 | 6.44 | 6.92 | 7.79 | 8.29 | 7.81 | 8.64 |  |  |
| 1a2f (291) | RMSD ($\downarrow$) | Chroma* | 0.39 | 1.21 | 3.42 | 7.45 | 14.77 | 16.99 | 17.29 |
| ADP-3D | 0.13 | 0.68 | 1.91 | 5.48 | 12.48 | 17.78 | 19.68 |  |  |
| ELBO ($\uparrow$) | Chroma* | 5.39 | 5.33 | 6.83 | 7.09 | 8.74 | 9.28 | 9.48 |  |
| ADP-3D | 5.24 | 6.54 | 7.62 | 7.85 | 8.20 | 8.27 | 9.12 |  |  |

### E.2 Gradient Descent for Linear Constraints

In Figure 6, we compare different gradient-based techniques for minimizing over $\mathbf{z}$ the following objective function:

$$ $\mathcal{L}(\mathbf{z})=\|\mathbf{M}\mathbf{x}^{*}-\mathbf{M}\mathbf{R}\mathbf {z}\|_{2}^{2}.$ (24) $$

$\mathbf{x}^{*}$ is the backbone structure of the ATAD2 protein (PDB:7qum). $\mathbf{M}$ is a masking matrix with a subsampling factor of 2. The preconditioning strategy is described in Section 4.2. Gradient descent with preconditioning and momentum leads to the fastest convergence.

Figure: Figure 6: Gradient Descent for Linear Constraint. Convergence speed of different techniques on a linear inverse problem (structure completion with a subsampling factor of 2), without the diffusion prior. Using preconditioning with momentum leads to the fastest convergence. The “loss” corresponds to the negative log-likelihood.
Refer to caption: x6.png

### E.3 Ablation Study for Model Refinement

In Figure 7, we analyze the importance of the different input measurements for atomic model refinement. Removing the partial atomic model leads to the largest drop in accuracy. The cryo-EM density map is the second most important measurement, followed by the generative prior and the sequence.

Figure: Figure 7: Ablation Study. For atomic model refinement, we combine three sources of conditioning information (incomplete model, density map and sequence) with the data-driven prior of the diffusion model. Here we highlight the importance of each conditioning information, and that of the generative prior. (a) Qualitative reconstructions with the target structure in transparency. (b) RMSD of alpha carbons vs. completeness. We use the same structure as in Fig. 4 (PDB:7pzt), and a cryo-EM map at 2.0 Å resolution.
Refer to caption: x7.png

### E.4 Additional video

We provide one additional video showing the predicted structure throughout the diffusion process, for the three tasks explored in this paper.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
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.
soundness: 2
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.
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.
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).

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
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.
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 2
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
soundness: 3
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
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
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)

Meta-review: 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.

Decision: Reject