Joint Edge Optimization Deep Unfolding Network
for Accelerated MRI Reconstruction

In accelerated parallel imaging, the algorithm aims to reconstruct high-quality MR images from partially sampled multi-coil k-space data. The methods of MRI reconstruction in parallel imaging scenarios can be categorized into traditional model-based and data-driven methods. In traditional methods, GRAPPA [16] uses interpolation in k-space; SENSE [17] exploits sensitivity estimation maps in the spatial domain to eliminate aliasing artifacts. In addition to this, some of the methods solve the optimization model (1) based on the theory of compressed sensing in parallel imaging scenarios, such as ESPIRiT [18], which combines a multi-coil degradation mechanism and a constrained regularity term for the reconstruction. However, the reconstructed image quality of traditional methods is not stable and the imaging speed is slow.

With the development of deep learning techniques, artificial neural networks have achieved excellence in many tasks in the field of computer vision [19, 20, 21]. Meanwhile, the data-driven methods appear and the researchers explore the potential of neural networks in the field of medical imaging. In the accelerated MRI reconstruction task, Wang et al. [22] and Schlemper et al. [23] proposed for the early time to use convolutional neural networks to learn the mapping relationship between zero-filled undersampled images to high-quality reconstructed images. Hyun et al.[24] utilized U-Net [25] which is an efficient encoder-decoder architecture in segmentation tasks, as the reconstruction modeling framework. Similar works have been done by Sun et al. [26] utilizing recursive dilated networks. Some methods [27][28] restore zero-filled degraded images based on the generative adversarial network architecture as a means of removing aliasing artifacts and redundant noise. In addition, since the original MRI data is in complex form, researchers [29] applied complex convolution operation to jointly process the real and imaginary parts of the MR image to preserve the magnitude and phase information of the image. Moreover, Considering the importance of edge information restoration, EAMRI [15] utilizes the edge detector to extract the image edge, which is used as guidance for constructing the attention map. However, the ”off-the-shelf” edge detectors have their own limitations which will migrate to the reconstruction algorithm. Despite their reconstruction result is generally better than traditional methods with lower inference time, the data-driven reconstruction methods described above are not conducive to the clinical practice due to the lack of interpretability resulting from the black-box nature of neural networks.

The innovative deep unfolding methods are initially introduced to sparse coding approximations [30], fusing optimization-driven iterative algorithms with the expressive power of neural networks. This methodology has seen successful application across a variety of low-level image processing tasks [31].

For instance, the Denoising Convolutional Neural Network (DPDNN) [32] and the Model-based Deep Learning architecture (MoDL) [33] utilize deep unfolding schema by employing neural networks to embody model-based regularization terms within image domain. E2EVarNet [34] integrate a gradient descent optimization for k-space data enhancements in MRI reconstruction, unfolding the first-order derivatives of regularization function with the neural network. HQS-Net [35] uses a half quadratic splitting method to process the reconstruction model and integrates the neural network module into the optimization process to design an efficient and lightweight reconstruction method; HUMUS-Net [36] integrates the multi-scale network structure and the attention mechanism into the depth-expanded reconstruction method to achieve high-quality reconstruction results, but due to the large size of the parameters of the network model, it is not that efficient in training. Meanwhile, methods like ISTA-Net [37] and ADMM-CSNET [38] take a leap forward by unfolding iterative soft-thresholding and alternating direction multiplier algorithms, respectively, into learning-based models that combine optimized algorithm with data-driven component. VS-Net [39] unfolding the iterative refinement optimization process with neural network to constructing a lightweight architecture, while MGDUN [40] uses neural network structures to semantically replace traditional sampling and modality conversion processes, particularly highlighting its prowess in tasks such as cross-modality super-resolution.

A noteworthy aspect of deep unfolding schema is its inherent reduction of hand-craft design, leaning towards an autonomous optimization. Moreover, the deployment of deep unfolding has not only augmented the clarity and interpretability of computational models but also enhanced the precision of the eventual solutions.

The signal intensity of MR images typically demonstrates a piecewise smooth distribution, which is attributed to the correlation with the underlying tissue structure. With this piecewise smooth property, a redundant system can provide a sparse representation of the MR image, generating coefficients with large amplitude along the locations of image edges and coefficients with near-zero (or zero) amplitude within the smooth regions (non-edge regions). Based on this, the edges and the image can be effectively associated with each other. In this paper, we introduce a probability map named $P_{ne}$ to quantify the probability of a region being smooth (non-edge) and utilize it to constrain the values of transformation coefficient inside the smooth regions to be close to zero. In addition, as MR images depict the structural information of anatomical tissues, they have certain geometrical structures. Consequently, the smooth regions within these MR images are not randomly distributed but conform to specific patterns and characteristics. Therefore, the probability map $P_{ne}$ that measures the likelihood of a region being smooth also should fulfill certain properties and can be imposed with regularizations to satisfy the corresponding properties.

