Spatio-Angular Noise2Noise for Self-Supervised Denoising of Diffusion MRI Data

Introduction

Diffusion MRI (DMRI) suffers from heavy noise, especially when the $$$b$$$-value is large. This is because DMRI measures the attenuation signals, resulting in a decrease in Signal-to-Noise Ratio (SNR). Existing traditional methods^1-2 for DMRI denoising show unsatisfactory performance in challenging scenarios, while deep learning methods³ usually require paired noisy-clean data. To this end, we propose a novel self-supervised DMRI denoising model, called Spatio-Angular Noise2Noise (SAN2N). Our SAN2N jointly considers the spatial and angular information of DMRI data in a unified framework. It generates angular neighboring DMRI data based on the geometric relationships between $$$q$$$-space sampling points. The resulting data and the original one are then fed to two neighborhood-based $$$x$$$-space sub-samplers to extract 4D similar patches in the spatio-angular domain. Extensive experiments on two datasets demonstrate the effectiveness of our SAN2N.

Methods

An overview of our SAN2N is shown in Figure 1(a, b). First, we create the DMRI data $$$\mathbf{y}^{'}=a(\mathbf{y}) \in \mathbb{R}^{L\times W\times H\times C}$$$ that is similar to the original DMRI data $$$\mathbf{y} \in \mathbb{R}^{L\times W\times H\times C}$$$ in $$$q$$$-space with an angular sampler $$$a(\cdot)$$$, where $$$L$$$, $$$W$$$, and $$$H$$$ denote the spatial dimensions of DMRI data, i.e., length, width, and height, while $$$C$$$ denotes the number of gradient directions. Afterward, these two sets of DMRI data $$$(\mathbf{y}^{'}, \mathbf{y})$$$ are fed to the spatial sub-samplers $$$\mathcal{S}=\{s_{1}(\cdot), s_{2}(\cdot)\}$$$ of our SAN2N. The samplers generate angular noisy data pairs $$$\{s_{1}(\mathbf{y}), s_{1}(\mathbf{y}^{'})\}$$$ and spatial noisy data pairs $$$\left\{s_{1}(\mathbf{y}), s_{2}(\mathbf{y})\right\}$$$. The denoising network $$$f_{\theta}(\cdot)$$$ parameterized by $$$\theta$$$ is trained with these noisy pairs by minimizing the following loss:
$$\underset{\theta}{\text{arg min}}~\mathbb{E}_{\mathbf{x}, \mathbf{y}}\left\|f_{\theta}\left(s_{1}(\mathbf{y})\right)-s_{2}(\mathbf{y})\right\|^{2} +\|f_{\theta}\left(s_{1}(\mathbf{y})\right)-s_{1}(\mathbf{y}^{'})\|^{2},$$
where $$$x$$$ is the underlying ground-truth noise-free DMRI data.
The angular sampler shown in Figure 1(c) is detailed as follows:
1. For each gradient $$$\mathbf{q}_i$$$, find its nearest neighboring gradient $$$\mathbf{q}_{i^{'}}$$$ on the same shell based on their angle $$$\theta(\mathbf{q}_i, \mathbf{q}_{i^{'}})$$$.
2. If $$$\theta(\mathbf{q}_i, \mathbf{q}_{i^{'}}) > \delta$$$, assign the same Diffusion-Weighted Image (DWI) in $$$\mathbf{y}$$$ to $$$\mathbf{y}^{'}$$$, i.e., $$$\mathbf{y}^{'}(\mathbf{q}_i) = \mathbf{y}(\mathbf{q}_i)$$$, otherwise adopt the angular neighboring one $$$\mathbf{y}^{'}(\mathbf{q}_i) = \mathbf{y}(\mathbf{q}_{i^{'}})$$$.
3. Repeat steps one and two for all gradients to obtain $$$\mathbf{y}^{'}$$$.
The spatial sampler shown in Figure 1(d) is detailed as follows:
1. Divide $$$\mathbf{y}$$$ into $$$K^{3}$$$ neighborhoods with the same size of $$$\frac{L}{K}\times \frac{W}{K}\times \frac{H}{K}\times C$$$ according to the spatial information.
2. Repeat selecting a voxel in each neighbor randomly until all neighbors are visited.
3. Construct a sub-sampled data $$$s_{1}(\mathbf{y})$$$ with a size of $$$\frac{L}{K}\times \frac{W}{K}\times \frac{H}{K}\times C$$$ using the voxels from the last step.
4. Repeat steps two and three for another subsampled data $$$s_{2}(\mathbf{y})$$$.
We further design a regularized denoising loss to consider the difference between the angular and spatial noisy data pairs:
$$\mathcal{L}_{\text{rx}} = \left\|f_{\theta}\left(s_{1}(\mathbf{y})\right)-s_{2}(\mathbf{y})-\left(s_{1}\left(f_{\theta}(\mathbf{y})\right)-s_{2}\left(f_{\theta}(\mathbf{y})\right)\right)\right\|_{2}^{2}.$$
Similarly, we define the regularization term for angular noisy data pair as:
$$\mathcal{L}_{\text{rq}} = \|f_{\theta}\left(s_{1}(\mathbf{y})\right)-s_{1}(\mathbf{y}^{'})-(s_{1}\left(f_{\theta}(\mathbf{y})\right)-s_{1}(f_{\theta}(\mathbf{y}^{'})))\|_{2}^{2}.$$
Finally, we define a total loss $$$\mathcal{L}$$$ with four terms, including an $$$x$$$-space loss $$$\mathcal{L}_{\text{x}}$$$_, a $$$q$$$-space loss $$$\mathcal{L}_{\text{q}}$$$, and their regularization losses $$$\mathcal{L}_{\text{rx}}$$$ and $$$\mathcal{L}_{\text{rq}}$$$, i.e.,
$$\mathcal{L} =\underbrace{\left\|f_{\theta}\left(s_{1}(\mathbf{y})\right)-s_{2}(\mathbf{y})\right\|_{2}^{2}}_{\mathcal{L}_{\text{x}}} +\gamma_{\text{x}} \mathcal{L}_{\text{rx}}\\+ \lambda (\underbrace{\|f_{\theta}\left(s_{1}(\mathbf{y})\right)-s_{1}(\mathbf{y}^{'})\|_{2}^{2}}_{\mathcal{L}_{\text{q}}} + \gamma_{\text{q}}\mathcal{L}_{\text{rq}}),$$
where $$$\gamma_{\text{x}}$$$ and $$$\gamma_{\text{q}}$$$ are two parameters controlling the contributions of the regularization loss $$$\mathcal{L}_{\text{rx}}$$$ and $$$\mathcal{L}_{\text{rq}}$$$. $$$\lambda$$$ balances the contributions of $$$x$$$-space losses and $$$q$$$-space losses. After the model is trained, the denoising data $$$\mathbf{\hat{x}}$$$ is obtained using $$$\mathbf{\hat{x}}=f_{\theta}(\mathbf{y})$$$ in the testing stage.

Results

Dataset: We use two datasets for our experiments. The simulated dataset is generated using DW-POSSUM^4-5 with SNRs of 15, 17.5, 20, and 22.5. Data from the Developing Human Connectome Project (DHCP) constitutes the real dataset used for validation.
Experimental Results: We evaluate the results using the RMSE, PSNR, and SSIM. SAN2N outperforms all competing methods on DW-POSSUM across all evaluation metrics and different SNRs in Figure 2. Furthermore, we show the visual comparison of DWIs in Figure 3. SAN2N effectively improves the quality of DWIs by reducing the noise. Clear structures can be observed after removing the noise by our method in the close-up views. Finally, we evaluate our SAN2N on the real dataset and show the qualitative visual comparison of DWIs in Figure 4. Our method provides a denoised image with more clear structures. Furthermore, it preserves the structural details during noise reduction demonstrated by the residual maps shown in Figure 5.
Conclusion
We proposed a self-supervised deep learning model, SAN2N, for noise reduction of DMRI data. SAN2N integrates the spatial and angular denoising of DMRI data into a unified framework with a $$$q$$$-space sampler and two $$$x$$$-space sub-samplers. We perform extensive experiments on simulated and real datasets. The results demonstrate that SAN2N achieves the best performance in comparison with cutting-edge methods.