Self-Score: Self-Supervised Learning on Score-Based Models for MRI Reconstruction

Motivated by the abovementioned problems, this paper proposes a new self-supervised DL method for MRI reconstruction. Specifically, the main contributions and observations of this work are summarized as follows.

1. We propose a self-supervised learning score-based diffusion model for the scenario without a fully sampled MRI training set and derive the corresponding conditioned Langevin MCMC sampling for MRI reconstruction.

2. In the proposed model, we show that the score of fully sampled MR image probability density can be estimated accurately under certain conditions depending only on the undersampled data. It provides theoretical guarantees for the proposed self-supervised learning MRI reconstruction.

3. In terms of reconstruction accuracy and generalization ability (including sampling trajectories and data shifts), experimental results show that our proposed method outperforms traditional parallel imaging, self-supervised DL, and conventional supervised DL methods and achieves comparable performance with conventional (fully sampled data trained) score-based diffusion methods.

The remainder of the paper is organized as follows. Section 2 reviews some related work. Section 3 discusses the proposed self-supervised score-based diffusion model. The implementation details are presented in Section 4. Experiments performed on several data sets are presented in Section 5. The discussions are presented in Section 6. The last section 7 gives some concluding remarks. All the proofs are presented in the Appendix.

Suppose a certain data set $\{x_i\in C^d{\}}_{i=1}^N$ contains i.i.d. sampled data from an unknown distribution $p\left(x\right)$ and the score function $s_{\varphi }:C^d\to C^d$ with parameter $\varphi$ is an approximation of the gradient of its log probability density, i.e., $\nabla \log p\left(x\right)$. Then, Langevin MCMC sampling is performed according to the score function to obtain samples that obey the distribution $p\left(x\right)$, i.e.,

where ${\eta }_i>0$ is the stepsize, and $z_t$ is standard normal. $\nabla \log p\left(x\right)$ is difficult to calculate directly, the denoising score matching method [31] first perturbs the original data $x$ to $\tilde{x}$, and the score function on $\tilde{x}$ can be obtained by solving the following optimization problem:

where $q_{\epsilon }\left(\tilde{x}|x\right):=N(\tilde{x};x,{\epsilon }^{2}I)$. Given a sequence of positive noise scales ${\epsilon }_{min}={\epsilon }_1<\cdots <{\epsilon }_L={\epsilon }_{max}$, the data noise perturbation process can be considered as a discretization of a continuous diffusion process. In particular, several works model this diffusion process through Markov chains [11] and Itô SDEs [30]. Leveraging the learned score function, we can perform the reverse process (Langevin MCMC sampling) of above diffusion model to generate the data $\tilde{x}\sim q_{{\epsilon }_{min}}\left(\tilde{x}\right)$ from random noise $n\sim q_{{\epsilon }_{max}}\left(n\right)$, where $q_{{\epsilon }_i}(\cdot )=\int q_{{\epsilon }_i}(\cdot |x)p\left(x\right)\text{d}x$. Typically, if ${\epsilon }_{min}$ is sufficiently small, it can be assumed that $q_{{\epsilon }_{min}}(\cdot )\approx p(\cdot )$.

Focusing on MRI, data $\{x_i\in C^d{\}}_{i=1}^N$ represents the fully sampled MR images. In many practical applications, a large training set of fully sampled MRI images is difficult to acquire, which limits the application of score-based diffusion models.

Self-supervised learning has been studied in MRI for a longer period. For example, the classical $k$-space interpolation method, which first fully samples a calibrated region in the central part of the $k$-space, learns a linear kernel and uses it in a translation-invariant manner to interpolate in the missing $k$-space data [9, 21]. Inspired by deep learning, linear kernels are generalized to convolutional neural networks to improve the accuracy of missing data interpolation [1, 15].

On the one hand, inspired by self-supervised learning denoising methods in computer vision, self-supervised MRI reconstruction methods that do not require fully sampled calibration data have been proposed [34, 22, 7]. More specifically, an undersampled $k$-space data pair is constructed by drawing on the self-supervised denoising method to construct a noisy data pair, and then the interpolation relationship between the missing data and the sampled data is learned from it to achieve missing data interpolation (MRI reconstruction).

