Deep learning prediction of diffusion MRI data with microstructure-sensitive loss functions

Abstract

Deep learning prediction of diffusion MRI (DMRI) data relies on the utilization of effective loss functions. Existing losses typically measure the signal-wise differences between the predicted and target DMRI data without considering the quality of derived diffusion scalars that are eventually utilized for quantification of tissue microstructure. Here, we propose two novel loss functions, called microstructural loss and spherical variance loss, to explicitly consider the quality of both the predicted DMRI data and derived diffusion scalars. We apply these loss functions to the prediction of multi-shell data and enhancement of angular resolution. Evaluation based on infant and adult DMRI data indicates that both microstructural loss and spherical variance loss improve the quality of derived diffusion scalars.

1. Introduction

Prediction of diffusion MRI (DMRI) data is central to problems such as reconstruction (Hong et al., 2019a; Mani et al., 2020), phase correction (Aggarwal et al., 2019), denoising (Fadnavis et al., 2020), super-resolution (Tanno et al., 2017), harmonization (Mirzaalian et al., 2016, 2018; Tax et al., 2019; Ning et al., 2020; Moyer et al., 2020), longitudinal data imputation (Hong et al., 2019b), and microstructural analysis (Ye et al., 2019; Raffelt et al., 2012). Predicting unacquired data from acquired data is often regarded as a regression problem, which can be tackled using methods such as sparse representation (St-Jean et al., 2016) and kernel regression (Chen et al., 2019a). However, the problem can be ill-posed and challenging to solve when the amount of data that needs to be predicted increases. Learning-based algorithms, particularly those based on deep learning (Schultz, 2012; Ye, 2017; Koppers et al., 2017; Ye et al., 2019; Gibbons et al., 2019; Zhang et al., 2020), can be more effective for ill-posed regression problems since they can leverage information learned from training data (Hong et al., 2019b).

A key factor in determining success in prediction is the loss function, which defines the optimization target (LeCun et al., 2015). Common loss functions, including ℓ₁, ℓ₂, and Huber losses (Huber, 1992; Kim et al., 2018), measure point-wise differences between the prediction and the target and can yield good peak signal-to-noise ratio (PSNR) performance (Johnson et al., 2016; Dosovitskiy and Brox, 2016; Ledig et al., 2017). DMRI analysis often focuses on derived diffusion scalars rather than raw DMRI data; therefore, prediction accuracy of diffusion scalars needs to be taken into consideration when designing loss functions.

In this work, we propose two novel loss functions, called the microstructural loss (ML) and the spherical variance loss (SVL), to explicitly consider the quality of derived diffusion scalars during DMRI data prediction. The ML includes an ℓ₁ loss for DMRI data and an ℓ₁ loss for a diffusion scalar of interest, estimated by a microstructure estimator sub-network. The SVL is based on the spherical variance of the diffusion signal profile — a signal statistic of each shell that is related to tissue diffusion properties (Minh Huynh et al., 2019; Huynh et al., 2019). The SVL is a composite of an ℓ₁ loss for the DMRI data and an ℓ₁ loss for the spherical variance.

To demonstrate the effectiveness of the two loss functions, we apply them to two DMRI data prediction tasks: Multi-shell (MS) data prediction and angular resolution enhancement. The first task aims to predict MS data by estimating additional shells from data with fewer shells. MS data allows the use of advanced diffusion models, e.g., neurite orientation dispersion and density imaging (NODDI) (Zhang et al., 2012) and diffusion kurtosis imaging (DKI) (Jensen et al., 2005), for probing tissue microstructure. However, collecting MS data requires longer acquisition times. Few methods were proposed for predicting MS data. A popular method described in Koppers et al. (2016) employs a multi-layer perceptron (MLP) to learn data prediction with an ℓ₁ loss function.

Angular resolution enhancement can be achieved by predicting high angular resolution (HAR) DMRI data from its low-resolution counterpart. HAR DMRI data allows the use of sophisticated diffusion models for better estimation of local fiber geometric structures. Angular resolution enhancement allows increasing the angular resolution post-acquisition without prolonging the acquisition time (Yin et al., 2018; Chen et al., 2019b; Lyon et al., 2022).

