---
title: 'Adam-PnP: Diffusion-Based Protein Reconstruction'
url: https://www.emergentmind.com/topics/adam-pnp
type: topic
---

# Adam-PnP: Diffusion-Based Protein Reconstruction

Adam-PnP is a Plug-and-Play framework for protein structure reconstruction that guides a pre-trained protein diffusion model with gradients from multiple, heterogeneous experimental sources while adapting to unknown and modality-specific noise levels on the fly. In the formulation introduced in "Adaptive Multimodal Protein Plug-and-Play with Diffusion-Based Priors" [2507.21260], the method targets inverse problems in which a protein backbone is represented by 3D coordinates and observed through lossy or noisy measurement operators, and it combines a diffusion-based structural prior with adaptive noise estimation, dynamic modality weighting, and guided reverse-time sampling.

## 1. Problem formulation and reconstruction objective

Adam-PnP models a protein backbone by coordinates $x \in \mathbb{R}^{4N\times 3}$ for $N$ residues. For each modality $m \in \{1,\dots,M\}$, the observation model is
$$
y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),
$$
where $A_m$ is a known differentiable operator, $y_m$ may have modality-specific dimensionality, and the noise level $\sigma_m$ is unknown. The corresponding Bayesian posterior factorizes as
$$
p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),
$$
with $p_\theta(x)$ induced by a pre-trained diffusion model and Gaussian likelihoods of the form
$$
p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).
$$
The MAP estimate is therefore
$$
\hat{x}=\arg\min_x\left\{-\sum_m \log p(y_m\mid x,\sigma_m)-\log p_\theta(x)\right\}.
$$

The central practical difficulty is that the $\sigma_m$ are both unknown and heterogeneous across modalities. Classical approaches often require manual weighting of modality-specific guidance terms to balance gradient magnitudes. In the Adam-PnP construction, this is treated as brittle and labor-intensive when the noise profile is misestimated. The framework addresses the issue through two coupled mechanisms: adaptive noise estimation, which estimates modality-specific $\sigma_m$ from residuals while correcting for denoiser-induced bias, and dynamic modality weighting, which sets
$$
w_m(t)\propto \frac{1}{\hat{\sigma}_m^2(t)+\epsilon}
$$
during reverse diffusion, thereby favoring modalities with higher estimated precision [2507.21260].

## 2. Diffusion prior and guided reverse dynamics

The prior is implemented with a pre-trained score-based SDE prior, specifically the Chroma backbone model, although the framework is described as compatible with DDPMs and other score-based models. In the generic score-based setting, the forward and reverse SDEs are
$$
d x_t = f(t)x_t\,dt + g(t)\,dW_t,
$$
and
$$
d x_t = [f(t)x_t - g(t)^2 s_\theta(x_t,t)]\,dt + g(t)\,d\hat{W}_t,
$$
where $s_\theta(x_t,t)\approx \nabla_{x_t}\log p_t(x_t)$ is the score network.

Chroma employs a non-isotropic diffusion process to preserve biophysical constraints. Its noise covariance is represented by a matrix $R\in\mathbb{R}^{4N\times 4N}$, with whitening latent $z=R^{-1}x$. The Chroma forward SDE is
$$
d x_t = -(\beta_t/2)x_t\,dt + \sqrt{\beta_t}\,R\,dW_t,
$$
and the hybrid Langevin reverse SDE is
$$
d x_t =
\left[
-\frac{1}{2}x_t -
\left(\lambda_t+\frac{\lambda_0\psi}{2}\right)RR^T s_\theta(x_t,t)
\right]\beta_t\,dt
+ \sqrt{\beta_t(1+\psi)}\,R\,d\hat{W}_t.
$$

