Diffusion Model-based FOD Restoration from High Distortion in dMRI^††thanks: This work is supported by the National Institute of Health (NIH) under grants R01EB022744, RF1AG077578, RF1AG056573, RF1AG064584, R21AG064776, U19AG078109, and P41EB015922.

We propose a volume order-aware diffusion model to enable the diffusion model to generate FODs across all different frequency orders. Firstly, inspired by the time encoding in the diffusion model [9, 8], we propose a volume and frequency order-aware encoding (referred to as volume encoding) to enhance the model’s ability to distinguish FODs of different frequencies and orientations. The volume encoding encodes an array $[L,V]$, where $L$ and $V$ are the frequency order number and the volume number of the single FOD volume, respectively.

Moreover, we also use the low-signal loss regions as a condition to provide the model with sufficient constraints during generation. In the low-signal loss data, we assign high-signal loss regions to number 1 to distinguish them from the background.

We use the $L_{1}$ loss during training, and enhance the loss of the high-distortion regions, as shown in Eq. 1:

Since FODs at different orders have different gray-level ranges, they are normalized before generation. Components with $L=0$ are normalized to the range $[0,1]$, while other FODs are normalized to $[-1,1]$. This ensures that the regions outside the brain remain 0 for all normalized FODs.

The FODs represented by SPHARM have interdependence between different volumes, and we model the relationship between different FOD volumes to help the generation. We propose a frequency-balanced cross-attention method that extracts features from all FOD volumes and feeds them to the U-Net in the diffusion model through the cross-attention [5]. This method achieves balanced attention across each frequency order.

As shown in Fig. 3, for FOD volumes in each frequency order, we calculate their average to reduce the volume number to one in each frequency order. Specifically, for order $L=0$, as it only has one volume, it can be directly used. Then, we use copied U-Nets from the diffusion model with time encoding $t=0$ and volume encoding $V$ to extract features from the averaged FOD volumes in each frequency order. These features correspond to the output of the self-attention of each decoding block. Afterward, a frequency-order aware convolution with kernel size 1 is performed to select and combine the information that is most relevant to the FOD volume that needs to be generated.

We use copied U-Nets from the diffusion model’s U-Net to extract cross-attention features. The parameters in copied U-Nets are frozen during the training. We copy the parameters to the copied U-Net when we update the U-Net in the diffusion model. This approach reduces the number of parameters that need to be trained; hence, it reduces the training time and prevents overfitting.

More detail of the frequency order-aware convolution is shown in Fig. 4. Firstly, we use sine and cosine functions to embed the order number. Afterward, a multilayer perceptron (MLP) layer is added to calculate the encoding weights $[e_{1},e_{2},...,e_{n}]^{T}$. Eq. 2 is used to encode the features $F$ from each frequency order by adding the encoding weights:

The combined features from the 1×1 Conv are fed into the cross-attention through the Key and the Value in the U-Net of the diffusion model, and the Query is the feature from the single volume. The most useful information from other FOD volumes is selected from the Value to enhance the restoration of the current FOD volume under consideration [5].

Our FOD-Diffusion method was trained, validated, and tested on the UK Biobank dataset. Overall, this study utilized $n=1315$ data points. They were extracted from the dataset randomly, with ages ranging from 30 years old to 80 years old.

All extracted data were pre-processed using Topup [18] and Eddy [20, 21] for correcting distortion, eddy current, and head motion. Then we calculated the FODs with the highest order $L=8$ using the method in Ref. [6]. To extract low-residual distortion cases for model training, we used the method in Ref. [4] to calculate the severity map of residual distortion, and we extracted 143 data that have both low mean signal loss ($<0.25$) and high mean FOD integrity ($>0.09$), which contains about 10% of the data, for model training, validating and testing. We used $n=90$ data for training, $n=10$ data for validation, and $n=43$ data for testing. We masked out regions in the brainstem of these data to act as the high-signal loss regions. The remaining $n=1172$ data that have high residual distortion formed a large test set to test the performance of the model in real data with high distortion. The masks of high-distortion regions were extracted based on the residual distortion severity map.

The model was trained for 100,000 iterations with the batch size of 8. We used AdamW optimizer in the training, with a learning rate of $10^{-5}$ initially and $10^{-6}$ after 70,000 iterations. We validated the model in every epoch, and the model with the least validation loss was selected. The training took about 68.5 hours on an NVIDIA RTX A5000 GPU, with about 22 GB of GPU memory. We used the DDPM for 1000 time steps in the inference, and the v-prediction [27] was used in all experiments. The inference time for a single FOD volume was about 90 seconds.

Fig. 5 shows two examples of the FODs for the ground truth and the results of our FOD-Diffusion model. These examples are from the test set. FODs generated by our FOD-diffusion model have high angular similarity with the ground truth. For the similarity of intensity, although some FODs have different intensities with the ground truth, the intensity is similar for most FODs.

We then quantitatively compared our model with the unconditional DDPM and two ablation studies. The first ablation only takes the low-signal loss FOD volume (named vol) as the condition, and the second one also has the volume encoding (named enc). Table 1 shows the mean squared error of brain stem results from different methods. The unconditional diffusion model failed in FODs’ generation. Results of the ablation study show that all the improvements we added helped to improve the accuracy of the generated FODs.

We also compared the geometric differences between different methods, following the work in Ref. [22]. We first calculated the angles of the FOD peaks with the top three largest amplitudes that have peak values larger than 0.5 from both the ground truth FODs and the results FODs. Afterward, we calculated the corresponding smallest angular differences for the highest and second highest peak from the ground truth FODs (called “$1^{st}$ peak” and “$2^{nd}$ peak”) to the peaks of the result FODs. The results are shown in Table 2. All ablation studies significantly overcome the unconditional diffusion model, and our FOD-Diffusion model has less angular differences and standard derivations than other models in the ablation studies.

Because we do not have ground truth for high distortion data, we will use the group-wise distribution on the FODs’ integrity at the pons regions to evaluate the FOD-Diffusion model in the high signal loss UK Biobank data ($n=1172$). Here we compare our FOD-Diffusion method with the Topup method, as it is a representative registration-based method and the most widely used method in connectome research.

Fig. 6 shows the distribution of the FODs’ integrity at the pons region before and after the signal loss recovery. The 1172 data are evenly divided into five groups based on the severity of signal loss, named groups of “very slight”, “slight”, “medium”, “severe”, and “very severe” residual distortion, respectively. Each of the first four groups has 236 subjects, and the last group has 228 subjects. The mean intensity of the $L=0$ component at the pons ROIs (the atlas is defined in Ref. [28]) in Fig. 6(a) was calculated and used as a measure of FOD integrity in the evaluation following the work in Ref. [4]. This evaluation method is efficient because the intensity of the $L=0$ component represents the mean FOD value [6], and the mean FOD decreases with the increase of the residual distortion. Fig. 6 (b) shows that while the FOD integrity of the original FODs reduces with the increase of the severity in signal loss, the integrity of restored FODs have similar distributions across the groups. This shows that our method can recover signal loss for data with high residual distortions.

The trachography results for two high signal loss cases are shown in Fig. 7. Here, we calculated the corticospinal tract (CST) [11] and the middle cerebellar peduncle (MCP) [12] for both the Topup results and the results of our FOD-Diffusion method. These two fiber bundles share the same region of interest at the pons (the red mask in Fig. 6(a)) for generating the seeds. Fig. 7 shows that the tractography results of CST using FOD-Diffusion cover more pons region, and the MCP generated using FOD-Diffusion can cover the whole pons region. Therefore, our FOD-Diffusion method can successfully restore FOD-based fiber connectivity at the pons region.