Generative AI for rapid diffusion MRI with improved image quality, reliability, and generalizability

Since its introduction in 1985, diffusion MRI (dMRI) has seen widespread adoption. The number of publications on diffusion MRI has exploded from a few hundred per year in the 1990s to tens of thousands in recent years driven by rapidly expanding scientific and clinical applications, especially to neuroscience. DMRI can provide valuable clinical information and assess tissue microstructure; however, its low signal-to-noise ratio (SNR) can limit its diagnostic and quantitative accuracy (Jones, 2012). To improve SNR, most dMRI protocols require low angular and spatial resolution or a long scan time, which limits usage in many important clinical settings. Therefore, there is great interest in having short patient scan times without compromising SNR or spatial and angular resolution.

Supervised methods have been proposed to denoise brain dMRI scans, but they are limited by their lack of generalizability to different b-values, diffusion-encoding directions, signal representations or models, scanners, and/or patient populations. They are typically trained and validated on the same dataset, such as the Human Connectome Project (HCP) (Karimi & Gholipour, 2022; Tian et al., 2020). Therefore, unsupervised/self-supervised denoising methods are preferred (Fadnavis et al., 2020), yet these methods lack the performance of supervised techniques within the trained domain and can perform variably on out-of-domain (OOD) datasets.

Most dMRI denoising methods are evaluated qualitatively or by denoising a subset of the data and evaluating accuracy compared with the full dataset. In this paper, we also evaluate denoising using external validation via test–retest reliability for Diffusion Tensor Imaging (DTI) (Mukherjee et al., 2008) and Neurite Orientation Dispersion and Density Imaging (NODDI) (Zhang et al., 2012) metrics to ensure precision and via known Structural Covariance Networks (SCNs) (Li et al., 2012; Wahl et al., 2010) of those metrics in white matter (WM) and gray matter (GM) to ensure biological accuracy. Finally, we emphasize removing heavy-tail noise, which can lead to biased biophysical metrics.

We propose to use a Swin UNEt TRansformer (SWIN) model (Hatamizadeh et al., 2022) to denoise dMRI data conditioned on registered T1 scans. Unlike other supervised methods, which train with a small subset of HCP data (Karimi & Gholipour, 2022; Tian et al., 2020), we use the full HCP dataset, training with all b-values and diffusion-encoding directions, including data augmentations: random flipping, rotation, scaling, and k-space downsampling. Training on close to 300,000 3D volumes across 1021 subjects allows SWIN to learn a denoising function that generalizes well. We also train a UNet convolutional neural network without a transformer component to determine whether network architecture affects denoising performance. We validate our approach on a held-out HCP Retest dataset and three external datasets acquired in different patient populations using different scanners and dMRI protocols, outperforming state-of-the-art unsupervised/self-supervised methods in accuracy and repeatability. Fine-tuning SWIN on even a single subject improves performance on out-of-domain datasets. Finally, our approach can also super-resolve dMRI data from HCP.

Experiments utilized four datasets, including one of normal young adult volunteers separated into a training dataset and held-out dataset. The other three are OOD patient datasets for generalizability and clinical applicability. All datasets were acquired with the written informed consent of the participants or their legal guardians. Only deidentified data were used.

• —

  HCP: Young adults (Essen et al., 2013) scanned with 90 diffusion-encoding directions at b-values of b = 1000, 2000, 3000 s/mm² with 1.25 mm resolution. We used 1021 subjects for training and held out 44 subjects in the HCP Retest dataset.

• —

  TBI: 45 adult mild traumatic brain injury patients acquired two decades ago (Kuceyeski et al., 2019; Wahl et al., 2010) on a 3T GE scanner with 55 diffusion-encoding directions at b = 1000 s/mm² and 1.8 mm voxels (0.9 mm in xy after zero-interpolation in k-space). We randomly selected 5 subjects for fine-tuning and 40 subjects for testing.

• —

  SPIN: 45 children ages 8-12 years with neurodevelopmental disorders acquired on a 3T Siemens Prisma scanner with 64 and 96 diffusion-encoding directions at $b=1000,2500$ s/mm², respectively, with 2.00 mm voxels (Mark et al., 2023). We randomly selected 5 subjects for fine-tuning and 40 subjects for testing.

• —

  AHA: 8 children ages 8-18 years with hemorrhagic vascular malformations, primarily arteriovenous malformations (AVMs) acquired on a 3T GE MR750 with 55 diffusion-encoding directions at b = 2000 s/mm² and 2.00 mm voxels (zero-interpolated in-plane to 1.00 mm) over three sessions: prior to resection, 6 months postsurgery, and 1 year postsurgery (one subject had only two sessions). A ninth subject was used for fine-tuning.

