Expectation Consistent Plug-and-Play for MRI

For image recovery problems, plug-and-play (PnP) methods have been developed that replace the proximal step in an optimization algorithm with a call to an application-specific denoiser, often implemented using a deep neural network. Although such methods have been successful, they can be improved. For example, the denoiser is often trained using white Gaussian noise, while PnP's denoiser input error is often far from white and Gaussian, with statistics that are difficult to predict from iteration to iteration. PnP methods based on approximate message passing (AMP) are an exception, but only when the forward operator behaves like a large random matrix. In this work, we design a PnP method using the expectation consistent (EC) approximation algorithm, a generalization of AMP, that offers predictable error statistics at each iteration, from which a deep-net denoiser can be effectively trained.

READ FULL TEXT VIEW PDF

page 1

page 3

page 4

06/09/2022

Denoising Generalized Expectation-Consistent Approximation for MRI Image Recovery

To solve inverse problems, plug-and-play (PnP) methods have been develop...
09/17/2018

Approximate message-passing for convex optimization with non-separable penalties

We introduce an iterative optimization scheme for convex objectives cons...
10/21/2020

MRI Image Recovery using Damped Denoising Vector AMP

Motivated by image recovery in magnetic resonance imaging (MRI), we prop...
07/04/2019

A New Insight into GAMP and AMP