The effectiveness of the ML and SVL will be demonstrated with data from the Baby Connectome Project (BCP) (Howell et al., 2019) and Human Connectome Project (HCP) (Van Essen et al., 2013). The experimental results indicate that these losses improve data prediction accuracy, both qualitatively and quantitatively. In what follows, we will first give a detailed description of our losses in Section 2 and report experiment results in Section 3. We will then provide additional discussions in Section 4 and finally conclude this work in Section 5.

2. Methods

In this section, we will first briefly introduce our data prediction network. We will then explain the details of our losses, i.e., the network-based ML and the statistic-based SVL. Finally, we will summarize and compare these two loss functions.

2.1. Data predictor

The data predictor (left in Fig. 1) estimates a non-linear function f_P(·) for predicting $\widehat{T}$, an estimation of the ground truth target data T, from input data S. To be consistent, we employ the same data predictor for the ML, the SVL, and all baseline losses.

Fig. 1.

Network Overview. The network consists of a data predictor and a microstructure estimator. The data predictor, shown on the left, is designed using an MLP. The concatenation branch, marked in yellow, reduces trainable parameters and preserves the source data. The microstructure estimator, shown on the right, is a key component of the ML and is realized using an MLP. f_E, T, and $\widehat{T}$ denote the microstructure estimator, target data, and predicted data, respectively.

In practice, the data predictor can be designed using a variety of network architectures, e.g., MLPs, convolutional neural networks (LeCun et al., 1998), and graph convolutional neural networks (Bronstein et al., 2017). In this work, we implement the data predictor utilizing an MLP for three main reasons: (i) MLP is widely employed in DMRI data prediction (Golkov et al., 2016; Koppers et al., 2016); (ii) MLP is straightforward to implement; and (iii) Most importantly, the main focus of our work is to investigate the effectiveness of novel loss functions; therefore we opt to adopt a simple and well-studied network for the data predictor. A comprehensive investigation of different data predictor architectures is beyond the scope of this paper.

MLP is a feedforward artificial neural network consisting of at least three layers. The first/last layer is an input/output layer. Therefore, the number of neural nodes in the input/output layer is determined by the dimension of the input/output data. The intermediate layers are hidden layers that encode features of different levels. Following (Golkov et al., 2016), we utilize three hidden layers, each with 150 neural nodes, matching approximately the number of our input measurements. The feature vector l_i of layer i is computed using

$$l_i=a\left(W_i{l}_{i-1}+b_i\right),$$ (1)

where a(·) is a non-linear activation function, ReLU. W_i and b_i are learned weights and biases. The linear representations of the features of a layer are non-linearly transformed with a ReLU activation function to give the features of the next layer. MLP is trained with a backpropagation algorithm (Rumelhart et al., 1988), which propagates errors, computed with respect to a loss function, through the network in a back-to-front manner so that trainable network parameters can be updated to eventually minimize the loss.

We predict data missing in the source and then utilize a concatenation branch to combine the predicted data and the source data for a complete output (Fig. 1). This strategy is useful for angular resolution enhancement, where the low angular resolution data is a subset of its high angular resolution counterpart. However, this strategy needs to be modified for MS data prediction (Section 3.2), since the source and target data are acquired using two noncollinear sets of different gradients. In this case, we remove the concatenation branch and directly feed the outcome of the data predictor to the microstructure estimator.

2.2. Microstructure Loss (ML)

Microstructure loss is based on a microstructure estimator network (right in Fig. 1) that estimates a microstructure scalar of interest based on the output of the data predictor. The microstructure loss consists of an ℓ₁ term for DMRI data and an ℓ₁ term for the microstructure scalar, i.e.,

$$L_{\text{ML}}=\frac{1}{N(T)}\Vert \widehat{T}-T{\Vert }_1+{\lambda }_{\text{ML}}{\Vert f_{\text{E}}(\widehat{T})-f_{\text{E}}(T)\Vert }_1,$$ (2)

where $\widehat{T}$ is the prediction of target data T, N(T) is the number of elements in T, λ_ML is a parameter balancing the contributions of the two ℓ₁ terms, and f_E(·) is a non-linear function realized via the microstructure estimator.