Based on the aforementioned discussion, our proposed joint edge optimization MRI reconstruction model can be formulated as follows:

where $*$ denotes a pixel-wise product¹¹1$*$ is omitted in the following text for brevity., ${P_{ne}}$ is the introduced non-edge probability map related to the edge information. The value of ${P_{ne}}$ ranges from $0$ to $1$. When its value is $0$, it indicates that the likelihood of the pixel belonging to a non-edge is $0$, and we do not constrain the corresponding transformation coefficient. On the other hand, when the value is $1$, it suggests that there is a high probability that this pixel belongs to a smooth (non-edge) region, and we expect the transformation coefficient to be close to $0$. Conversely, when the images are updated, the value of ${P_{ne}}$ will be decreased or increased accordingly based on whether the corresponding pixel belongs to edge or not. $W$ denotes the stationary Haar wavelet transformation. ${\lambda_{1}}$, ${\lambda_{2}}$, and $\rho$ are the trade-off parameters of different regularization terms.

The optimization model (2) consists of four terms. The first is the fidelity term, which ensures the data consistency between the reconstructed MR image $x$ and the measured k-space data. The third term is the co-regularizer between image $x$ and edge-related variables ${P_{ne}}$, linked by a transformation $W$, aiming to promote the piece-wise smoothness property of image $x$. The second term $\Phi(x)$ and the forth term $\Psi({P_{ne}})$ are the regularizers of the image $x$ and the probability map ${P_{ne}}$ respectively. Instead of imposing hand-crafted priors, we learn the inherent priors of images and edges automatically through the network.

In Model (2), two variables are included, i.e., the reconstructed image $x$ and the non-edge probability map $P_{ne}$. To solve this optimal model, we decouple the two variables $\{x,P_{ne}\}$ and optimize them alternatively by solving the following subproblem (3) and (4):

where $k$ denotes the stage index of the iteration. As seen in the above subproblems, the co-regularizer between the reconstructed image and the non-edge probability map exists in both two subproblems. Hence in the alternating optimization process of two variables, they will guide each other for better recovery. For the solution of the two subproblems, we will introduce them in detail in the following.

The principle of deep unfolding is that the modules of optimize-based approach which need to be manually designed are difficult to determine, so the deep unfolding schema utilize the neural networks to replaced the undetermined module or operator, introducing the representing capability of deep neural networks into the optimizing process. In traditional model-based methods, the design of the regularization terms $\Phi$ and $\Psi$ are crucial for the final reconstruction. In this paper, to avoid the reliance on hand-crafted prior and the interference of inaccurate assumptions, the neural networks are utilized to automatically learn the inherent priors. Specifically, since the regularization terms $\Phi$ and $\Psi$ only occur in the proximal operator in Eq. (7) and Eq. (12), we consider the substitution of neural networks for the proximal operator. Thus for Eq. (7) we utilize an edge recovery network (ERN) to replace the proximal gradient operator $prox_{\Psi}$. And for Eq. (12), the proximal gradient operator $prox_{\Phi}$ is substituted using image denoising network (IDN). Therefore, the optimization for the auxiliary variables $V$ and $Z$ can be rewritten in the following form:

where ${\theta_{e}^{(k)}}$ and ${\theta_{i}^{(k)}}$ are the parameters of the edge recovery network ERN and the image denoising network IDN at the ${k}^{th}$ iteration, respectively. As to the specific network structure design of ERN and IDN, we utilize the standard U-Net architecture [25] which is suitable for the image-to-image conversion task due to its encoder-decoder structure that extracts features from images at different scales through multi-layer up-sampling and down-sampling operations. Moreover, the standard U-Net architecture is combined with the residual connection from the module input to output in our method so as to avoid gradient vanishing issue [42].

