Technical Summary: Physics Matters in PnP: Recovery Guarantees with the MMSE and NN Denoisers

1. Problem Statement

The paper addresses the linear ill-posed inverse problem of recovering a random vector $x \in \mathbb{R}^n$ from an observation $y \in \mathbb{R}^m$ governed by the forward model:
$$y = Ax + e$$
where $A: \mathbb{R}^n \to \mathbb{R}^m$ is a linear operator, and $e$ is Gaussian noise with a zero-mean and a possibly non-diagonal, positive definite covariance matrix $\Sigma_E$. The problem is ill-conditioned and potentially non-unique if $A$ is not injective.

The authors investigate the Plug-and-Play (PnP) paradigm, specifically the forward-backward splitting (FBS) iteration. Unlike classical approaches that rely on explicit regularization functionals, PnP replaces the proximal operator with a denoising operator $\hat{D}$. The general iteration considered is:
$$x_{k+1} = \hat{D}(x_k + \gamma B(y - Ax_k))$$
where $\gamma > 0$ is a step size and $B \in \mathbb{R}^{n \times m}$ is a linear operator. The central challenge addressed is establishing rigorous recovery guarantees (error bounds) for this iteration when the denoiser is the Minimum Mean Square Error (MMSE) estimator, and extending these results to cases where the MMSE is approximated by a neural network (NN).

A critical distinction of this work is the treatment of the noise distribution. While many existing PnP analyses assume the denoiser is trained on isotropic Gaussian noise ($\sigma^2 I$), this paper considers degenerate Gaussian noise with general covariance structures, acknowledging that the noise injected into the denoiser during the PnP iteration depends on the forward operator $A$ and the choice of $B$.

2. Methodology

2.1 The MMSE Denoiser Framework

The authors define the optimal denoiser $\hat{D}$ as the MMSE estimator of $X$ given a noisy observation $Z = X + \Upsilon$, where $\Upsilon \sim \mathcal{N}(0, \Sigma_\Upsilon)$. Crucially, $\Sigma_\Upsilon$ is not assumed to be isotropic; it is allowed to be degenerate (rank-deficient) and is chosen to match the distribution of the noise term $\gamma B e$ generated during the PnP iteration.

The MMSE estimator is defined as the posterior conditional mean:
$$\hat{D}_\mu(z) = \mathbb{E}[X | Z=z] = \frac{\int x \phi(z-x; \tilde{\Sigma}_\Upsilon) d\mu(x)}{\int \phi(z-x; \tilde{\Sigma}_\Upsilon) d\mu(x)}$$
where $\phi$ is the Gaussian density kernel adapted to the subspace $V = \text{Im}(\Sigma_\Upsilon)$.

2.2 Key Assumptions

The analysis relies on several structural assumptions:

• Compactness: The support of the prior distribution $\mu_X$ is compact and bounded by $M$.
• Subspace Consistency: The support of $\mu_X$ lies within the subspace $V$ spanned by the noise covariance $\Sigma_\Upsilon$.
• Operator Properties: The matrix $BA$ is symmetric and positive semidefinite. The operator $B$ maps into $V$.
• Restricted Injectivity: The intersection of the kernel of $BA$ and the tangent cone of the convex hull of the support of $\mu_X$ is trivial, ensuring invertibility on the relevant set.

2.3 Analysis of Neural Network Approximations

Recognizing that the exact MMSE estimator is computationally intractable in high dimensions, the authors extend their analysis to denoisers parametrized by ReLU neural networks. They utilize approximation theory to bound the necessary width, depth, and weight magnitudes of a network required to approximate the identity map (and thus the MMSE denoiser) on the support of the prior with a specific precision $\epsilon$.

3. Key Contributions

The paper makes three primary contributions, moving beyond the "black-box" view of PnP denoisers:

1. Regularity and Stability of MMSE Denoisers:
   The authors establish that the MMSE denoiser is Lipschitz continuous and, under specific conditions (Gaussian noise with non-degenerate covariance), cocoercive. This provides a rigorous theoretical justification for the "firm non-expansiveness" often enforced heuristically in PnP methods. They also prove the stability of the denoiser with respect to the prior measure (Wasserstein continuity).