We first train the microstructure estimator using T and the microstructure scalar estimated from T. We then fix the microstructure estimator and train the data predictor using paired source and target data.

Similar to the data predictor, we design the microstructure estimator f_E using MLP with three hidden layers, each with 150 neural nodes. The number of neural nodes in the input layer is determined by the dimension of T.

2.3. Spherical Variance Loss (SVL)

Unlike the network-based ML, the SVL is a mathematical loss based on signal statistics (Fig. 2). The spherical variance of the signal is related to tissue microstructural properties (Novikov et al., 2018; Xu et al., 2018; Minh Huynh et al., 2019). Unlike point-wise losses, such as ℓ₁, ℓ₂, and Huber losses (Johnson et al., 2016; Dosovitskiy and Brox, 2016; Ledig et al., 2017), the SVL promotes q-space sharpness and avoids bias in derived diffusion scalars, such as fractional anisotropy, due to q-space smoothing.

Fig. 2.

Spherical variance loss. The spherical variance of the predicted DMRI data is employed to preserve q-space signal sharpness. T and $\widehat{T}$ denote the target data (ground truth) and predicted data, respectively. V and $\widehat{V}$ are the corresponding spherical variances.

We denote a target diffusion signal as T_b(g), where b and $g\in {\mathbb{S}}^2$ are the b-value and gradient direction, respectively. The spherical variance V_b is computed as the second-order central moment of signals acquired for the same b-shell, i.e.,

$$V_b=\frac{1}{4\pi }{\int }_{{\mathbb{S}}^2}{\left(T_b(g)-M_b\right)}^{2}dg,$$ (3)

where M_b is the corresponding first order raw moment defined as

$$M_b=\frac{1}{4\pi }{\int }_{{\mathbb{S}}^2}{T}_b(g)dg.$$ (4)

We combine the ℓ₁ terms for DMRI data and spherical variances to form the SVL:

$$L_{\text{SVL}}=\frac{1}{N(T)}\Vert \widehat{T}-T{\Vert }_1+\frac{{\lambda }_{\text{SVL}}}{N(V)}\Vert \widehat{V}-V{\Vert }_1,$$ (5)

where $\widehat{V}$ is a vector containing the spherical variances for different shells of the predicted DMRI data. Correspondingly, V is the vector for the target DMRI data. N(V) is the number of elements in V, which is determined by the number of shells. λ_SVL is a parameter trading off the contributions of the two ℓ₁ terms.

The SVL constrains the prediction to match not only the DMRI data, but also higher-level signal statistics, such as the spherical variance. In our implementation of the SVL, we utilize the same data predictor as in the ML for fair comparison.

2.4. Method summary

Two loss functions are introduced for improving the quality of diffusion scalars derived from the predicted DMRI data. They are however different in three aspects:

1. Motivation: The ML is a network-based loss function with a microstructure estimator sub-network to infer microstructural properties based on the predicted DMRI data. In contrast, the SVL is formulated based on the fact that tissue properties are related to the spherical variance.

2. Training: The SVL is trained in a standard end-to-end manner. Training the ML involves first training the microstructure estimator and then training the data predictor with the microstructure estimator fixed. Therefore, the SVL is more training friendly than ML.

3. Implementation Complexity: The ML integrates the microstructure estimator in its framework. In contrast, the SVL directly computes the loss value using a closed-form equation, which is easier to implement than the ML.

3. Experiments

3.1. Datasets

3.1.1. Dataset – Multi-shell data prediction

We demonstrate the effectiveness of the ML and SVL in MS data prediction with the BCP data (Howell et al., 2019). All data were acquired using a Siemens 3T Magnetom Prisma MRI scanner with two imaging protocols: (i) 2-Shell: b = 1500, 3000 s/mm² with a total of 74 non-collinear gradient directions and (ii) 6-Shell: b = 500, 1000, 1500, 2000, 2500, 3000 s/mm² with a total of 144 non-collinear gradient directions. The two protocols shared some common imaging parameters, including 140 × 105 imaging matrix, 1.5×1.5×1.5 mm³ resolution, TE = 88 ms, TR = 2365 ms, and 32-channel receiver coil. Note that the 2-shell and 6-shell data were acquired in the same imaging session. EDDY (Andersson and Sotiropoulos, 2016) and TOPUP (Andersson et al., 2003) were performed on the 2-shell and 6-shell data combined so that they are aligned.

