Diffusion Model-based FOD Restoration from High Distortion in dMRI

TL;DR

This work addresses the challenge of susceptibility-induced distortion in diffusion MRI that corrupts fiber orientation distributions (FODs) and hampers tractography. It proposes FOD-Diffusion, a diffusion-model-based restoration framework that employs volume-order encoding to generate SPHARM-based FOD volumes across orders up to $L=8$ and a frequency-balanced cross-attention mechanism to fuse information from all orders, conditioned on surrounding low-distortion FODs. On UK Biobank data, the method achieves high-fidelity FOD restoration with reduced RMSE and angular errors on ground-truth tests and demonstrably restores FOD integrity in high-distortion brainstem regions, improving CST and MCP tractography. The approach is memory-efficient, capable of training on a single GPU with $24$ GB, and offers practical benefits for reliable brain connectivity analysis in distorted regions, with potential extension to additional clinical datasets.

Abstract

Fiber orientation distributions (FODs) is a popular model to represent the diffusion MRI (dMRI) data. However, imaging artifacts such as susceptibility-induced distortion in dMRI can cause signal loss and lead to the corrupted reconstruction of FODs, which prohibits successful fiber tracking and connectivity analysis in affected brain regions such as the brain stem. Generative models, such as the diffusion models, have been successfully applied in various image restoration tasks. However, their application on FOD images poses unique challenges since FODs are 4-dimensional data represented by spherical harmonics (SPHARM) with the 4-th dimension exhibiting order-related dependency. In this paper, we propose a novel diffusion model for FOD restoration that can recover the signal loss caused by distortion artifacts. We use volume-order encoding to enhance the ability of the diffusion model to generate individual FOD volumes at all SPHARM orders. Moreover, we add cross-attention features extracted across all SPHARM orders in generating every individual FOD volume to capture the order-related dependency across FOD volumes. We also condition the diffusion model with low-distortion FODs surrounding high-distortion areas to maintain the geometric coherence of the generated FODs. We trained and tested our model using data from the UK Biobank (n = 1315). On a test set with ground truth (n = 43), we demonstrate the high accuracy of the generated FODs in terms of root mean square errors of FOD volumes and angular errors of FOD peaks. We also apply our method to a test set with large distortion in the brain stem area (n = 1172) and demonstrate the efficacy of our method in restoring the FOD integrity and, hence, greatly improving tractography performance in affected brain regions.

Figures (7)

• Figure 1: The residual distortion affects FODs and tractography. (a) The Topup method from FSL topup aligns the $b=0$ images from 2 opposite phase encoding directions to correct the susceptibility-induced distortion. (b) Corrupted FODs and failed tractography of the data in (a) illustrate the impact of the signal loss that cannot be recovered by distortion correction. (c) FODs from a low signal loss case and the successful fiber tracking results.
• Figure 2: The unconditioned DDPM model that only generates one FOD volume at each time failed in restoring the FODs. (a) The FODs have different numbers of volumes for each order. (b) Errors in individual volumes lead to erroneous FOD representations.
• Figure 3: The framework of the FOD-Diffusion model. Our FOD-Diffusion model takes the low-signal loss FODs as a condition in generation and uses the volume and order number encoding (denoted as V) to generate FODs in different frequency order. We also use low-signal loss FODs from all FOD volumes to extract the cross-attention information and help the generation of each FOD volume.
• Figure 4: Details of the frequency order-aware cross-attention calculation. We use the frequency order encoding to adjust the features from each order, and then we use a convolution layer to combine and select the features.
• Figure 5: Examples of the ground truth FODs and the FODs of our FOD-Diffusion model from the test set. Especially, (a) is the same case in Fig. \ref{['Figure2']}