2. Recovery Guarantees for PnP with MMSE:
   The paper derives pointwise and Wasserstein distance recovery error bounds for the PnP-FBS iteration.

   • Pointwise Bounds: The error bound depends on the condition number of the operator $BA$ restricted to the tangent cone of the solution set, the noise level, and the sampling error of the empirical prior.
   • Wasserstein Bounds: The distribution of the iterates converges to the true prior distribution. The bound explicitly includes a Bures metric term, which measures the discrepancy between the noise distribution used to train the denoiser ($\Upsilon$) and the effective noise in the iteration ($\gamma B E$).

3. The Necessity of Physics-Aware Denoisers:
   A central theoretical finding is that the denoiser cannot be chosen in a physics-agnostic way. The choice of the operator $B$ and the noise covariance $\Sigma_\Upsilon$ used to construct the denoiser must be coupled. Specifically, for optimal performance, the denoiser must be trained on noise statistics that match the transformed observation noise $\gamma B E$. The standard choice $B = A^\top$ implicitly assumes the denoiser is agnostic to the forward model, which the authors show leads to suboptimal recovery bounds unless specific conditions are met.

4. Main Results

4.1 Pointwise Recovery

Theorem 3.2 provides a bound on the reconstruction error $\|x_k - x\|$. The error converges to a neighborhood of the solution determined by:

• Noise Amplification: Proportional to $\gamma \|B\| \|e\|$.
• Denoiser Performance: Related to the distance between the true signal and the empirical training samples.
• Conditioning: The convergence rate $q$ depends on the condition number of $BA$ on the solution set.
  Crucially, the bound holds even if the denoiser is not contractive, provided the condition number of $BA$ is sufficiently favorable relative to the Lipschitz constant of the denoiser.

4.2 Distributional Recovery

Theorem 3.6 establishes bounds on the expected squared Wasserstein distance $W_2^2(\hat{\mu}_k, \mu_X)$ between the distribution of the iterates and the true prior. The bound reveals that the error is minimized when the covariance of the training noise $\Sigma_\Upsilon$ matches the covariance of the iteration noise $\gamma^2 B \Sigma_E B^\top$. If this match is perfect, the Bures metric term vanishes, significantly improving the bound.

4.3 Neural Network Approximation

Theorem 3.11 extends the pointwise bounds to the case where the MMSE denoiser is replaced by a trained neural network. The error bound includes terms for:

• The approximation error of the network (dependent on $\epsilon$).
• The Lipschitz constants of the network.
• The sampling error of the training data.
  The authors derive specific requirements for the network's depth and width to achieve a desired approximation precision, showing that the recovery guarantees hold provided the network is sufficiently expressive.

5. Significance and Claims

The paper claims that its analysis fundamentally shifts the understanding of PnP methods by demonstrating that physics matters.

• Rejection of Agnosticism: The authors argue that the common practice of training denoisers on generic noise (e.g., isotropic Gaussian) and applying them to inverse problems with arbitrary forward operators $A$ and noise covariances $\Sigma_E$ is theoretically flawed. The denoiser must be "physics-aware," meaning its training noise distribution must align with the effective noise introduced by the operator $B$ in the PnP iteration.
• Relaxed Assumptions: By leveraging the specific properties of the MMSE estimator (cocoercivity and Lipschitz bounds derived from the prior), the authors derive recovery guarantees without requiring the strong assumption that the denoiser is a contraction mapping or a proximal operator of a convex functional.
• Quantitative Bounds: The paper provides explicit, non-asymptotic error bounds that quantify the trade-offs between the conditioning of the inverse problem, the quality of the denoiser (and its training data), and the noise level.

In conclusion, the work provides a rigorous theoretical foundation for Plug-and-Play methods, emphasizing that the design of the denoiser and the choice of the splitting operator $B$ are inextricably linked to the forward model and the noise statistics of the specific inverse problem being solved.