For preprocessing, dMRI data were skull-stripped with Synthstrip (Hoopes et al., 2022), corrected for eddy current-induced distortions and subject movements with Eddy (Andersson & Sotiropoulos, 2016), and aligned to structural 3D T1 scans with Boundary-Based Registration (Greve & Fischl, 2009). T1 scans were skull-stripped using Synthstrip and segmented using SynthSeg (Billot et al., 2022). Due to lesions impacting SynthSeg performance in AHA data, T1 scans were instead segmented using Freesurfer’s recon-all command with the SynthSeg option enabled or with recon-all-clinical if recon-all failed. Finally, coregistered T1 and dMRI scans were resampled with 5th order spline interpolation at 1.25 mm, concatenated together, and used as inputs to the model. Negative values were clipped.

We measure denoising and scan speedup performance in fully sampled and subsampled data. We use the minimum directions for unique fits: 6 for DTI, 15 for 4th order spherical harmonic, and 28 for 6th order spherical harmonic along with 1 b = 0 s/mm² volume. We select the directions to minimize the condition number of the design matrix (Skare et al., 2000; Tian et al., 2022). For NODDI, we use the fully sampled HCP acquisition. We compare SWIN and UNet with five state-of-the-art unsupervised/self-supervised dMRI denoising methods: block-matching and 4D filtering (BM4D) (Maggioni et al., 2013), Marchenko-Pastur Principal Component Analysis (MPPCA) (Veraart et al., 2016), Tensor Marchenko-Pastur Principal Component Analysis (TPCA) (Olesen et al., 2023), DDM² (Xiang et al., 2023), and Patch2Self (P2S) (Fadnavis et al., 2020).

We measure mean absolute error (MAE) between ground truth fully sampled data and model predictions from subsampled data for DTI, including principal eigenvector (V1), fractional anisotropy (FA), axial diffusivity (AD), radial diffusivity (RD), and mean diffusivity (MD). For higher order spherical harmonics, we used the Jensen–Shannon distance (JSD) between fully sampled ground truth and model predictions, projected onto a uniformly distributed 362-direction hemisphere (Cohen-Adad et al., 2011). DTI only utilized the lowest shell (b = 1000 s/mm²) for multishell datasets, whereas spherical harmonic estimation used every shell. For super-resolution, dMRI data from an HCP subject were k-space downsampled by a factor of 2 and then upsampled to emulate a low resolution acquisition.

We explore clinical translation using the AHA dataset of children with intracerebral hemorrhage before and after intervention. We measure global WM and GM DTI metrics and perilesional changes in DTI microstructure over time from both subsampled and fully sampled shells. To investigate signal rectification due to low SNR, we measure the heavily attenuated free water signal in the lateral ventricles and compare the resulting intensity histograms produced by each denoising method.

We assess repeatability using HCP Retest data by measuring within-subject coefficient of variation (CoV) for DTI and NODDI across two sessions. We consider both fully sampled and subsampled data for DTI test–retest reliability, but use full multishell data for NODDI. For repeatability in WM regions, we perform tract-based spatial statistics (TBSS) using the Johns Hopkins University (JHU) atlas (Smith et al., 2006).

SCNs are derived by computing region-wise coefficients of determination among subjects. We measure mean absolute difference between first and second sessions for SCN repeatability, using Desikan–Killiany–Tourville atlas regions for cortical GM SCNs and JHU WM atlas regions for WM SCNs. We evaluate DTI SCN repeatability on subsampled b = 1000 s/mm² data and NODDI SCN repeatability on fully sampled multishell data. Accuracy is determined from MAE between denoised subsampled DTI SCNs and ground truth fully sampled DTI SCNs.

For both DTI estimation and microstructural repeatability validation, we test for better SWIN performance using subject-wise, one-tailed, paired t-tests between it and the self-supervised methods for the HCP dataset, and between SWIN model with fine-tuning on one subject and self-supervised denoising for the three OOD datasets. For test–retest reliability of regional results, we average CoV across WM and GM regions, respectively, when conducting these statistical inference tests. We set a significance level of 0.05.