In the following, we present an overall reconstruction framework for our joint edge optimization model. The optimization process consists of $K$ iterative stages, each containing four optimization steps. First, the edge auxiliary variable $V$ is optimized using the edge recovery network (ERN). Then, the non-edge probability map ${P_{ne}}$ is optimized based on the closed-form solution of Eq. (9), known as edge optimization (EO). The reconstructed image auxiliary variable $Z$ is optimized using the image denoising network (IDN), and finally, the reconstructed image $x$ is optimized using single-step gradient descent by Eq. (14), known as image optimization (IO). The detailed optimizing steps at each stage are described in Algorithm 1. As for the initialization of the image ${x^{(0)}}$, we process the partially sampled k-space data of the multi-coils using the zero-filled method to obtain the image domain data and later fuse the multi-coils data into one channel using sensitivity estimation maps, i.e., ${x^{(0)}}=\sum\limits_{i=1}^{{n}}{S_{i}^{H}{F^{H}}{y_{i}}}$; As for the initialization of the non-edge probability map $P_{ne}^{(0)}$, first the stationary haar wavelet transform operator $W$ transforms the initial MR image to the wavelet domain, then the coefficients of the high-frequency channel are normalized to [0,1] with a maximum-minimum normalization operator $N$. Finally, the initial $P_{ne}^{0}$ is obtained by inverting the normalized coefficients as follows:

where $|W{x^{(0)}}|$ is the coefficients of the high-frequency channel by pixel-wise absolute operation. The iterative reconstruction framework is shown in Figure 2, where $P_{ne}^{(0)}$ and ${x^{(0)}}$ are used as inputs to the $K$-stage deep unfolding reconstruction network to obtain the optimized edge information $P_{ne}^{(K)}$ and the final reconstructed MR image ${x^{(K)}}$. In Figure 3, we visualize the non-edge probability map while showing the evolution procedure at different iteration stages.

Subsequent to the construction of the deep unfolding iterative framework, the training of our reconstruction model follows an end-to-end supervised learning paradigm. Similar with the process of initialization, fully sampled k-space data are exploited to create ground truth image $\hat{x}$ and the non-edge probability map $\hat{P}_{ne}$. The training strategy utilizes a multi-component loss function, composed of the mean square error (MSE) for fidelity in image reconstruction and the L1 norm for precision in edge information recovery. Consequently, the total loss function $L$ for the joint edge optimization within the reconstruction model is formulated as follows:

Here, $N_{s}$ symbolizes the number of samples in the training dataset sample scale while ${\gamma_{1}}$ and ${\gamma_{2}}$ are the respective trade-off parameters that modulate the loss components’ contribution to $L$. In addition, adjusting all of the undetermined parameters in the iterative process manually is a cumbersome task accompanied by a degree of uncertainty. Therefore, to elevate the model training efficiency, all of the undetermined parameters along with the parameters from the edge recovery network (ERN) and the image denoising network (IDN), denoted by $\{{\theta_{e}^{(k)}}\}_{k=1}^{K}$ and $\{{\theta_{i}^{(k)}}\}_{k=1}^{K}$, are collected as trainable entities. This framework ensures the self-adjusting capability of the hyper-parameters throughout the back-propagation process, fostering their inherent adaptability within the reconstruction architecture. With regards to multi-stage iterative model configurations, our investigative results utilizes a non-shared parameter strategy. Such an approach amplifies the quality of the reconstructed image in the context of our joint edge information optimization framework. Accordingly, the learnable parameter set for the proposed reconstruction model is expressed as $\Theta=\{{\theta_{e}^{(k)}},{\theta_{i}^{(k)}},{\rho^{(k)}},{\alpha^{(k)}},{\beta^{(k)}},{s^{(k)}}\}_{k=1}^{K}$, encompassing all the iterations and facilitating a tailored optimization at each discrete stage of the unfolding process.

For our experiments, we employed two NVIDIA RTX 3090 Ti GPUs for our hardware and Pytorch is used as the deep learning framework for building the training and inferring pipeline. The training process utilized the ADAM optimizer [43], which is known for its adaptive learning capabilities. An initial learning rate of 0.01 was chosen to ensure a strong starting point for optimization, and we applied a cosine decay learning rate schedule to reduce the learning rate smoothly over epochs for enhancing the model’s convergence ability. Training was conducted with a small batch size of 2. We repeatedly passed the dataset through the model for 180 epochs to ensure thorough learning and convergence. Our model architecture chosen K = 7 stages in total, based on empirical evidence suggesting this to be optimal for the performance of our method. Each stage’s parameters operated under a non-shared strategy, which provided the flexibility needed to adapt to each specific set of the optimized stages. The ERN and IDN networks were built upon the original U-Net architecture, capitalizing on its symmetrical design with max-pooling layers in the encoder for undersampling and an equal number of transconvolution layers in the decoder for upsampling, ensuring detailed feature extraction and restoration. We opted for Kaiming initialization [44] for setting up the neural network parameters, a method designed to maintain adaptive variance as a countermeasure against gradients diminishing too rapidly during training. Regarding the loss function, we integrated trade-off parameters ${\gamma_{1}}$ and ${\gamma_{2}}$, weighted as 1 and 0.1 to balance the total loss function.