A concise expectation propagation (EP) based message passing algorithm (...
02/25/2016

Expectation Consistent Approximate Inference: Generalizations and Convergence

Approximations of loopy belief propagation, including expectation propag...
02/10/2022

A Plug-and-Play Approach to Multiparametric Quantitative MRI: Image Reconstruction using Pre-Trained Deep Denoisers

Current spatiotemporal deep learning approaches to Magnetic Resonance Fi...
08/11/2018

Statistics of the Eigenvalues of a Noisy Multi-Soliton Pulse

For Nonlinear-Frequency Division-Multiplexed (NFDM) systems, the statist...

1 Introduction

Magnetic resonance imaging (MRI) is a medical imaging approach that uses magnetic fields to create detailed anatomical images. Although MRI provides excellent soft-tissue contrast without the use of ionizing radiation, it takes a long time to fully sample the measurement space. Thus, it is common to take relatively few measurements and apply sophisticated post-processing to reconstruct an accurate image. Although our paper focuses on MRI, the methods we propose apply to any application where the goal is to recover a signal from undersampled Fourier measurements.

The measurements collected in -coil MRI, known as “k-space” measurements, can be modeled as

(1)

where

is a vectorized version of the

-pixel image we wish to recover,

is a unitary 2D discrete Fourier transform (DFT),

is a sampling mask formed from

rows of the identity matrix

, is the th coil-sensitivity map, and is additive white Gaussian noise (AWGN) with precision

(i.e., variance

). In the special case of single-coil MRI, and , the all-ones vector. In this paper, we use the variable-density sampling masks like that shown in Fig. 1. In MRI, the pixels-to-measurement ratio, , is known as the “acceleration rate.” When , one cannot uniquely determine from due to the nullspace of , and so prior information about is needed for its recovery.

To recover images from undersampled MRI measurements, many methods have been proposed. Some are based on iterative optimization [1, 2], as described in the sequel. More recently, deep neural networks (DNNs) that directly map MRI measurements to an image have been proposed, e.g., [3, 4]. Although such DNNs work well, training them requires huge fully-sampled k-space datasets (which may be difficult or impossible to obtain) and changes in the acquisition parameters (e.g., sampling mask) from training to testing can degrade performance [5].

In this work, we focus on the “plug and play” approach that iteratively calls a DNN for image denoising, which brings several advantages. First, DNN denoisers can be trained using image patches, implying the need for relatively few images and no k-space data. Second, the denoiser is trained independently of the acquisition parameters, so that it generalizes to any acquisition scenario. Our approach is based on the generalized expectation consistent (GEC) approximation algorithm from [6], which lives in the family of approximate message passing (AMP) algorithms like [7, 8].

Figure 1: A variable-density sampling mask at .

2 Background

2.1 Compressed-sensing-based methods

The conventional approach to MRI image recovery [1, 2] is to pose and solve an optimization problem of the form

(2)

where promotes measurement fidelity and is an image-based regularizer. Typical choices are

(3)

for (1) and with a suitable transform (e.g., wavelet or total-variation) and carefully chosen . Such encourage sparsity in the transform coefficients .

Many algorithms have been proposed to solve (2) with convex and [2]. For example, the alternating directions method of multipliers (ADMM) [9] iterates

(4a)
(4b)
(4c)

where the proximal operator is defined as

(5)

In (4), is a tunable stepsize that affects the speed of ADMM’s convergence but not its fixed point.

2.2 Plug-and-play methods

The prox operator (5

) can interpreted as a denoiser, in particular, the maximum a posteriori (MAP) estimator of

with prior from an observation with -precision AWGN . Leveraging this fact, Bouman et al. [10] proposed to replace ADMM line (4b) with a call to a high-performance image denoiser like BM3D [11] or DnCNN [12], giving rise to “plug-and-play” (PnP) ADMM. PnP extensions of other algorithms, such as primal-dual splitting (PDS) [13] and proximal gradient (PG) [14], have also been proposed. As shown in the recent overview paper, PnP methods have been shown to significantly outperform compressed-sensing-based approaches in MRI [5]. Note, however, that when (4b) is replaced with a denoising step of the form “,” the stepsize does affect the fixed-point [5] and thus must be tuned.

Although PnP algorithms work well in MRI, there is room for improvement. For example, while image denoisers are typically designed/trained to remove the effects of AWGN, PnP algorithms do not provide the denoiser with an AWGN-corrupted input at each iteration. Rather, the denoiser’s input error has iteration-dependent statistics that are difficult to analyze or predict.

2.3 Approximate message passing

In (2), if we interpret as a log-likelihood and as a log prior, then can be interpreted as the MAP estimate of from . However, because image recovery results are often judged by mean-squared error (MSE), one may be more interested in the minimum MSE (MMSE) estimate of from . Interestingly, both MMSE and MAP estimation are facilitated by approximate message passing (AMP) methods like [7, 8].

For example, the AMP algorithm from [7] iterates

(6a)
(6b)
(6c)

over , starting from , where is a Lipschitz denoising function, is the trace of the Jacobian of at , and . When configured for MAP estimation, AMP uses the MAP denoiser . When configured for MMSE estimation, AMP instead uses the MMSE denoiser for with .

Importantly, when the forward operator is large and i.i.d. sub-Gaussian, and when is Lipschitz, the macroscopic behavior of AMP is rigorously characterized by a scalar state-evolution [15, 16]. When is the MMSE denoiser and the state-evolution has a unique fixed point, AMP provably converges to the MMSE-optimal estimate [15, 16]. For images, the MMSE denoiser can be approximated by BM3D or a DNN, as proposed in [17], leading to “denoising-AMP” (D-AMP). There, the trace-Jacobian in (6a) is approximated using the Monte-Carlo approach [18]

(7)

with random and small .

More recently, the vector AMP (VAMP) algorithm [8] was proposed, with similar properties as AMP (e.g., rigorous state evolution and provable MMSE estimation) but applicability to a wider class of random matrices: right orthogonally invariant (ROI) ones. Inspired by D-AMP, a denoising VAMP (D-VAMP) was proposed in [19] and analyzed in [20].

2.4 AMP for MRI

Neither AMP nor VAMP works as intended in MRI because in (1) lacks sufficient randomness. In fact, these algorithms tend to diverge in MRI if applied without modification.

The failure of AMP and VAMP can be understood from their error recursions. For AMP, the error recursion is [21]

(8)
(9)

It is important to keep in mind that images have much more energy at low Fourier frequencies than at high ones. The same tends to be true of the output error of an image denoiser. Even so, if was large and i.i.d. (with zero mean and elementwise variance ), then the term in (8) would randomize such that the denoiser input error vector looks like AWGN. In MRI, however, both and have Fourier structure, this randomization does not happen, and AMP behaves unpredictably. A similar behavior plagues VAMP.

Several MRI-specific variations of AMP and VAMP have been proposed to counter these deficiencies. For example, [22] proposed D-AMP with a very small , which helps the algorithm converge, but at the cost of degrading its fixed points. [23] proposed a damped D-VAMP that, combined with a novel initialization, showed improved performance and runtime over PnP-ADMM for MRI.

Several other VAMP-based algorithms for MRI have been designed to recover the wavelet coefficients of the image rather than the image itself. The motivation is that, in this case, is a Fourier-Wavelet matrix, which is approximately block diagonal [24], where the blocks correspond to the wavelet subbands. With an appropriate modification of VAMP, the subband error vectors can be made to behave more like AWGN, albeit with different variances. The first incarnation of this idea appeared in [25], where a fixed bandwise normalization procedure was used. Later, for single-coil MRI with variable-density sampling masks, a “variable density AMP” (VDAMP) algorithm with band-specific adaptive wavelet thresholding was proposed in [26], which was able to successfully predict the noise variance in each subband at each iteration. More recently, the D-VDAMP algorithm [27] extended VDAMP to DNN denoising in each subband.

Although D-VDAMP is the state-of-the-art AMP algorithm for MRI, it is based on a non-standard modification of VAMP with degraded fixed points, which makes early stopping critical for good performance. Also, it is not clear how to extend D-VDAMP to multi-coil MRI. These issues motivate our approach, which is described next.

3 Proposed Approach

3.1 Denoising GEC

Our approach uses the GEC framework from [6], which is summarized in Alg. 1. When solving a convex optimization problem of the form (2), the functions in Alg. 1 take the form

(10)

for the generalized proximal operator

(11)

where and creates a diagonal matrix from its vector argument. Note that if , then . Furthermore, if the vectors were held fixed over the iterations and took the form , then Alg. 1 reduces to a variant of ADMM (4) with two dual updates: (4c) and a similar step between (4a) and (4b). So, GEC can be interpreted as an ADMM-like algorithm with two adaptive vector-valued stepsizes, and .

1:.
2:Select initial
3:repeat
4:     // Measurement fidelity
5:     
6:     
7:     
8:     
9:     // Denoising
10:     
11:     
12:     
13:     
14:until Terminated
Algorithm 1 Generalized EC (GEC)

In lines 6 and 11, GEC averages the diagonal of the Jacobian separately over coefficient subsets using :

(12)

In (12), is the size of the th subset and is the th diagonal subblock of the matrix input . When , GEC reduces to VAMP.

We focus on the quadratic loss (3), which yields

(13)

For , we propose to “plug in” a DNN denoiser. For both and , we approximate the term in (12) using

(14)

where the th coefficient subset in is i.i.d. unit-variance Gaussian and the others are zero. Inspired by D-AMP and D-VAMP, we call this approach “denoising GEC” (D-GEC).

3.2 D-GEC for Image Recovery

Like [25, 26, 27], we recover the wavelet coefficients rather than the image pixels . For orthogonal wavelet transform , we have and , so that we can rewrite (1) as

(15)

To apply D-GEC to -recovery, we choose as in (13), but with in place of , and for the diagonalization subsets we choose the subbands of a depth- 2D wavelet transform. As in [27], we perform denoising in the wavelet domain using a denoiser that can exploit knowledge of the noise variance in each wavelet subband, as provided by the precision vector .

4 Numerical Experiments

We now compare the proposed D-GEC algorithm to the existing D-VDAMP [27] and PnP-PDS [13] algorithms. Based on the extensive experiments in [27], D-VDAMP is state-of-the-art among PnP algorithms. PnP-PDS is a useful baseline, since it has the same fixed points as PnP-ADMM and PnP-PG.

Denoisers: For a fair comparison to D-VDAMP [27], we use the DNN denoiser proposed in [27], which is modification of DnCNN [12]

that accepts the noise standard deviation (SD) in each wavelet subband. The denoiser was trained using noise that was white in each subband but with SD that varies across subbands. In particular, 5 copies of the denoiser were trained using 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.) The DNNs were trained using patches from 70 MRI images of the Stanford 2D FSE dataset available at