SWIN (Hatamizadeh et al., 2022; Liu et al., 2021) was implemented using PyTorch (Paszke et al., 2019) and MONAI with default model hyperparameters (Cardoso et al., 2022) (Fig. 1) and trained on an NVIDIA V100 GPU using mean-squared error loss between model output and ground truth. Ground truth was obtained from a 6th order spherical harmonic fit for each shell projected onto the acquired directions. AdamW (Loshchilov & Hutter, 2017) optimization was conducted with a learning rate of 1e-5, gradient clipping with a maximal norm of 1.0, and 16-bit precision for 14 epochs. During training, we first downsampled the dMRI scan with a probability of 0.5 in frequency space to an anisotropic resolution between 1.25 and 3 mm and linearly upsampled back to 1.25 mm resolution (Lyu et al., 2019). Patches of 128 x 128 x 128 were randomly cropped, flipped, and rotated in 90° increments along all axes. The input dMRI patch was standardized to unit variance and the input T1 patch was standardized to a standard deviation uniformly log-scaled between 0.25 and 4.0. For inference, we use a sliding window with an overlap of 0.875. Fine-tuning was achieved via additional training on external data from one held-out subject out of five with a learning rate of 1e-6 for three epochs, reporting the average result for validation. Since training and evaluation were on 1.25-mm HCP data, model predictions were 5th order spline interpolated to native dMRI resolution for external OOD dataset validation. The UNet (Kerfoot et al., 2019) was implemented using Pytorch and Monai with default model hyperparameters. UNet training, fine-tuning, and validation were conducted identically, except training continued for 20 epochs until training loss stabilized.

The SWIN model achieves the lowest MAE for almost all DTI metrics in WM and GM in the HCP and TBI validation datasets, even without any fine-tuning (Table 1). For the SPIN dataset, the SWIN model also outperforms all other models, except for FA in GM where BM4D did best. In AHA data, excepting MD and RD, SWIN achieves the lowest MAE. While UNet exceeds BM4D, MPPCA, and P2S in most settings, its performance still lags behind the SWIN model, especially after fine-tuning. Patch2Self does worse than other denoising methods, especially for V1 estimation. We include quantitative comparison with DDM² with one example subject from each dataset (Supplementary Table S6). Qualitative comparisons reinforce these quantitative results. SWIN captures finer features in WM and GM microstructure without the excessive smoothing of BM4D, TPCA, and MPPCA (Figs. 2 and 3; Supplementary Figs. S1 and S2). We examine the noise distribution of each denoising algorithm in the lateral ventricles of one AHA subject (Fig. 4B). Unlike other denoising algorithms, SWIN is able to transform the original data from a heavy-tailed Rician-like distribution into a more Gaussian-like distribution with significantly smaller variance and skew but greater kurtosis.

SWIN achieves the lowest CoV between test and retest datasets for all metrics in GM and AD and MD in WM using the 6-direction subsampled shell, as well as MD and RD in both GM and WM for the 90-direction fully sampled shell (Table 2). For 90 directions, applying no denoising yields the lowest CoV for AD, FA, and MD in WM and AD in GM, while MPPCA does best for FA in GM. We include region-wise DTI CoV values for the six-direction subsampled shell in GM (Supplementary Tables S9, S10, S11, and S12) and WM (Supplementary Tables S13, S14, S15, and S16).

For repeatability of intracellular volume fraction (ICVF), fiber orientation dispersion index (ODI), and free water fraction (ISOVF), where all acquired data were used, SWIN outperforms all other denoising methods except for ODI estimation in GM, where P2S is slightly better. In particular, SWIN excels at ICVF repeatability, achieving close to 50% lower CoV than the next best method on average and achieves dramatically lower CoV on regional GM and WM measurements (Fig. 4A). SWIN also achieves considerably better test–retest reliability than the other denoising approaches for ISOVF as well, especially in global and regional GM measurements (Fig. 4A). We include region-wise NODDI CoV values in GM (Supplementary Tables S17, S18, and S19) and WM (Supplementary Tables S20, S21, and S22).

SWIN generates the most accurate FA, MD, and RD SCNs in WM and FA and RD SCNs in GM. Patch2Self achieves the lowest error for AD and MD SCN in GM, while MPPCA has the most accurate AD SCN in WM. MPPCA yields the lowest SCN repeatability error in WM for AD, FA, and RD, while Patch2Self has the most repeatable SCN in WM for ODI and GM for AD and FA. SWIN has the lowest SCN repeatability error for MD, ICVF, and ISOVF in WM plus MD, RD, ICVF, ODI, and ISOVF in GM. SWIN performed remarkably better than all competing methods for ICVF and ISOVF, achieving CoV values one-sixth to one-third of those from no denoising or denoising with P2S, MPPCA, or BM4D.