Here we aim to predict the 6-shell data from its 2-shell counterpart. Our dataset consisted of 10 randomly selected individuals 1 year to 2 years of age. We extracted voxels to construct a large-scale training dataset with more than one hundred thousand samples.

3.1.2. Dataset – Angular resolution enhancement

We utilized the HCP data (Van Essen et al., 2013) for angular resolution enhancement. The data was acquired using a customized Siemens 3T Connectome Skyra MR scanner with the following imaging protocol: b = 1000, 2000, 3000 s/mm², 90 gradient directions per shell, 145 × 174 imaging matrix, 1.25×1.25×1.25 mm³ resolution, TE = 89 ms, TR = 5500 ms, and 32-channel receiver coil (Van Essen et al., 2013).

We downsampled the data angularly with a factor of 5 by uniformly selecting 18 gradient directions for each shell to form the low angular resolution (LAR) data with a total of 54 gradient directions. The data predictor was trained to estimate the data of the remaining 216 gradient directions, which were combined with the LAR data for a complete output of 270 gradient directions.

Our dataset consisted of 16 randomly selected subjects. We utilized 6 subjects for training and the other 10 subjects for testing. We constructed the training dataset using a sufficiently large number of voxel samples.

3.2. Implementation details

For training, we randomly extracted 20,000 voxels from the brain region of every training subject to construct a training dataset containing in total 6 × 20,000 = 120,000 samples. The two tasks were trained in a supervised manner. The proposed methods were implemented using TensorFlow 1.2 (Abadi et al., 2016). We trained the networks using the ADAM optimizer (Kingma and Ba, 2014) with an initial learning rate of 0.0001 and an exponential decay rate of 0.95. The batch size was set to 1024. Over-fitting was avoided by early stopping. The networks were trained using an NVIDIA GeForce GTX 1080 Ti GPU with 11 GB RAM. For ML, we utilized generalized fractional anisotropy (GFA) (Tuch, 2004), a widely-adopted model-free diffusion scalar, as the prediction target of the microstructure estimator.

3.3. Parameter settings

We set the two tuning parameters in the ML and SVL as λ_ML = 0.5 and λ_SVL = 1.0. Parameter selection was based on grid search with [0.01, 0.1, 0.5, 1.0, 1.5, 2.0, 10.0]. More details are described in Section 3.6.2.

3.4. Baseline losses

We compared our method with various baseline losses described below.

1. ℓ₁ loss: Due to its robustness to outliers, the ℓ₁ loss is one of the most popular loss functions used in regression tasks. We define the ℓ₁ loss as

   $$L_1=\frac{1}{N(T)}\Vert \widehat{T}-T{\Vert }_1.$$ (6)
2. ℓ₂ loss: The ℓ₂ loss is another popular loss function and is defined as

   $$L_2=\frac{1}{N(T)}\Vert \widehat{T}-T{\Vert }_2^2.$$ (7)

   Compared with the ℓ₁ loss, the ℓ₂ loss is less robust to outliers, but eases training (Zhao et al., 2016).