http://mridata.org

. For PnP-PDS, we used a standard DnCNN denoiser trained on the same data with white noise of SD uniformly distributed in 20-50, as in

[27]. Because we used real-valued images, the denoisers use only the real part of the input and generate a real-valued output.

Test data: For evaluation, we used the ten 352352 MRI images in Fig. 2, which were not in the training dataset. The measurements were constructed using (1) with complex AWGN whose variance was adjusted to give a pre-masking SNR of 40 dB. For the multicoil experiments, we used coil sensitivities simulated using the Biot-Savart law, while in the single-coil case, we used .

Figure 2: Test images from http://mridata.org.
Figure 3: SD of D-GEC’s denoiser input versus iteration.
Figure 4: Example single-coil recoveries and error maps.

                 

Figure 5: True coefficients, denoiser input, error (iteration 10).

Horizontal, Scale 3    Vertical, Scale 2   Diagonal, Scale 1

Figure 6: Wavelet coefficient QQ-plots showing Gaussianity.

Algorithm parameters: For D-GEC and D-VDAMP, we used a 2D Haar wavelet transform with levels, giving wavelet subbands. D-GEC used the auto-tuning scheme from [28] and the damping scheme from [23] with parameter . D-VDAMP code was obtained from the authors and run under default settings, which are detailed in [27]. PnP-PDS was run for 200 iterations using the stepsize that maximized PSNR on the training set.