However, in terms of imaging quality, many experiments show that self-supervised learning methods tend to perform inferior to supervised learning methods. Moreover, as in $k$-space interpolation methods, the self-supervised learning of the interpolation kernel relies on the sampling trajectory, thus leading to a lack of generalization ability. Therefore, a novel self-supervised learning method with high reconstruction quality and generalization ability has been desired for MRI reconstruction.

The forward model of MRI can be expressed as

where $A$ is the encoding matrix, $y$ is the measurement, $x$ is the image to be reconstructed, and $n$ is the measurement Gaussian noise with scare $\gamma$ i.e., $n\sim N(n;0,{\gamma }^{2}I)$. In particular, $A=PFS$ for the case of multichannel acquisition, where $S$ denotes the coil sensitivities and $F$ represents the Fourier transform, and $P$ denotes the undersampling pattern. Since $A$ is an ill-conditioned operator, it is difficult to reconstruct the image $x$ from $y$ accurately and stably. Therefore, it is necessary to introduce a prior $p\left(x\right)$ of $x$ to realize regularization. This paper’s main task is to accurately estimate the distribution $p\left(x\right)$ of the fully sampled image $x$ from the undersampled data $y$.

Before presenting our method, we make the following key assumption:

According to Assumption 1, we can obtain the conditional distribution $p(x|y,\theta )$, and if we can further obtain the distribution of the parameters $\theta$, we can infer the distribution of image $x$ using the following Bayesian formula

Therefore, our next objective is to estimate the distribution of $p\left(\theta \right)$. Defining undersampled $k$-space data pairs $D:=\left\{\right(y_i,{y^{\prime }}_i){\}}_{i=1}^N$, we estimate $p\left(\theta |D\right)$ by Bayesian inference, i.e., by using $q(\theta |{\mu }_{\theta },{\sigma }_{\theta })$ to estimate $p\left(\theta |D\right)$, where

$\theta =\{{\theta }_s{\}}_s$, ${\mu }_{\theta }=\{{\mu }_{\theta ,s}{\}}_s$, and ${\sigma }_{\theta }=\{{\sigma }_{\theta ,s}{\}}_s$. In this paper, we use a Bayesian convolutional neural network (BCNN) to represent $f_{\theta }$ and output the distribution $q(\theta |{\mu }_{\theta },{\sigma }_{\theta })$. The framework of self-supervised Bayesian DL is shown in Figure 1. In particular, the architecture of mapping $f_{\theta }$ or BCNN can be directly adopted from the POCS-SPIRiT model driven neural network [6]. Introduced randomness, the network parameters contain standard deviation ${\sigma }_{\theta }$ and mean ${\mu }_{\theta }$. For parameter $\theta$, we first generate normal distribution noise $ϵ\sim {\prod }_{s}N({ϵ}_s|0,1)$, and then determine the convolutional kernel by sampling, i.e., $\theta ={\mu }_{\theta }+{\sigma }_{\theta }ϵ$.

The tight approximation is achieved by minimizing the following KL-divergence:

The final calculation result is as follows:

The calculation process of Proposition 3.1 is shown in Appendix .1.

Bringing the above estimated $q\left(\theta |{\mu }_{\theta },{\sigma }_{\theta }\right)$ into the Bayesian equation (3), we can obtain the distribution $p\left(x\right)$. However, the integral in (3) is difficult to calculate exactly in general. If $\nabla \log p\left(x\right)$ is accessible, samples obeying the distribution $p\left(x\right)$ can be collected according to Langevin MCMC sampling (1).

Following the score matching method, we can estimate $\nabla \log p\left(x\right)$ in a self-supervised manner by the following optimization problem.

The calculation process of Proposition 3.2 is shown in Appendix .2.

Based on Proposition 3.2, the score function at different perturbation levels can be learned, and then the desired samples are obtained by performing MCMC sampling according to them. More specifically, we perturb the $f_{\theta }\left(y\right)$ by Gaussian noise with scales $\{{\epsilon }_i{\}}_{i=1}^T$ that satisfies ${\epsilon }_1<{\epsilon }_2<\cdots <{\epsilon }_T$. Let $p_{{\epsilon }_i}(\tilde{x},|y,\theta ))=N(\tilde{x}|f_{\theta }\left(y\right),{\epsilon }_i^{2}I)$ and perturbed data distribution is $p_{{\epsilon }_i}\left(\tilde{x}\right)=\int p_{{\epsilon }_i}(\tilde{x}|y,\theta )p\left(y\right)p\left(\theta \right)\text{d}y\text{d}\theta$. If ${\epsilon }_1$ is chosen such that ${\epsilon }_1={\gamma }_1$, then $p_{{\epsilon }_1}\left(x\right)=p\left(x\right)$ holds. Based on Proposition 3.2, we can estimate $\nabla \log p_{{\epsilon }_i}\left(\tilde{x}\right)$ at all the scales by training a joint score function $s_{\varphi }(\tilde{x},{\epsilon }_i)$ with the following loss:

where $q\left(\theta \right):=q\left(\theta |{\mu }_{\theta },{\sigma }_{\theta }\right)$ is the minimizer of objective (5).

After training an score function, we can perform the condition Langevin MCMC sampling to solve problem (2), i.e.,

In particular, the choice of parameter ${\eta }_i$ follows that of literature [29], and the conditional Langevin MCMC sampling is detailed in Algorithm 1.

The advantages of our proposed method over existing self-supervised learning methods (learning the mapping $f_{\theta }$ directly from undersampled data) are twofold.

1. Intuitively, the score matching model (6) works by separating out the noise from the current image $\tilde{x}$, so even if $f_{\theta }\left(y\right)$ is not a perfectly clean image, the impact on the noise separation mechanism ($s_{\varphi }(\cdot )$) learning is relatively small.

2. Methodologically, we utilize the prior $p\left(x\right)=\int p(x|y,\theta )p\left(y\right)q\left(\theta \right)\text{d}y\text{d}\theta$ on a set of model $\theta$ over a set of data $y$ rather than the prior of a single model over a single set of data i.e., $f_{{\theta }_i}\left(y_i\right)$, ${\theta }_i\sim q\left(\theta \right)$ and $y_i\sim p\left(y\right)$. It has been shown that such ensemble models can often outperform single models [37].

The schematic diagram of the BCNN (including $f_{\theta }$) architecture is illustrated in Figures 2. Since Assumption 1 stems from the generalization of ”linear predictability” and ”flat invariance” of the $k$-space data, we followed the SPIRiT-POCS model neural network [6] for representing function $f_{\theta }$. In particular, to introduce randomness, we replaced the last layer of the module with a Bayesian neural network. On the other hand, we used the NCSNv2 network [29] to learn the score function, whose specific parameters are set: ${\epsilon }_L=50$, ${\epsilon }_1=0.0066$, number of classes is 266, ema is true, ema rate is 0.999.

The ADAM [16] optimizer with ${\beta }_1=0.9,{\beta }_2=0.999$ is chosen for optimizing loss functions (6) and (5). The size of the mini batch is 1, the number of epochs is 200 and the learning rate is ${10}^{-4}$. As shown in Figure 3, in constructing undersampled data pairs, we followed [34] to obtain $y^{\prime }$ by further Gaussian undersampling of $y$. The models were implemented on an Ubuntu 20.04 operating system equipped with an NVIDIA A6000 Tensor Core (GPU, 48 GB memory) in the open PyTorch 1.10 framework [25] with CUDA 11.3 and CUDNN support.

In this study, the quantitative evaluations were all calculated on the image domain. The reconstructed and reference images were derived using an inverse Fourier transform followed by an elementwise square-root of sum-of-the squares (SSoS) operation, i.e. $z\left[n\right]=\left({\sum }_{i=1}^{N_c}|x_i\right[n]{|}^2{)}^{\frac{1}{2}}$, where $z\left[n\right]$ denotes the $n$-th element of image $z$, and $x_i\left[n\right]$ denotes the $n$-th element of the $i$th coil image $x_i$. For quantitative evaluation, the peak signal-to-noise ratio (PSNR), normalized mean square error (NMSE) value, and structural similarity (SSIM) index [33] were adopted.