SWIN achieves better denoising even in poor quality data, including the lowest CNR scan in the AHA dataset (Fig. 5). SWIN denoising with only 6 directions approaches the image quality of all 55 directions, resulting in a 9-fold speedup of scan time. Applying SWIN to all 55 directions removes much of the FA noise from the AVM and its surrounding hemorrhage. Fine-tuning leads to further improvements in image quality for both the UNet and SWIN models.

SWIN denoising with 6 directions produces DTI values consistent with those of other denoising algorithms with all 55 directions (Supplementary Fig. S3). SWIN achieves lower FA than the other methods when limited to 6 directions, barring P2S which yields unrealistically low values for WM. Hence, SWIN can reduce noise that artifactually inflates FA. Apart from FA in GM, SWIN denoising generates similar values for both 6 directions and 55 directions, indicating that microstructural metrics remain consistent, even as angular resolution is increased. SWIN also consistently produces lower FA, MD, RD, and AD values than other denoising methods in the perilesional space across all three sessions in AHA subjects (Supplementary Fig. S4). MPPCA, BM4D, and no denoising (RAW) tend to follow the same values.

Although SWIN was not trained for super-resolution, it can be used to resample a dMRI dataset to 1.25 mm resolution (Fig. 6). SWIN captures fine anatomic details in posterior periventricular WM, avoiding the excessive blurring of BM4D and MPPCA, while achieving significantly lower MAE in all DTI metrics (Supplementary Table S5).

To the best of our knowledge, this is the first supervised dMRI denoising method that can be applied, without modification, to dMRI from widely varying scanners, patient populations, and acquisition parameters. The novelty of our approach and our large improvement over the previous SOTA lie in the way we train our model. Instead of hard-coding diffusion-encoding directions, we trained our model to learn a denoising function that performs well in a wide range of q-space which increases its robustness to changes in SNR and contrast. By doing so, we also increased the size of our training set by a factor of 288 and made fine-tuning on even one example dMRI scan feasible. We believe that this training protocol should be feasible for other noise-limited MRI contrasts, such as fMRI, as well as other imaging modalities, such as CT. We validated our approach on held-out scans from the HCP dataset and three external OOD datasets: children with neurodevelopmental disorders, adults with TBI, and both children and adolescents with intracranial hemorrhage before and after resection of their vascular malformations. We showed that SWIN produces more accurate DTI metrics and spherical harmonic coefficients than other denoising methods with the minimum amount of data required for a unique fit. We also demonstrated the superior test–retest reliability of SWIN for both DTI and NODDI metrics as well as SCNs derived from those metrics. Finally, we showed that SWIN denoising can track brain microstructural changes with greater accuracy than other denoising models.

With SWIN denoising, drastic scan time speedup is possible. Most dMRI protocols require at least 5 b = 0 s/mm² and 30 b = 1000 s/mm² acquisitions to obtain accurate DTI metrics, taking at least 10 minutes for high spatial resolution on typical clinical MR scanners available in the community (Jones, 2004). With SWIN, accurate high-resolution DTI is achievable in under 2 minutes. Indeed, on the scanners which acquired the data, a minimal 6-direction acquisition takes 90 seconds for HCP, 100 seconds for TBI or AHA, and only 20 seconds for SPIN. This speedup enables high-resolution dMRI in uncooperative populations by mitigating motion artifacts.

SWIN showed high repeatability in all cases, especially for NODDI metrics of tissue intracellular volume fraction and free water fraction, often achieving better than 50% lower CoV than the next best method. This could reflect SWIN’s approximately Gaussian output distribution. Unlike other denoising methods, SWIN removes the heavy tail noise inherent in dMRI data. Even BM4D, which is designed to correct for Rician noise, fails at this task. SWIN and, to a lesser extent, UNet succeed partly because, unlike other denoising models, they leverage T1 anatomical information for tissue segmentation, differentiating GM from WM from CSF. Furthermore, both models are trained to reduce mean-squared error, which heavily penalizes outliers and thus reduces the heavy tail.

There were cases where SWIN was not best. For instance, SWIN underperforms MPPCA in WM for 6th order spherical harmonic fitting, even in HCP data, possibly because SWIN denoises one direction at a time, whereas MPPCA collectively denoises all directions (Supplementary Table S7). By processing each dMRI volume separately, we consider the full brain volume and avoid artifact propagation across volumes but do not utilize correlation across volumes, such as the transformer patch-based approach in Karimi & Gholipour (2022). Due to GPU memory constraints, a trade-off exists between utilizing spatial and angular correlations. Data compression techniques, such as quantized variational autoencoders, might overcome this obstacle (Duan et al., 2023).

