Model-Guided Multi-Contrast Deep Unfolding Network for MRI Super-resolution Reconstruction

Let $X\in\mathbb{R}^{n\times C}$ represents the degraded observations and $Z\in\mathbb{R}^{N\times C}$ represents the unknown original image, where $C$ denotes the number of channels, and $n=h\times w$ and $N=H\times W$. It is assumed that an LR image is obtained by down-sampling and blurring an HR image, so the linear relationship between the observed image and the original HR image can be typically formulated as follows:

where $D\in\mathbb{R}^{n\times N}$ and $K\in\mathbb{R}^{N\times N}$ represent the process of down-sampling and blurring, respectively, and ${\mathcal{N}}_{1}$ denotes the noise.

Transform modal relationship considering multi-contrast MR images in MCSR task, the transform relationship between the guide image $Y\in\mathbb{R}^{N\times C}$ and the unknown original image can be formulated by:

where $P\in\mathbb{R}^{N\times N}$ is the transform function, and the ${\mathcal{N}}_{2}$ represents the noise in this process. As a result, $Z$ can be obtained by solving the following objective function:

where the hyper parameters ($\eta,\lambda_{1},\lambda_{2}$) are the trade-off coefficient, and the last two regularization terms correspond to the prior domain knowledge of the MCSR task, which are noise prior in the typical degradation process and transform modal noise prior in multi-contrast image, respectively. The choice of various regularization functions reflects different ways of incorporating prior knowledge about the unknown original HR MR image.

In the next section, we describe how to solve the objective function with an iterative algorithm. Then, we unfold the iterative algorithm into neural networks for image SR, obtaining an end-to-end reconstruction architecture.

Following the framework of half-quadratic splitting (HQS) to introduce two auxiliary splitting parameters $U$ and $V$ for $Z$ with different prior knowledge of the MCSR task, the Eq. 5 can be formulated as a non-constrained optimization problem, which can be written as:

where $\beta_{1}$, $\beta_{2}$, $\lambda_{1}$, and $\lambda_{2}$ are the penalty parameters. To obtain an unrolling inference, Eq. 6 can be divided into the following three sub-problems and solved alternatively:

here, $t$ denotes the HQS iteration index.

For the objective function of Eq. 5 with the prior knowledge of the MCSR task, we employ the efficient Proximal Gradient Descent (PGD) to solve the above three sub-problems:

where $Prox_{{\mathfrak{R}}_{1}}(\cdot)$ and $Prox_{{\mathfrak{R}}_{2}}(\cdot)$ are proximal operators corresponding to penalty ${{\mathfrak{R}}_{1}}(\cdot)$ and ${{\mathfrak{R}}_{2}}(\cdot)$, which integrate the information coming from target contrast LR MRI and guidance contrast HR MRI. And the gradient related notations are detailed as:

In summary, the iterative algorithm for solving the MCSR task of Eq. 5 is given above, where we initialize $Z^{(0)}$ with a bicubic interpolated version of $X$. Under the framework of the interpretable MGDUN model, the PGD algorithm usually requires dozens of iterations to converge.

The main idea behind deep unfolding network is that conventional iterative soft-thresholding algorithm (ISTA) can be implemented equivalently by a stack of recurrent neural networks (Wisdom et al., 2017). Inspired by the principle of model-driven deep learning, we generalize Eq. 3.2.2 to Eq. 12 as a network block. Each step is translated with deep learning terminologies.

In each network, two auxiliary variables ($U$ and $V$) are updated first. Due to the existence of noise, the $Prox(\cdot)$ operator can be implemented by a deep denoising module ($\boldsymbol{DM}$). In Eq. 3.2.2, given the evaluated HR image $Z^{(t)}$ of the current stage and auxiliary splitting parameter $U^{(t-1)}$ of the previous stage, it generates the auxiliary splitting parameter $U^{(t)}$ of the current stage. The same as $V^{(t)}$ in the form of Eq. 3.2.2. In neural networks, these two steps are implemented by:

where $\boldsymbol{DM}(\cdot;C,C^{\prime},C)$ is used to obtain more expressive auxiliary splitting parameter. $C^{\prime}$ is the number of channels for feature maps.

Then, we reconstruct our estimated HR image according to Eq. 12 and Eq. 13. In our network, the reconstruction process which is implemented by a reconstruction module consists of cross-modal transform module, down-sampling block, and up-sampling block, takes auxiliary various $U^{(t)}$, $V^{(t)}$, the evaluated HR image $Z^{(t)}$, the LR MRI and the guide image $Y$ as inputs and outputs the reconstructed image $Z^{(t+1)}$. Therefore, the reconstruction process in the neural network is as follows:

where $\boldsymbol{Up}(\cdot)$ and $\boldsymbol{Down}(\cdot)$ denote the up-sampling and down-sampling operator in spatial resolution, respectively, and $\boldsymbol{P}(\cdot)$ and $\boldsymbol{P^{T}}(\cdot)$ perform the cross-modal transform functions.

Then, the updated $Z^{(t+1)}$ is fed into the next stage to refine the estimate $U$ and $V$ again. The denoising module and the reconstruction module are alternatively updated $T$ times until reaching the final reconstruction.

The overall network architecture of the MGDUN is shown in Fig. 1, which contains $T$ stages that are intentionally designed to correspond to $T$ iterations in the PGD optimization algorithm. Each stage of MGDUN consists of three specified network modules, containing a deep denoising module for denoising and updating for auxiliary variables, a cross-modal transform Module for cross-modal transform function, and a reconstruction module for reconstruction and updating for $Z$. We will elaborate on each module next.

The design of deep denoising modules corresponds to the computing the updated estimate $U^{(t)}$ and $V^{(t)}$ in Eq. 3.2.2 and Eq. 3.2.2, respectively. In general, any existing image denoising network can be used as the denoising module here.

As shown in Fig. 1, the intermediate estimates $U^{(t-1)}$ and $V^{(t-1)}$ are fed into the proximal operator after weighting with the intermediate estimate $Z^{(t)}$ for further refinement. In this paper, we have adopted a variant of U-net as the backbone of the deep denoising module. Other more effective networks for medical image denoising can be also adopted. The U-net denoising network consists of an encoder and a decoder. As shown in Fig. 2(a), the encoder consists of four encoding blocks, which contain two convolutional layers with $3\times 3$ kernels and ReLU nonlinearity. Corresponding to the encoder, the decoder also consists of four decoding blocks, which contain two convolutional layers with $3\times 3$ kernels and ReLU nonlinearity. Instead of predicting the refine auxiliary various directly, we enable the denoising module to predict the residual by adding a skip connection from the input to the output. To reduce the number of network parameters and the effect of overfitting, we opt to enforce all denoising modules sharing the same network parameters.

Noted that Eq. 16 involves the cross-modal transform matrix $\boldsymbol{P}$ and $\boldsymbol{P^{T}}$ that are expensive to calculate. We find that $\boldsymbol{P}$ and $\boldsymbol{P^{T}}$ perform a cross-modal transform operation, and the two processes are inverted: one from the target contrast image to the guide contrast image, and the other from the guide contrast image to the target contrast image. The process can be formulated as:

As a result, we design cross-modal transform modules using the principle of invertible neural networks (INNs). INNs have been adopted in various inference tasks and achieved excellent performance due to their flexibility (Kingma et al., 2016; Lu et al., 2021). We formulate an INN architecture design to serve as cross-modal transform operation. It consists of two pixel shuffling layers (Dinh et al., 2016) and several INN blocks. As shown in Figure 2(b), relevant invertible modules are embedded in the cross-modal transform module.

For the t-th stage, given an evaluation HR MRI ($Z^{(t)}$) to be refined, we first put it to one pixel shuffling layer for changing dimension, then pass through several INN blocks to execute the cross-modal transform function, and finally restore the original dimension through the other pixel shuffling layer.