3. Huber loss: The Huber loss can be viewed as a combination of ℓ₁ loss and ℓ₂ loss. Following Huber (1992) and Kim et al. (2018), we define the Huber loss as

   $$L_{\text{Huber}}=\{\begin{array}{ll}\frac{1}{2}\Vert \widehat{T}-T{\Vert }_2^2,\hfill & \text{for}\Vert \widehat{T}-T{\Vert }_1<\delta ,\hfill \\ \delta \Vert \widehat{T}-T{\Vert }_1-\frac{1}{2}{\delta }^2,\hfill & \text{otherwise}\hfill \end{array}$$ (8)

   where δ is a parameter balancing the ℓ₁ loss and the ℓ₂ loss. We set δ = 0.5 following Kim et al. (2018).

3.5. Evaluation methods

1. Peak signal-to-noise ratio (PSNR): Quantitative evaluation was performed using PSNR:

   $$\text{PSNR}=20{\log }_{10}\frac{\text{MAX}}{\text{RMSE}},$$ (9)

   where MAX is the maximum value of the data and RMSE is the root mean square error between the prediction and the ground truth.

2. Structural similarity index (SSIM): Quantitative evaluation was also performed using SSIM (Wang et al., 2004):

   $$\begin{array}{l}\mathrm{SSIM}(\widehat{T},T)=\\ \frac{\left(2\mu (\widehat{T})\mu (T)+C_1\right)\left(2\sigma (\widehat{T},T)+C_2\right)}{\left({\mu }^2(\widehat{T})+{\mu }^2(T)+C_1\right)\left({\sigma }^2(\widehat{T})+{\sigma }^2(T)+C_2\right)},\end{array}$$ (10)

   where are $\mu (\widehat{T})$ and $\sigma (\widehat{T})$ are the mean and variance of the prediction, $\sigma (\widehat{T},T)$ is the covariance of the prediction and the ground truth, C₁ and C₂ are two positive constants that ensure that the denominator is non-zero.

3. Microstructure scalars: We considered various microstructure models for evaluation. The corresponding microstructure scalars included GFA given by q-ball imaging (Tuch, 2004), mean kurtosis (MK) given by DKI (Jensen et al., 2005), as well as orientation dispersion (OD), neurite density (ND), and cerebrospinal fluid volume fraction (CSFVF) given by NODDI (Zhang et al., 2012). GFA and DKI scalars were computed using Dipy (Garyfallidis et al., 2014). NODDI scalars were computed using AMICO (Daducci et al., 2015).

4. Absolute difference (AD) map: The AD map reflects voxel-wise differences between the prediction and the ground truth.

5. Mean absolute difference (MAD): MAD was computed as the mean of the AD values in the brain region.

6. Fiber orientation distribution functions (ODFs): We generate fiber ODFs using the multi-shell multi-tissue model described in Yap et al. (2016).

3.6. Results – BCP multi-shell data prediction

3.6.1. Quantitative evaluation

Results shown in Figs. 3 and 4 indicate that both the ML and SVL significantly improve the average PSNR and SSIM for GFA, MK, and NODDI scalars (OD, ICVF, and CSFVF combined). Compared with the SVL, the ML performs better in improving the estimation of GFA, but not MK and NODDI scalars.

Fig. 3.

Quantitative evaluation via PSNR — MS data prediction. PSNR comparison of different losses in terms of GFA, MK, and NODDI.

Fig. 4.

Quantitative evaluation via SSIM — MS data prediction. SSIM comparison of different losses in terms of GFA, MK, and NODDI.

3.6.2. Influence of λ_ML and λ_SVL

Fig. 5 indicates that, with the increase of λ_ML, the GFA PSNR reaches maximum and then slightly fluctuates, while the MK PSNR continues to increase. The NODDI PSNR peaks at λ_ML = 0.5. We chose λ_ML = 0.5 based on NODDI scalars since GFA and MK improvements are already significant. Fig. 6 shows that the NODDI PSNR peaks at λ_SVL = 1.0. Experiments were performed using the dataset for MS data prediction.

Fig. 5.

Influence of λ_ML. PSNR changes with respect to λ_ML for GFA, MK, and NODDI. PSNR peaks at λ_ML =0.5 for NODDI scalars.

Fig. 6.

Influence of λ_SVL. PSNR changes with respect to λ_SVL for GFA, MK, and NODDI. PSNR peaks at λ_SVL =1.0 for NODDI scalars.

3.6.3. Qualitative evaluation

We randomly selected one subject for qualitative inspection. Fig. 7 indicates that the ML gives the GFA image that is closest to the ground truth. This is further confirmed by the AD maps shown in the figure. The SVL gives the second-best AD map with smaller AD values. The MAD values indicate that the ML and SVL yield better performance than the ℓ₁, ℓ₂, and Huber losses.

Fig. 7.

GFA Images — MS data prediction. Visual comparison of GFA images and associated AD maps. The numbers at top left corners are the mean AD values.

We further investigated the influence of different losses on fiber ODFs. The results, shown in Fig. 8, indicate that the SVL and ML give clean and coherent fiber ODFs that are closer to the ground truth than the ℓ₁, ℓ₂, and Huber losses.

Fig. 8.

Fiber ODFs — MS data prediction. The ML and SVL yield fiber ODFs that are closest to the ground truth.

3.7. Results – HCP angular resolution enhancement

3.7.1. Quantitative evaluation

The quantitative results, shown in Figs. 9 and 10, are consistent with our observations in Figs. 3 and 4, where the ML and SVL outperform baseline losses. Specifically, both the ML and SVL improve the GFA, MK, and NODDI results in comparison with the baseline losses. The ML gives the best results for GFA and the SVL gives the best results for MK and NODDI. Note that the optimal parameters determined in Section 3.6.2 were used.

Fig. 9.

Quantitative evaluation via PSNR — Angular resolution enhancement. PSNR comparison of different losses in terms of GFA, MK, and NODDI.

Fig. 10.

Quantitative evaluation via SSIM — Angular resolution enhancement. SSIM comparison of different losses in terms of GFA, MK, and NODDI.

3.7.2. Qualitative evaluation

Fig. 11 indicates that the ML yields the GFA image that is closest to the ground truth. The SVL is the second-best loss and reduces the AD values significantly.

Fig. 11.

GFA Images — Angular resolution enhancement. Visual comparison of GFA images and associated AD maps. The numbers at top left corners are the mean AD values.

Fig. 12 indicates that the fiber ODFs given by the SVL and ML are closer to the ground truth than the baseline losses. The improved fiber ODFs given by our losses can benefit subsequent processing, such as tractography (Aydogan and Shi, 2016) and connectomics (Li et al., 2016).

Fig. 12.

Fiber ODFs – Angular resolution enhancement. The ML and SVL yield fiber ODFs that are closest to the ground truth.

4. Discussion

A straightforward approach to take into account diffusion properties in the designing of loss functions is by directly computing the microstructure estimations from the predicted data. However, the fitting of microstructural models can be non-linear and the computational complexity can be too high for efficient training of deep neural networks, especially considering the fact the loss function needs to be repeatedly computed for numerous times. This limitation is overcome by the ML and SVL, which are computationally efficient.

In addition to the spherical variance, other spherical moments can be considered to design the loss. A typical example is the spherical mean, which is highly related to the tissue properties of brain white matter (Kaden et al., 2016). However, our evaluation shows that incorporating spherical mean into the loss affects the results only marginally. Currently, we did not consider other high-order spherical moments since (i) their connection with the diffusion properties of brain white matter is unclear and (ii) they are sensitive to noise (Flusser et al., 2009).

The advantages arising from our losses are further reflected by the derivatives of the predicted DMRI data. For example, by improving fiber ODFs via MS data prediction (Fig. 8), our losses improve the estimation of apparent fiber densities (Raffelt et al., 2012) (Fig. 13), orientations (Fig. 14), and (Fig. 14), and tractography (Fig. 15).

Fig. 13.

Apparent fiber densities. PSNR comparison of losses in terms of apparent fiber densities derived from fiber ODFs.

Fig. 14.

Axonal orientations. Orientational discrepancy (Yap et al., 2011) of fiber ODFs.

Fig. 15.

Tractography. Comparison of the inferior fronto-occipital fasciculus (IFO) given by the different losses.

5. Conclusion

In this paper, we have proposed two novel loss functions for deep learning prediction of DMRI data. The proposed losses sufficiently consider the microstructural properties of brain white matter during data prediction, avoiding the drawbacks of existing losses that only consider signal prediction error. We demonstrated the effectiveness of our losses, the ML and the SVL, by applying them to the prediction of 6-shell DMRI data from their 2-shell counterparts and the angular resolution enhancement of q-space downsampled DMRI data. Experimental results demonstrate that the ML and SVL substantially outperform standard loss functions.

Data availability

The authors do not have permission to share data.