coil coils
method PSNR SSIM PSNR SSIM PSNR SSIM PSNR SSIM
PnP-PDS 40.66 0.968 37.38 0.951 34.71 0.935 33.09 0.917
D-VDAMP 42.36 0.972 35.92 0.918 n/a n/a n/a n/a
D-GEC 42.97 0.977 37.65 0.946 45.18 0.993 41.13 0.982
Table 1: Recovery results averaged over the 10 test images.

Single-coil results: Table 1 shows that D-GEC outperformed D-VDAMP in all single-coil experiments and outperformed PnP-PDS in all but SSIM at . Figure 5 shows an example of the wavelet coefficients input to D-GEC’s denoiser at the 10th iteration, and their error relative to the true coefficients. Figure 3 shows the evolution of the standard deviation at the input to D-GEC’s denoiser in each subband; there is a good agreement between true and predicted values. Figure 6 suggests that the subband errors are Gaussian. Figure 4 shows image recoveries and error maps for one test image at .

Multi-coil results: Table 1 shows D-GEC significantly outperforming PnP-PDS in PSNR and SSIM in the 4-coil case. D-VDAMP does not support multi-coil recovery and thus is not shown.

5 Conclusion

We designed a GEC-based plug-and-play algorithm for MRI that renders the subband errors white and Gaussian with predictable variance, and used it with a denoiser trained to handle subband errors that are white and Gaussian with known variance. Experiments show good performance relative to previous approaches in single- and multi-coil settings.