Our results were achieved with simple data augmentations, although using more extensive simulations, such as those in Muckley et al. (2021), could lead to greater generalizability. In addition, all inputs were resampled to 1.25 mm resolution, but with further data augmentation or resolution-independent architectures (Dehghani et al., 2023), it may be possible to perform denoising on native resolution if super-resolution is not desired. Our results require a T1 anatomical volume to be acquired in conjunction with the dMRI data, which we found improved denoising fidelity. While most datasets which have dMRI data also collect T1, this might not be always the case. Further work is needed to determine the importance of adding anatomical information to denoising performance. Finally, our denoising method is designed to be performed after correction for patient movement, susceptibility-induced distortions, and eddy-current correction via the Eddy command as well as coregistration with a T1 scan. However, it may be fruitful to pursue 2D denoising which cleans raw data slice-by-slice to improve Eddy itself and yield better downstream results.

We trained a simple UNet which, while performing better than self-supervised methods, trailed the SWIN model. This could reflect the ability of SWIN transformers to capture long-range dependencies better than UNets, especially when trained with enough data (Dosovitskiy et al., 2020). Since the same hyperparameters were used for both SWIN and UNet, it may be that the hyperparameters were simply more optimal for the SWIN model. Further hyperparameter optimization is required to determine the best architecture. In principle, we believe that most advanced neural network architectures, with sufficient complexity, would perform adequately.

Some of the generalizability of our model might be attributed to grokking, a phenomenon that occurs when a neural network with good training loss but poor generalization will, upon further training, transition to perfect generalization (Power et al., 2022). We anecdotally observed grokking to occur after the sixth epoch of training and suspect this is due to using weight decay and AdamW adaptive stochastic gradient descent training on a sufficiently large dataset with data augmentation. A better understanding of what leads to grokking could be instrumental in designing generative AI models that generalize well at scale for healthcare.

Finally, fine-tuning on even one subject consistently led to improved denoising performance for DTI, but results were mixed for higher-order spherical harmonics. This could be because fine-tuning reduces model bias, but increases variance which can accumulate error over the 15 or 28 directions used to compute the 4th or 6th order spherical harmonics, respectively, compared with the six directions used in DTI estimation. Moreover, there were no significant benefits from fine-tuning on more than one subject. Recent work has demonstrated that large language models can effectively learn from a single example, and our model, while significantly smaller, might also exhibit similar behavior via successful domain adaptation with limited fine-tuning data. To the best of our knowledge, fine-tuning has never been reported before in the context of dMRI denoising and merits further investigation.

We make the source code and model publicly available at https://github.com/ucsfncl/dmri-swin. Of the four imaging datasets used in this report, the Human Connectome Project (Essen et al., 2013) is a widely used publicly available dataset, two others have been previously utilized for clinical research (Kuceyeski et al., 2019; Mark et al., 2023) but not for AI research, and the fourth has not been previously published.

The HCP Test and Retest datasets were anonymized and not collected by the investigators; this work is classified as nonhuman research. All participants for the AHA, TBI, and SPIN datasets provided written informed consent by the participants or their legal guardians, as well as assent in the case of minors, under a research protocol approved by the institutional review board at our institution.

A.S. performed the imaging and statistical analysis and writing of the manuscript. X.P. assisted in the data analysis. H.C. and L.T.C. assisted in the data curation. A.S. and P.M. designed the experiments and interpreted the findings. All authors contributed to the article and approved the submitted version.

The study was funded by a National Institute of Health R01 grant, award number MH116950, and by the American Heart Association Bugher Foundation.

The authors declare that the research was conducted in the absence of any commercial or financial relationships that could be construed as a potential conflict of interest.

HCP data were provided by the Human Connectome Project, WU-Minn Consortium (Principal Investigators: David Van Essen and Kamil Ugurbil; U54 MH091657) funded by the 16 NIH Institutes and Centers that support the NIH Blueprint for Neuroscience Research; and by the McDonnell Center for Systems Neuroscience at Washington University. TBI data were acquired as part of a research project funded by NIH R01NS060886 (Principal Investigator: Pratik Mukherjee). SPIN data were acquired as part of a research project funded by NIH R01 MH116950 (Principal Investigators: Pratik Mukherjee and Elysa J. Marco). AHA data were acquired with funding from the American Heart Association (AHA) Bugher Foundation (Principal Investigators: Heather Fullerton, Christine Fox, Helen Kim, and Pratik Mukherjee).

Supplementary material for this article is available with the online version here: https://doi.org/10.1162/imag_a_00193.