Within this prior, Adam-PnP performs guidance in a discrete-time reverse loop. The learned reverse step is augmented by data-consistency gradients:
$$
x_{t-1}=F_\theta(x_t,t)-\eta_t\sum_m w_m(t)\nabla_x \phi_m(x_t;y_m,\sigma_m(t)),
$$
where $\phi_m(x;y_m,\sigma_m)=-\log p(y_m\mid x;\sigma_m)$. The method applies this guidance in the whitened latent space. At time $t$, it first computes a denoised latent estimate
$$
\tilde{z}_0 \leftarrow R^{-1}D_\theta(Rz_t,t),
$$
then performs a weighted gradient step to obtain $\hat{z}_0$, which is used in the next reverse propagation [2507.21260].

This embedding of Plug-and-Play guidance into reverse diffusion is structurally analogous to diffusion posterior sampling, but the Adam-PnP variant is explicitly multimodal, variance-adaptive, and executed in a whitened latent. A plausible implication is that the latent-space implementation is not merely computationally convenient; it is coupled to Chroma’s non-isotropic noise model and thus to the geometry of the protein prior.

## 3. Adaptive noise estimation and dynamic modality weighting

For each modality, Adam-PnP uses residuals
$$
r_m(x):= y_m-A_m(x)
$$
under the Gaussian log-likelihood
$$
\log p(y_m\mid x;\sigma_m)= -\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2 + C.
$$
A direct residual-based estimator for $\sigma_m^2$ is biased because the residual contains both measurement noise and denoiser error:
$$
r_m = n_m - \bigl(A_m(R\tilde{z}_0)-A_m(Rz^\ast)\bigr).
$$
If $A_m$ is $L_m$-Lipschitz and $\tilde{z}_0$ approximates $z^\ast$, the paper states
$$
\mathbb{E}[\|r_m\|^2] \le \sigma_m^2 + c_t(L_m\tau_t)^2.
$$

To address this, Adam-PnP uses a bias-corrected robust estimator with EMA:
$$
\tilde{\sigma}_m^2(t):=\operatorname{median}\!\left([y_m-A_m(R\tilde{z}_0)]^2\right),
$$
$$
\hat{\sigma}_m^2(t):=\operatorname{EMA}_t\!\left[\max\!\left(\epsilon,\tilde{\sigma}_m^2(t)-\gamma(1-t/T)(L_m\tau_t)^2\right)\right].
$$
The estimator is described as robust because it uses the median, and stabilized because it uses exponential moving average across timesteps. The annealed correction term $(1-t/T)(L_m\tau_t)^2$ reduces bias as denoiser accuracy improves over time.

These estimated variances directly scale the modality-specific likelihood gradients:
$$
\nabla_x \log p(y_m\mid x;\sigma_m(t))
=
\frac{1}{\sigma_m^2(t)}A_m'(x)^T r_m(x),
$$
with $\hat{\sigma}_m^2(t)$ substituted for the unknown $\sigma_m^2(t)$. Dynamic weights are then assigned according to estimated precision,
$$
w_m(t)\propto \frac{1}{\hat{\sigma}_m^2(t)+\epsilon}, \qquad \sum_m w_m(t)=M.
$$
This normalization preserves overall guidance magnitude while prioritizing reliable modalities [2507.21260].

The method’s name can invite confusion with Adam-style adaptive optimization. The paper explicitly states that Adam-PnP does **not** maintain per-modality first or second moments for the weights. Instead, it uses EMA inside $\hat{\sigma}_m^2(t)$ and a single global Polyak momentum on the aggregated guidance. Under modality conflict, residual growth increases $\hat{\sigma}_m^2(t)$ and reduces the corresponding $w_m(t)$, so unreliable modalities are automatically down-weighted across timesteps.

## 4. Sampling algorithm, modality operators, and computational profile

The full algorithm is a guided sampling loop over $t=T,T-1,\dots,1$ in the whitened latent variable. Its inputs are measurements $\{y_m\}$, operators $\{A_m\}$, a pre-trained denoiser $D_\theta$, and whitening matrix $R$. The latent is initialized as $z_T\sim \mathcal{N}(0,I)$, with momentum $v_T=0$ and optional initialization of $\hat{\sigma}_m^2(0)$ from initial residuals. Each reverse step consists of the following operations [2507.21260]:

1. **Prior projection**: $\tilde{z}_0 \leftarrow R^{-1}D_\theta(Rz_t,t)$.
2. **Residual computation**: $r_m \leftarrow y_m-A_m(R\tilde{z}_0)$.
3. **Adaptive noise estimation**: compute $\tilde{\sigma}_m^2$ by median residual statistics, then update $\hat{\sigma}_m^2(t)$ by EMA with bias correction.
4. **Dynamic weighting**: set $w_m(t)\leftarrow c\cdot 1/(\hat{\sigma}_m^2(t)+\epsilon)$ with $\sum_m w_m(t)=M$.
5. **Guidance aggregation**: form
   $$
   g_{\text{total}} \leftarrow \sum_m w_m(t)\nabla_z \log p(y_m\mid Rz;\hat{\sigma}_m(t))\big|_{z=\tilde{z}_0}.
   $$
6. **Global momentum and guidance step**:
   $$
   v_t \leftarrow \rho v_{t+1} + (1-\rho)g_{\text{total}}, \qquad
   \hat{z}_0 \leftarrow \tilde{z}_0 - \eta v_t.
   $$
7. **Reverse propagation**:
   $$
   z_{t-1}\sim \mathcal{N}(\alpha_{t-1|t}\hat{z}_0,\tau_{t-1|t}^2 I),
   $$
   followed by the final output $\hat{x}=Rz_0$.

The experiments instantiate three differentiable modalities:

| Modality | Operator |
|---|---|
| Partial C$\alpha$ coordinates (P) | $A_P(x)=Px$ |
| Pairwise C$\alpha$ distances (D) | $A_D(x)$ returns selected distances $d_{ij}(x)=\|x_i-x_j\|$ |
| Simulated low-resolution cryo-EM density (E) | $A_E(x)$ renders a low-resolution electron density map or its Fourier coefficients |

For the coordinate modality,
$$
\log p(y_P\mid x;\sigma_P)= -\frac{1}{2\sigma_P^2}\|Px-y_P\|^2 + C,
\qquad
\nabla_x \log p(y_P\mid x;\sigma_P)= \frac{1}{\sigma_P^2}P^T(y_P-Px).
$$
For pairwise distances, the likelihood is
$$
\log p(y_D\mid x;\sigma_D)= -\frac{1}{2\sigma_D^2}\sum_{(i,j)\in\Omega}(d_{ij}(x)-y_{ij})^2 + C,
$$
with coordinate gradients derived from
$$
\frac{\partial}{\partial x_i}d_{ij}(x)=\frac{x_i-x_j}{\|x_i-x_j\|},\qquad
\frac{\partial}{\partial x_j}d_{ij}(x)= -\frac{x_i-x_j}{\|x_i-x_j\|}.
$$
For cryo-EM density,
$$
\log p(y_E\mid x;\sigma_E)= -\frac{1}{2\sigma_E^2}\|A_E(x)-y_E\|^2 + C,
\qquad
\nabla_x \log p(y_E\mid x;\sigma_E)= \frac{1}{\sigma_E^2}A_E'(x)^T(y_E-A_E(x)).
$$
All guidance is transferred to latent space by the chain rule,
$$
\nabla_z \log p(y_m\mid Rz;\sigma_m)=R^T\nabla_x \log p(y_m\mid x;\sigma_m)\big|_{x=Rz}.
$$

Relative to unguided diffusion, the overhead is dominated by per-modality forward evaluations $A_m(R\tilde{z}_0)$ and Jacobian-vector products or autodiff for the likelihood gradients. Additional memory consists mainly of the global momentum vector $v_t$ and per-modality variance estimates $\hat{\sigma}_m^2(t)$. The overall cost scales roughly linearly with the number of modalities $M$, plus renderer cost for complex operators such as cryo-EM [2507.21260].

## 5. Empirical evaluation in protein reconstruction

