Matching Plug-and-Play Algorithms to the Denoiser
Abstract
To solve inverse problems, plug-and-play (PnP) methods have been developed that replace the proximal step in a convex optimization algorithm with a call to an application-speciﬁc denoiser, often implemented using a deep neural network (DNN).
Although such methods have been successful, they can be improved.
For example, denoisers are usually designed/trained to remove white Gaussian noise, but the denoiser input error in PnP algorithms is usually far from white or Gaussian.
Approximate message passing (AMP) methods provide white and
Gaussian denoiser input error, but only when the forward operator is a large random matrix.
In this work, we propose a PnP algorithm based on generalized expectation consistent (GEC) approximation that offers predictable error statistics at each iteration, as well as a new DNN denoiser that leverages those statistics.
Introduction
Linear inverse problems aim to recover a signal x 0 ∈ R N from measurements y ∈ R M of the form
[Equation]
where A is a known linear operator and w is unknown noise.
Well known examples of linear inverse problems include deblurring; superresolution; inpainting; image recovery in magnetic resonance imaging (MRI), computed tomography or microscopy; and decoding in communications.
In many cases, one has knowledge of the class of signals from which x 0 is drawn.
If harnessed, such knowledge can be extremely beneﬁcial to signal recovery.
In fact, in problems where A is not full column rank, such knowledge is essential because many x 0 explain y even in the absence of noise.
The traditional approach to solving inverse problems relied on an expert-based model for the signal class (e.g., sparsity in a known basis) and used an algorithm to jointly enforce that model, together with (1).
The primary challenge with this approach is that some signal classes (e.g., natural images) are difﬁcult to model.
A more recent approach is to train a deep neural network (DNN) to predict x 0 from y .
Although this method has shown promise, it requires a huge number of ( x 0 , y ) pairs for training, which may be unavailable in some applications.
Also, the trained model may not generalize to test scenarios where A and/or the statistics of w have changed.
So-called “plug-and-play” (PnP) approaches give a middle-ground between traditional algorithmic approaches and end-to-end DNNs.
There, a DNN is trained as a signal denoiser, and that denoiser is called iteratively by an algorithm that knows A and the statistics of w .
Convolutional DNN denoisers can be trained using only patches of the signal, so that relatively few full signal examples are required.
Also, because the denoiser is trained independently of A and w , generalization is not an issue.
To understand the conventional approach to PnP algorithm design, it helps to ﬁrst pose the inverse problem as an optimization problem of the form
[Equation]
where g 1 ( x ) and g 2 ( x ) promote measurement and signal-class ﬁdelity, respectively.
For example,
Algorithm 1 ADMM
Require: g 1 ( · ) , g 2 ( · ) and γ .
[Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation]
is an appropriate choice when w is white Gaussian noise (WGN) with precision γ w .
For convex g 1 and g 2 , many algorithms have been proposed to solve (2).
One of the best known is ADMM [1], shown as Alg. 1.
In lines 3–4, γ is a positive stepsize that affects convergence speed but not the ﬁxed point, and
[Equation]
Using (3), line 4 can be interpreted as MAP estimation of x 0 with prior p ( x 0 )
∝ e − g 2 ( x 0 ) from an observation r = x 0 + e with γ -precision AWGN e , i.e.,
MAP denoising.
This observation led Bouman et al.
to propose that the prox in line 4 be replaced by state-of-the-art denoiser f 2 : R N → R N , giving rise to PnP-ADMM [2].
The PnP formula was then applied to other algorithms, such as primal-dual splitting (PDS) in [3].
Regularization-bydenoising (RED) [4,5] yields a related class of algorithms that iteratively call a denoiser.
See [6] for a comprehensive overview of PnP and RED.
With a well-designed DNN denoiser f 2 , PnP and RED signiﬁcantly outperform optimization-based approaches, as well as end-to-end DNNs in limited-data and mismatched- A scenarios (see, e.g., [6]).
However, there is room for improvement.
For example, while DNN denoisers are typically trained to remove the effects of additive WGN (AWGN), PnP and RED algorithms do not provide the denoiser with an AWGN-corrupted input at each iteration.
Thus, one might wonder: Is it possible to design an iterative denoising algorithm that presents the denoiser with AWGN at each iteration?
Approximate Message Passing
For Lipschitz denoisers f 2 , there is indeed a family of iterative denoising algorithms that presents the denoiser with WGN of known variance at each iteration: approximate message passing (AMP) algorithms like [7, 8].
In fact, when used with an MMSE denoiser f 2 , one can prove that these algorithm converge (exponentially fast) to the MMSE estimate of x 0 from y [9,10].
But there is a catch: for these algorithms to work as intended, the forward operator A should be large and random.
For example, the original AMP algorithm [7] wants A to be large and i.i.d.
sub-Gaussian, while the
VAMP algorithm [8] imposes the looser requirement that A be large and right-orthogonally invariant (ROI), i.e., that the right singular-vector (RSV) matrix of A should be drawn uniformly from the set of orthogonal matrices.
Essentially, these algorithms want that the RSV matrix of A randomly rotates the denoiser’s output error at each iteration, rendering it statistically equivalent to WGN.
In many linear inverse problems, however, A does not have sufﬁcient randomness for these AMP algorithms to work as intended.
For example, in many imaging applications, A is constructed from a
Fourier operator, but the denoiser output error also has signiﬁcant Fourier structure, and so the desired randomization does not happen.
As a result, when these AMP algorithms are used with Fourierstructured A and x 0 , they tend to diverge.
Although heuristic damping and initialization strategies can be applied to stabilize VAMP in these applications, and even outperform PnP-ADMM [11], they do not yield denoiser input errors that are white and Gaussian at each iteration.
That said, progress in VAMP algorithm design with Fourier-structured A and natural-image x 0 has been made in [12–14].
These works all recover the wavelet coefﬁcients of the image rather than the image pixels.
The motivation is that the Fourier-Wavelet matrix is known to be approximately block diagonal [15], where the blocks are formed by the wavelet subbands.
Furthermore, the subblocks of the Fourier-Wavelet matrix behave as randomizing operators for the wavelet coefﬁcients in each wavelet subband.
This gives the potential to design a VAMP-like algorithm with WGN error in each subband, as long as differences across subbands are properly handled.
In the ﬁrst of these works, Schniter, Rangan, and Fletcher [12] used VAMP to recover the whitened wavelet coefﬁcients in an attempt to equalize the subband error variances, but did not track the error variances over the iterations.
Later, for the special case of variable-density (VD) subsampled Fourier measurements, Millard et al.
[13] proposed a VAMP modiﬁcation called VDAMP that uses wavelet thresholding to track the subband error variances across iterations.
More recently, Metzler and
Wetzstein [14] extended VDAMP to generic denoising, and designed a DNN denoiser that exploits known error-variance within each wavelet subband.
The resulting technique, coined “denoisingVDAMP” (D-VDAMP), yielded subband errors that were white and Gaussian with predictable variance at each iteration, and showed a signiﬁcant boost in MRI image recovery performance over
VDAMP, PnP-ADMM, and RED [14].
Although D-VDAMP works reasonably well, it is based on a heuristic modiﬁcation of VAMP with poor ﬁxed points.
Also, it is not clear how to use D-VDAMP with forward operators other than
VD-subsampled Fourier.
These issues motivate our proposed approach, which is described next.
Proposed Approach
Our approach uses the GEC framework from [16], as summarized in Alg. 2.
GEC is similar to VAMP,
Algorithm 2 Generalized EC (GEC)
Require: f 1 ( · , · ) , f 2 ( · , · ) , and gdiag( · ) .
1: Select initial r 1 , γ2: repeat 3: //
Measurement ﬁdelity 4: (cid:98) x 1 ← f 1 ( r 1 , γ 1 ) 5: η 1 ← Diag(gdiag( ∇ f 1 ( r 1 , γ 1 ))) − 1 γ6: γ 2 ← η 1 − γ7: r 2 ← Diag( γ 2 ) − 1 (Diag( η 1 ) (cid:98) x 1 − Diag( γ 1 ) r 1 ) 8: //
Denoising 9: (cid:98) x 2 ← f 2 ( r 2 , γ 2 ) 10: η 2 ← Diag(gdiag( ∇ f 2 ( r 2 , γ 2 )))
− 1 γ11: γ 1 ← η 2 − γ12:
r 1 ← Diag( γ 1 ) − 1 (Diag( η 2 ) (cid:98) x 2 − Diag( γ 2 ) r 2 ) 13: until Terminated
but it averages the diagonal of the Jacobian ∇ f i separately over L coefﬁcient subsets using gdiag:
R N × N → R N : gdiag( Q ) (cid:44) [ d 1 1 T N 1 , . . .
, d L 1 T N L ] T (4a) d (cid:96) =
tr { Q (cid:96)(cid:96) }
N (cid:96) .
(4b)
In (4), N (cid:96) is the size of the (cid:96) th subset and Q (cid:96)(cid:96)
∈ R N (cid:96) × N (cid:96) is the (cid:96) th diagonal subblock of the matrix input Q .
In Alg. 2,
Diag( · ) creates a diagonal matrix from its vector argument.
When L = 1 , GEC reduces to VAMP.
When GEC is used to solve a convex optimization problem of the form (2), the functions f i take the form f i ( r , γ ) =
gprox g i , γ ( r ) for gprox
[Equation]
(5) where (cid:107) q (cid:107) γ
(cid:44) (cid:112) q H Diag( γ ) q .
Note that if γ = γ 1 , then gprox ρ, γ = prox γ − 1 ρ .
Furthermore, if the γ i vectors are held ﬁxed over the iterations and take the form γ i = γ 1 for some γ > 0 , then GEC reduces to the “Peaceman Rachford” variant of ADMM, which uses two dual updates.
Thus, GEC can be interpreted an ADMM-like algorithm with adaptive vector-valued stepsizes, γ 1 and γ 2 .
For our proposed “denoising GEC” (D-GEC) approach, we use a DNN denoiser for f 2 .
Also, we focus on WGN w (i.e., quadratic log-likelihood), which gives
[Equation]
For both f 1 and f 2 , we approximate the tr { Q (cid:96)(cid:96) } term in (4) using the Monte Carlo approach [17]
[Equation]
where the coefﬁcients in the (cid:96) th subset of q (cid:96) are i.i.d
unit-variance Gaussian and the others are zero.
For Fourier-structured A , we recover the wavelet coefﬁcients c
= Ψ x 0 of x 0 , like in [12–14].
For an orthogonal wavelet transform Ψ , we have that x 0 = Ψ T c 0 , and so (1) implies the model
[Equation]
To apply D-GEC to c 0 recovery, we use f 1 as in (6), but with B in place of A .
For the diagonalization subsets in (4), we choose the L subbands of a depth- D wavelet transform.
As in [14], we denoise the wavelet coefﬁcient estimate r 2 by converting to the pixel domain, denoising in that domain, and converting back to the wavelet domain.
In the pixel-domain, GEC predicts that the noise will be correlated with correlation matrix Ψ Diag( γ ) − 1 Ψ T .
To handle this correlated noise, we propose a new approach to denoising called “corr+corr,” which presents the (pixel-domain) denoiser with the noisy image in one channel and an independent noise realization in another.
During training, the denoiser learns to extract statistics from this realization and use them productively for denoising.
Numerical Experiments
We compare the proposed D-GEC algorithm to the existing D-VDAMP [14] and PnP-PDS [3] algorithms for MRI image recovery.
For C -coil MRI, we constructed A =
[ A T 1 , . . .
, A T C ] T with
A c = MF Diag( s c ) ∈
C M × N , where F is the unitary 2D Fourier matrix, sampling mask M contains M rows of I N , and s c ∈ C N is a coil sensitivity map generated from the Biot-Savart Law.
In the single-coil case, C =1 and s 1 = 1 .
Like in [14], we used VD masks (see Fig.
4 of Appendix).
Denoisers: For D-VDAMP, we used the denoiser proposed in [14], which is a modiﬁcation of
DnCNN [18] that accepts the noise standard deviation (SD) in each wavelet subband.
Five versions of this denoiser were trained using independent subband noise SDs uniformly distributed in the ranges 0-10, 10-20, 20-50, 50-120, and 120-500, respectively.
(Pixel values ranged from 0-255.)
For PnP-PDS, we trained bias-free DnCNN [19] with white noise of SD uniformly distributed in the range 0-55.
For D-GEC, we trained bias-free DnCNN using the previously described corr+corr approach using independent subband SDs uniformly distributed in the ranges 0-10, 10-20, 20-50, 50-120, and 120-500, respectively.
All DNNs were trained using patches from 70 MRI images of the
Stanford 2D FSE dataset available at http://mridata.org .
All denoisers use only the real part of the input and generate a real-valued output.
Test data: For evaluation, we used ten 352 × 352 MRI images from the Stanford 2D FSE dataset that were not used as training images (see Fig.
3 of Appendix).
The measurements y were constructed using (1) with complex AWGN w whose variance was adjusted to give a pre-masking SNR of 40 dB.
Algorithm parameters: For D-GEC and D-VDAMP, we used a 2D Haar wavelet transform with
D =4 levels, giving L =13 wavelet subbands.
D-VDAMP code was obtained from the authors and run under default settings, which are detailed in [14].
PnP-PDS was run for 200 iterations using the parameters that maximized PSNR on the training set.
For D-GEC, we used the damping scheme from [11] and the auto-tuning scheme from [20].
Single-coil results: Table 1 shows that D-GEC signiﬁcantly outperformed D-VDAMP and PnP-PDS
[Table]
coil
[Table]
in all experiments.
Figureshows the evolution of the SD of the error at the input to DGEC’s denoiser in each subband.
There, we see a close agreement between true and predicted values.
Figure 2 shows QQ plots that indicate the subband errors are Gaussian.
In the Appendix, Fig. 5 shows an example of the wavelet coefﬁcients input to D-GEC’s denoiser at the 10th iteration, and their error relative to the true coefﬁcients, while Fig.
6 shows image recoveries and error maps for one of the test images at M/N =1 / 4 .
Multi-coil results: Table 1 shows D-GEC signiﬁcantly outperforming PnP-PDS in PSNR and SSIM in the 4-coil case.
D-VDAMP does not support multi-coil MRI and thus is not shown in the table.
Conclusion
We designed a GEC-based PnP algorithm that aims to present the denoiser with subband errors that are white and Gaussian with known variance at each iteration, and a denoiser that handles the resulting noise, which has a known correlation structure in the pixel domain.
Experiments with MRI image recovery show signiﬁcantly improved recovery performance relative to previous PnP methods.