Adam-PnP: Diffusion-Based Protein Reconstruction

• Adam-PnP is a plug-and-play framework that reconstructs protein structures by guiding a pre-trained diffusion model with gradients from heterogeneous experimental sources.
• It integrates multiple modalities—such as partial coordinates, pairwise distances, and cryo-EM density—using dynamic weighting based on adaptive noise estimation.
• Empirical evaluations demonstrate improved folding fidelity with lower RMSD values when high-resolution data are used, while unreliable modalities are down-weighted.

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" (Banerjee et al., 28 Jul 2025), 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

$N$0

The central practical difficulty is that the $N$1 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 $N$2 from residuals while correcting for denoiser-induced bias, and dynamic modality weighting, which sets

$N$3

during reverse diffusion, thereby favoring modalities with higher estimated precision (Banerjee et al., 28 Jul 2025).

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

$N$4

and

$N$5

where $N$6 is the score network.

Chroma employs a non-isotropic diffusion process to preserve biophysical constraints. Its noise covariance is represented by a matrix $N$7, with whitening latent $N$8. The Chroma forward SDE is

$N$9

and the hybrid Langevin reverse SDE is

$m \in \{1,\dots,M\}$0

Within this prior, Adam-PnP performs guidance in a discrete-time reverse loop. The learned reverse step is augmented by data-consistency gradients:

$m \in \{1,\dots,M\}$1

where $m \in \{1,\dots,M\}$2. The method applies this guidance in the whitened latent space. At time $m \in \{1,\dots,M\}$3, it first computes a denoised latent estimate

$m \in \{1,\dots,M\}$4

then performs a weighted gradient step to obtain $m \in \{1,\dots,M\}$5, which is used in the next reverse propagation (Banerjee et al., 28 Jul 2025).

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

$m \in \{1,\dots,M\}$6

under the Gaussian log-likelihood

$m \in \{1,\dots,M\}$7

A direct residual-based estimator for $m \in \{1,\dots,M\}$8 is biased because the residual contains both measurement noise and denoiser error:

$m \in \{1,\dots,M\}$9

If $$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$0 is $$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$1-Lipschitz and $$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$2 approximates $$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$3, the paper states

$$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$4

To address this, Adam-PnP uses a bias-corrected robust estimator with EMA:

$$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$5

$$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$6

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 $$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$7 reduces bias as denoiser accuracy improves over time.

These estimated variances directly scale the modality-specific likelihood gradients:

$$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$8

with $$y_m = A_m(x^\ast) + n_m, \qquad n_m \sim \mathcal{N}(0,\sigma_m^2 I),$$9 substituted for the unknown $A_m$0. Dynamic weights are then assigned according to estimated precision,

$A_m$1

This normalization preserves overall guidance magnitude while prioritizing reliable modalities (Banerjee et al., 28 Jul 2025).

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 $A_m$2 and a single global Polyak momentum on the aggregated guidance. Under modality conflict, residual growth increases $A_m$3 and reduces the corresponding $A_m$4, 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 $A_m$5 in the whitened latent variable. Its inputs are measurements $A_m$6, operators $A_m$7, a pre-trained denoiser $A_m$8, and whitening matrix $A_m$9. The latent is initialized as $y_m$0, with momentum $y_m$1 and optional initialization of $y_m$2 from initial residuals. Each reverse step consists of the following operations (Banerjee et al., 28 Jul 2025):

1. Prior projection: $y_m$3.
2. Residual computation: $y_m$4.
3. Adaptive noise estimation: compute $y_m$5 by median residual statistics, then update $y_m$6 by EMA with bias correction.
4. Dynamic weighting: set $y_m$7 with $y_m$8.
5. Guidance aggregation: form

$y_m$9

1. Global momentum and guidance step:

$\sigma_m$0

2. Reverse propagation:

$\sigma_m$1

followed by the final output $\sigma_m$2.

The experiments instantiate three differentiable modalities:

Modality	Operator
Partial C$\sigma_m$3 coordinates (P)	$\sigma_m$4
Pairwise C$\sigma_m$5 distances (D)	$\sigma_m$6 returns selected distances $\sigma_m$7
Simulated low-resolution cryo-EM density (E)	$\sigma_m$8 renders a low-resolution electron density map or its Fourier coefficients

For the coordinate modality,

$\sigma_m$9

For pairwise distances, the likelihood is

$$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$0

with coordinate gradients derived from

$$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$1

For cryo-EM density,

$$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$2

All guidance is transferred to latent space by the chain rule,

$$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$3

Relative to unguided diffusion, the overhead is dominated by per-modality forward evaluations $$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$4 and Jacobian-vector products or autodiff for the likelihood gradients. Additional memory consists mainly of the global momentum vector $$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$5 and per-modality variance estimates $$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$6. The overall cost scales roughly linearly with the number of modalities $$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$7, plus renderer cost for complex operators such as cryo-EM (Banerjee et al., 28 Jul 2025).

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$$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$8 RMSD in ångströms:

$$p(x \mid \{y_m\}) \propto p_\theta(x)\prod_{m=1}^M p(y_m \mid x,\sigma_m),$$9

The prior is a pre-trained Chroma diffusion model, and comparisons are made across modality combinations with fixed schedules (Banerjee et al., 28 Jul 2025).

Modality combination	Backbone C$p_\theta(x)$0 RMSD (Å)
P	$p_\theta(x)$1
D	$p_\theta(x)$2
E	$p_\theta(x)$3
P + D	$p_\theta(x)$4
P + D + E	$p_\theta(x)$5
P + E	$p_\theta(x)$6
D + E	$p_\theta(x)$7

Among the reported combinations, $p_\theta(x)$8 gives the best mean RMSD, while $p_\theta(x)$9 alone fails to fold. The paper states that low-resolution $$p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$0 alone or with $$p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$1 is weak, and that Adam-PnP dynamically down-weights $$p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$2 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(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$3 active shows that increasing high-resolution data, specifically more distances and partial coordinates, monotonically improves RMSD and reaches $$p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$4 Å at high coverage and low noise $$p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$5 Å). For adaptive variance estimation on the distance modality, the reported true-versus-estimated values are:

• $$p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$6 with absolute error $$p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$7,
• $$p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$8 with absolute error $$p(y_m \mid x,\sigma_m) \propto \exp\!\left(-\frac{1}{2\sigma_m^2}\|y_m-A_m(x)\|^2\right).$$9,
• $N$00 with absolute error $N$01.

Qualitative examples are summarized as improved folding fidelity when fusing $N$02 with $N$03, whereas $N$04 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 (Banerjee et al., 28 Jul 2025).

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" (Shoushtari et al., 2023), “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 (Banerjee et al., 28 Jul 2025). 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 $N$05 and $N$06; 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 (Banerjee et al., 28 Jul 2025).