For a comprehensive assessment, we engage in comparative experiments using both multi-coil and single-coil MRI data. Our ablation study and parameter strategy analysis are conducted on the multi-coil dataset for the parallel imaging scenario.

The multi-coil dataset employed is sourced from the dataset made public in conjunction with the MoDL study [33]. This dataset is derived from a 3D T2 CUBE sequence captured using a 12-coil parallel imaging setup. It encompasses a total of 360 slices allocated for training, complemented by a validation set consisting of 164 slices. Each image slice adheres to a resolution of 256x232. The sensitivity maps, which are indispensable for various parallel imaging reconstruction techniques, were calculated using the ESPIRIT algorithm [18] that utilize the center region of k-space to reconstruct the MR image, aiming at considering the most informative part of the MRI signal for image reconstruction.

In addition to multi-coil data, our methodology was also validated on the fastMRI single-coil knee dataset [46] including varied contrasts such as proton-density weighting with fat-saturated (PDFS) and proton-density (PD) imaging modalities. We randomly sampled 20% of both training and testing samples from the aforementioned repository. The single-coil dataset setting in our experiment included 4034 slices for training purposes and 745 slices designated for testing. Adhering to common evaluation protocols, all slices were uniformly cropped centrally to a standardized dimension of 320x320 pixels.

To comprehensively assess the effectiveness of our proposed methodology, we conduct evaluations using both multi-coil and single-coil MRI datasets. The Peak Signal-to-Noise Ratio (PSNR) is utilized as the principal metric for quantitative analysis. In total, we have conducted comparisons with nine additional methods. Among them, U-Net[24], DCCNN[23] and RefineGAN[27] use a pure data-driven paradigm to reconstruct MR images in the image domain using neural network models. RefineGAN uses generative adversarial network and bi-domain cyclic loss function. MoDL[33], VS-Net[39], RecurrentVarNet[45], ADMM-CSNet[38] and HQS-Net[35] use deep unfolding paradigm, which utilize neural networks for unfolding the optimization process of MR images. EAMRI[15] takes the image edge information into account in the reconstruction method and they extract the edge information based on the edge detection operator to guide the generation of the attention map. All competing methods are fine-tuned to their optimal settings for a fair comparison.

In the context of multi-coil MRI data, we explore scenarios involving Cartesian and random undersampling schemes at acceleration factors of 6x and 10x, as documented in Table 1. The experiment results indicate that regardless of the undersampling templates or acceleration rates, our method outperforms the other techniques in terms of PSNR values. This is further supported by error maps presented in Figure 4, which clearly demonstrate that the proposed approach achieves the minimum reconstruction error. For single-coil MRI data, experiments are imparted at both 6x and 10x acceleration using random and equidistant Cartesian undersampling schemes. The quantitative results presented in Table 2 reveal our method’s superior performance in PSNR metric outpacing other algorithms. Supplementing our quantitative analysis is a visual examination of the single-coil reconstruction outcomes, shown in Figure 5, which showcase that our method not only produces reconstructions with precise edge details but also exhibits minimal errors when compared to the fully-sampled reference. Overall, these visual results illustrate our method’s adaptability and performance across variously accelerated MRI scenarios on both multi-coil and single-coil configurations.

In a multi-stage deep unfolding reconstruction architecture, there are typically two parameter strategies available for selection: the recursive shared-parameter form and the non-shared form that assigns independent parameters to each iteration stage. The shared-parameter form is characterized by a compact parameter set, promoting speed up inference by minimizing the computational footprint. Conversely, the non-shared form delivers enhanced flexibility, with each stage fine-tuned for optimal performance at the potential cost of increased parameter count.

Our experimental involved testing the performance impact of both strategies across different iteration stages ranging from 1 to 10 on a multi-coil dataset subjected to 6x random sampling. The outcomes of this investigation are graphically presented in Figure 7, providing quantitative evidence for the comparative analysis. The result curves illustrate a superior performance yield when deploying the non-shared parameter configuration within our proposed model. This suggests that the capability to customize each stage’s parameters independently facilitates more effective optimization. Furthermore, an analysis of the trend in reconstruction quality as a function of iterative stages reveals a positive trajectory up until the seventh iteration in the non-shared configuration. Additional iterations beyond this threshold will lead to a decrease in overall performance, a phenomenon we attribute to the increasing parameter scale that lead the model to be overfitting. These observations affirm the importance of tailoring the iterative parameter strategy to the specific requirements of the reconstruction task. Our proposed method, employing the non-shared parameter strategy, appears to reconcile the model’s complexity with the imperative for high-quality image reconstruction.