The reported reconstruction task is PDB 7r5b with 127 residues, using combinations of the three modalities above. Performance is measured by aligned backbone C$\alpha$ RMSD in ångströms:
$$
\operatorname{RMSD}(\hat{x},x^\ast)=\min_{T\in SO(3)}\|T\hat{x}-x^\ast\|.
$$
The prior is a pre-trained Chroma diffusion model, and comparisons are made across modality combinations with fixed schedules [2507.21260].

| Modality combination | Backbone C$\alpha$ RMSD (Å) |
|---|---|
| P | $0.74 \pm 0.26$ |
| D | $1.17 \pm 0.95$ |
| E | $11.67 \pm 1.00$ |
| P + D | $0.65 \pm 0.18$ |
| P + D + E | $0.67 \pm 0.19$ |
| P + E | $0.76 \pm 0.27$ |
| D + E | $1.04 \pm 0.08$ |

Among the reported combinations, $P + D$ gives the best mean RMSD, while $E$ alone fails to fold. The paper states that low-resolution $E$ alone or with $D$ is weak, and that Adam-PnP dynamically down-weights $E$ when conflicting. This suggests that the adaptive weighting mechanism is functioning as a reliability estimator rather than as a fixed fusion rule.

A data sparsity ablation with $P + D + E$ active shows that increasing high-resolution data, specifically more distances and partial coordinates, monotonically improves RMSD and reaches $0.10$ Å at high coverage and low noise $(\sigma=0.05$ Å). For adaptive variance estimation on the distance modality, the reported true-versus-estimated values are:
- $0.05 \rightarrow 0.086 \pm 0.009$ with absolute error $0.036$,
- $0.10 \rightarrow 0.120 \pm 0.013$ with absolute error $0.020$,
- $0.20 \rightarrow 0.247 \pm 0.056$ with absolute error $0.047$.

Qualitative examples are summarized as improved folding fidelity when fusing $P$ with $D$, whereas $E$ alone produces unfocused conformations. The paper also reports failure cases: extremely low-quality or severely mis-specified modalities can still impede progress if their gradients dominate early, although adaptive weighting mitigates this effect.

## 6. Position within Plug-and-Play research, nomenclature, and limitations

Adam-PnP is explicitly situated within the Plug-and-Play prior tradition, in which a strong prior operator is alternated with a data-consistency step. In classical PnP, the prior operator is often a proximal map or denoiser; in Adam-PnP, the cycle is embedded in discretized reverse diffusion, with a denoiser-driven prior projection followed by multimodal likelihood guidance. The paper describes this as analogous to diffusion posterior sampling, but differentiated by multimodal guidance, variance adaptation, and operation in a whitened latent [2507.21260].

A nomenclature issue arises because the label “Adam-PnP” is not uniformly used in the literature. In "Prior Mismatch and Adaptation in PnP-ADMM with a Nonconvex Convergence Analysis" [2310.00133], “Adam-PnP” refers to ADMM-based Plug-and-Play, not the Adam optimizer. That work studies PnP-ADMM under prior mismatch, establishes a nonconvex convergence analysis for mismatched MMSE denoisers, and clarifies that ADMM-based PnP alternates a data-consistency step with denoising. This is a distinct usage from the protein-focused diffusion framework of [2507.21260]. The shared terminology reflects the broader Plug-and-Play lineage rather than a common algorithmic core.

The diffusion-based Adam-PnP paper does not provide formal convergence guarantees. It instead gives stability and convergence intuition grounded in Plug-and-Play principles: a strong pre-trained prior projection is alternated with a likelihood gradient step, while adaptive noise estimation and precision-based weighting mitigate gradient-scale mismatch and conflicting modalities. The reported limitations are specific. The method requires differentiable measurement operators and tractable Jacobian-vector products; bias correction depends on approximate $L_m$ and $\tau_t$; strong diffusion priors are needed under sparse or noisy conditions; and very low-resolution modalities may be down-weighted to near-irrelevance, limiting their contribution. Future work identified in the paper includes formal analysis of stability for multimodal adaptive guidance, per-modality adaptive moment schemes for weights, and improved physics-based renderers for cryo-EM and other modalities [2507.21260].

Source: https://www.emergentmind.com/topics/adam-pnp