For the forward operation, one pixel shuffling layer executes dimension addition first. Then the input ($Z^{(t)}$) is divided into ($Z_{1}^{(t)}$) and ($Z_{2}^{(t)}$) along the channel axis, and the corresponding cross-modal transform output is $\mathfrak{T}_{1}^{(t)}$ and $\mathfrak{T}_{2}^{(t)}$ (two components of $\mathfrak{T}^{(t)}$). This process corresponds to the operation of cross-modal transform matrix $\boldsymbol{P}$, in which an INN block can be expressed as :

where $\varphi(\cdot)$, $\eta(\cdot)$ and $\rho(\cdot)$ are arbitrary functions, $exp(\cdot)$ is Exponential functions, and $\odot$ is the Hadamard product.

Accordingly, for the backward operation, given [$\mathfrak{T}_{1}^{(t)}$, $\mathfrak{T}_{2}^{(t)}$], it is easy to calculate [$Z_{1}^{(t)}$, $Z_{2}^{(t)}$] as:

This process corresponds to the operation of matrix $\boldsymbol{P^{T}}$ in Eq. 18. The forward and backward operations are shown in Fig. 2.

The design of the reconstruction module corresponds to the update of intermediate evaluated $Z^{(T)}$ as described in Eq. 16. With the output of denoising modules ($U^{(t)}$, $V^{(t)}$), the evaluated HR image $Z^{(t)}$, the LR MRI and the guide image $Y$, we can reconstruct the updated image $Z^{(t+1)}$. The architecture of the reconstruction module is shown in Fig. 1 and Eq. 16, which still involve the degradation operations (${(DK)^{T}}$ and ${DK}$). The pair of operation ${(DK)^{T}}$ and ${DK}$ can be implemented by up-sampling and down-sampling layers for modeling capability.

The operators ${(DK)^{T}}$ and ${DK}$ are simulated using a convolution network layer respectively. Specifically, ${DK}$ is simulated by a network called down-sampling-blocks ($Down$) consisting of a convolutional layer with $3\times 3$ kernels and 64 channels, one max pool layer to decrease the spatial resolution, and two convolutional layers with $3\times 3$ kernels for reprojection to the original dimension (as shown in Figure 2(c)). Similarly, the ${(DK)^{T}}$ is simulated by a network call Up-sampling-blocks ($Up$) consisting of a convolutional layer with $3\times 3$ kernels and 64 channels, and one upsample layer to increase the spatial resolution and two convolutional layers with $3\times 3$ kernels for reprojection to the original dimension as shown in Figure 2(d).

We will apply this model for the super-resolution reconstruction of LR T2WI with the aid of HR PDWI (or HR T1WI).

Our MGDUN is supervised by the $\mathcal{L}_{1}$ loss between $Z^{(T)}$ and the groundtruth $Z$. Then, the overall network is trained by minimizing the following loss function:

where $X_{i}$, $Y_{i}$, and $Z_{i}$ denote the $i^{th}$ pair of target contrast LR T2WI, guide HR PDWI, and the original target contrast HR T2WI, respectively. $g(\cdot;\Theta)$ denotes the reconstructed HR T2WI by the network with parameter $\Theta$.

For quantitative evaluation, we utilize the average PSNR, SSIM, and MSE. Tab. 1 reports the target contrast reconstruction performance on different datasets under $2\times$ and $4\times$ enlargements where the best and second best values are highlighted in bold red and underline blue, respectively. It is clearly noted that our method achieves the best performance on different enlargements of different datasets. The results demonstrate that our model can effectively fuse the two contrasts, which is beneficial to the SR reconstruction of the target contrast. This substantiates the effectiveness and flexibility of our method with a certain degree of generalization.

We provide qualitative comparison results on the IXI dataset as well as the BraTs dataset and their corresponding error maps in Fig. 3 and Fig. 4. The texture of error maps represents the restoration error, the smoother the texture, the better the reconstruction. As we can see, the input has significant aliasing artifacts and lacks anatomical details. It can be noted that our model recovers the image with fewer visible artifacts and reconstructs more details than other competing methods. The quality improvement achieved by MDCUN may be associated with the full usage of the feature maps from the former stages to refine the final results.