In this section, we tested the performance of various methods when sampling patterns and data types (anatomies) are consistent in training and testing. Figure 4 shows the reconstruction results for 5-fold random undersampling of fastMRI knee data and 4-fold uniform sampling of SIAT brain data. For the traditional PI method SENSE, the aliasing pattern remains in the reconstructed image and noise is amplified to the point that image details are obliterated. An obvious aliasing pattern remains in the reconstructed images for the self-supervised learning method. Although the noise is well suppressed for the supervised learning method, a tiny aliasing pattern remains in the reconstructed image where the red arrow points. The reconstructed image is also blurrier, and some texture information is lost from the visual perception. Our proposed method performs well in aliasing pattern suppression and image texture detail recovery. In particular, it is worth mentioning that the proposed fully-sampled-data-free method outperforms the conventional supervised DL method that requires a fully sampled training dataset, which is of practical significance.

The competitive quantitative results of the above methods are shown in Table 1. Our method consistently outperforms traditional SENSE, conventional self-supervised and supervised methods characterized by quantitative evaluations. The above experiments confirm our method’s competitiveness under consistent training and testing environments.

Good generalization ability is necessary for realistic applications of DL methods. However, as mentioned earlier, some existing self-supervised learning MR reconstruction methods are generalized from $k$-space interpolation, which will not work when sampling patterns and $k$-space data structures are inconsistent in training and testing. The proposed method eventually learns the distribution of fully sampled MR images, free from sampling patterns and data structure limitations. Therefore, the proposed method tends to have better generalization ability. To verify this claim, we designed two experiments as follows.

Theoretically, the accuracy of the proposed method regarding the estimation of data distribution depends mainly on Assumption 1. Although it is a natural generalization of ”linear predictability” and ”translational invariance” in $k$-space interpolation methods, there may still be some cases where it is not fully accurate. Therefore, we designed comparison experiments with the conventional score-based diffusion method trained on fully sampled data to verify the accuracy of the proposed self-supervised learning method on the data distribution estimation.

Figure 7 shows the reconstruction results of the conventional score-based method (termed score) and our proposed self-supervised score-based method (termed self-score) for fastMRI knee and SIAT brain data under 6-fold and 4-fold uniform undersampling, respectively. In terms of visual perception, the two methods perform almost identically. The quantitative metrics are shown in Table 4. It can be found that the proposed self-score method performs almost identically to the score method on the fastMRI knee dataset and even slightly better than the score on the SIAT brain dataset. This experiment validates the accuracy of the proposed self-supervised learning method for data distribution estimation.

We also compared the generalization capabilities of the two methods described above, i.e., trained on the fastMRI knee dataset but applied to reconstruct 6-fold uniformly undersampled fastMRI brain MR images. From the results presented in Figure 8, we can find that the proposed self-supervised score-based method trained on undersampled data performs very close to the method trained on fully sampled data in terms of the visual perception and error of the reconstructed images. The above experimental result further validates the accuracy of the proposed self-supervised learning method for estimating the distribution of the fully sampled MRI data.

As mentioned above, although Assumption 1 is a common assumption in $k$-space interpolation methods, there is no rigorous theory to guarantee its correctness. Therefore, it is a future work for us to design a self-supervised score-based method for MRI reconstruction without relying on Assumption 1.

For the parameter $\theta$ distribution estimation, we use a simple Gaussian distribution $q(\theta |{\mu }_{\theta },{\sigma }_{\theta })$ to approximate it. On the other hand, it is also possible to collect the sample ${\theta }_i$ by MCMC sampling and then approximate the distribution of $\theta$ by ${\sum }_i{\delta }_{\theta ={\theta }_i}$. Given the successful application of $q(\theta |{\mu }_{\theta },{\sigma }_{\theta })$ approximation in computer vision [24, 18], such an approach is adopted in this paper. It is worth noting that the MCMC method and the more accurate distribution approximation methods will be reserved as our options.

Finally, for conventional DL MRI reconstruction methods, the learned network can map the undersampled $k$-space data directly to the reconstructed image, which takes little time. Our method (including other score-based diffusion methods) needs to perform an iteration (MCMC sampling) to reconstruct the image, which takes a relatively long time. Reducing the reconstruction (sampling) time by introducing acceleration means, such as momentum, is a direction worth investigating in the future.

By Bayesian formula, the KL-divergence can be calculated as

For the first part, we have

For the second part, we have

Combining above two equations together, proof is completed.

This proof is based on literature [31]. First, we have

where $c_1$ denotes a constant. For the second term, we have

where the first and fourth equations make use of the derivative property of the log function. On the other hand, we have

where $c_2$ denotes a constant. The proof is completed.