References

  • [1] M. Lustig, D. L. Donoho, J. M. Santos, and J. M. Pauly, ‘‘Compressed sensing MRI,’’ IEEE Signal Process. Mag., vol. 25, no. 2, pp. 72--82, Mar. 2008.
  • [2] J. A. Fessler, ‘‘Optimization methods for magnetic resonance image reconstruction,’’ IEEE Signal Process. Mag., vol. 37, no. 1, pp. 33--40, 2020.
  • [3]

    K. H. Jin, M. T. McCann, E. Froustey, and M. Unser, ‘‘Deep convolutional neural network for inverse problems in imaging,’’

    IEEE Trans. Image Process., vol. 26, no. 9, pp. 4509--4522, 2017.
  • [4] K. Hammernik, T. Klatzer, E. Kobler, M. P. Rec ht, D. K. Sodickson, T. Pock, and F. Knoll, ‘‘Learning a variational network for reconstruction of accelerated MRI data,’’ Magnetic Resonance Med., vol. 79, no. 6, pp. 3055--3071, 2018.
  • [5] R. Ahmad, C. A. Bouman, G. T. Buzzard, S. Chan, S. Liu, E. T. Reehorst, and P. Schniter, ‘‘Plug and play methods for magnetic resonance imaging,’’ IEEE Signal Process. Mag., vol. 37, no. 1, pp. 105--116, 2020.
  • [6] A. K. Fletcher, M. Sahraee-Ardakan, S. Rangan, and P. Schniter, ‘‘Expectation consistent approximate inference: Generalizations and convergence,’’ in Proc. IEEE Int. Symp. Inform. Thy., 2016, pp. 190--194.
  • [7] D. L. Donoho, A. Maleki, and A. Montanari, ‘‘Message passing algorithms for compressed sensing,’’ Proc. Nat. Acad. Sci., vol. 106, no. 45, pp. 18 914--18 919, Nov. 2009.
  • [8] S. Rangan, P. Schniter, and A. K. Fletcher, ‘‘Vector approximate message passing,’’ IEEE Trans. Inform. Theory, pp. 6664--6684, 2019.
  • [9] S. Boyd, N. Parikh, E. Chu, B. Peleato, and J. Eckstein, ‘‘Distributed optimization and statistical learning via the alternating direction method of multipliers,’’ Found. Trends Mach. Learn., vol. 3, no. 1, pp. 1--122, 2011.
  • [10] S. V. Venkatakrishnan, C. A. Bouman, and B. Wohlberg, ‘‘Plug-and-play priors for model based reconstruction,’’ in Proc. IEEE Global Conf. Signal Info. Process., 2013, pp. 945--948.
  • [11] K. Dabov, A. Foi, V. Katkovnik, and K. Egiazarian, ‘‘Image denoising by sparse 3-D transform-domain collaborative filtering,’’ IEEE Trans. Image Process., vol. 16, no. 8, pp. 2080--2095, 2007.
  • [12] K. Zhang, W. Zuo, Y. Chen, D. Meng, and L. Zhang, ‘‘Beyond a Gaussian denoiser: Residual learning of deep CNN for image denoising,’’ IEEE Trans. Image Process., vol. 26, no. 7, pp. 3142--3155, 2017.
  • [13] S. Ono, ‘‘Primal-dual plug-and-play image restoration,’’ IEEE Signal Process. Lett., vol. 24, no. 8, pp. 1108--1112, 2017.
  • [14] U. Kamilov, H. Mansour, and B. Wohlberg, ‘‘A plug-and-play priors approach for solving nonlinear imaging inverse problems,’’ IEEE Signal Process. Lett., vol. 24, no. 12, pp. 1872--1876, May 2017.
  • [15] R. Berthier, A. Montanari, and P.-M. Nguyen, ‘‘State evolution for approximate message passing with non-separable functions,’’ Inform. Inference, 2019.
  • [16] M. Bayati and A. Montanari, ‘‘The dynamics of message passing on dense graphs, with applications to compressed sensing,’’ IEEE Trans. Inform. Theory, vol. 57, no. 2, pp. 764--785, Feb. 2011.
  • [17] C. A. Metzler, A. Maleki, and R. G. Baraniuk, ‘‘BM3D-AMP: A new image recovery algorithm based on BM3D denoising,’’ in Proc. IEEE Int. Conf. Image Process., 2015, pp. 3116--3120.
  • [18] S. Ramani, T. Blu, and M. Unser, ‘‘Monte-Carlo SURE: A black-box optimization of regularization parameters for general denoising algorithms,’’ IEEE Trans. Image Process., vol. 17, no. 9, pp. 1540--1554, 2008.
  • [19] P. Schniter, S. Rangan, and A. K. Fletcher, ‘‘Denoising-based vector approximate message passing,’’ in Proc. Intl. Biomed. Astronom. Signal Process. (BASP) Frontiers Workshop, 2017, p. 77.
  • [20] A. K. Fletcher, P. Pandit, S. Rangan, S. Sarkar, and P. Schniter, ‘‘Plug-in estimation in high-dimensional linear inverse problems: A rigorous analysis,’’ in Proc. Neural Inform. Process. Syst. Conf., 2018, pp. 7440--7449.
  • [21] P. Schniter, ‘‘A simple derivation of AMP and its state evolution via first-order cancellation,’’ IEEE Trans. Signal Process., vol. 68, pp. 4283--4292, 2020.
  • [22] E. M. Eksioglu and A. K. Tanc, ‘‘Denoising AMP for MRI reconstruction: BM3D-AMP-MRI,’’ SIAM J. Imag. Sci., vol. 11, no. 3, pp. 2090--2109, 2018.
  • [23] S. Sarkar, R. Ahmad, and P. Schniter, ‘‘MRI image recovery using damped denoising vector AMP,’’ in Proc. IEEE Int. Conf. Acoust. Speech & Signal Process., 2021, pp. 8108--8112.
  • [24] B. Adcock, A. C. Hansen, C. Poon, and B. Roman, ‘‘Breaking the coherence barrier: A new theory for compressed sensing,’’ Forum of Mathematics, Sigma, vol. 5, no. E4, doi:10.1017/fms.2016.32.
  • [25] P. Schniter, S. Rangan, and A. K. Fletcher, ‘‘Plug-and-play image recovery using vector AMP,’’ presented at the Intl. Biomedical and Astronomical Signal Processing (BASP) Frontiers Workshop, Villars-sur-Ollon, Switzerland, Jan. 2017. [Online]. Available: http://www2.ece.ohio-state.edu/~schniter/pdf/basp17_poster.pdf
  • [26] C. Millard, A. T. Hess, B. Mailhé, and J. Tanner, ‘‘Approximate message passing with a colored aliasing model for variable density Fourier sampled images,’’ arXiv:2003.02701, 2020.
  • [27] C. A. Metzler and G. Wetzstein, ‘‘D-VDAMP: Denoising-based approximate message passing for compressive MRI,’’ in Proc. IEEE Int. Conf. Acoust. Speech & Signal Process., 2021, pp. 1410--1414.
  • [28] A. K. Fletcher, M. Sahraee-Ardakan, S. Rangan, and P. Schniter, ‘‘Rigorous dynamics and consistent estimation in arbitrarily conditioned linear systems,’’ in Proc. Neural Inform. Process. Syst. Conf., 2017, pp. 2542--2551.