2223
Uncertainty-Aware Anatomical Brain Parcellation using Diffusion MRI1University of Electronic Science and Technology of China, Chengdu, China, 2Nanjing University of Science and Technology, Nanjing, China, 3East China University of Science and Technology, Shanghai, China, 4The Second Affiliated Hospital of Zhejiang University School of Medicine, Hangzhou, China, 5The First Affiliated Hospital, Sun Yat-sen University, Guangzhou, China, 6Harvard Medical School, Boston, MA, United States
Synopsis
Keywords: Analysis/Processing, Machine Learning/Artificial Intelligence, Diffusion MRI
Motivation: While anatomical brain parcellation has long been performed using anatomical MRI and atlas-based approaches, deep learning methods together with diffusion MRI techniques can improve parcellation performance and interpretation of prediction uncertainty.
Goal(s): Our goal is to design an uncertainty-aware deep learning network to utilize multiple diffusion MRI parameters for accurate brain parcellation while enabling voxel-level uncertainty estimation.
Approach: We include five evidential deep learning subnetworks and perform an evidence-based ensemble for parcellation prediction and uncertainty estimation.
Results: The results demonstrate our method’s superior parcellation performance over several state-of-the-art methods, its promising results in unseen patient scans, and potential applications in brain abnormality detection.
Impact: The proposed approach enables improved accuracy in brain parcellation from diffusion MRI, facilitating the understanding of the human brain in health and disease. It may also serve as an effective tool for brain abnormality detection, fostering inquiries into uncertainty-quantified diagnostics.
Introduction
Anatomical brain parcellation is a vital step in neuroimaging analysis. While it is conventionally derived from atlas-based tools such as FreeSurfer1, recent studies have shown the advantages of deep learning for improved parcellation accuracy and efficiency2-4. Diffusion MRI (dMRI)5 is a specialized MRI technique that can characterize brain microstructure using dMRI-derived parameters such as fractional anisotropy (FA), mean diffusivity (MD), and diffusion tensor eigenvalues (E1, E2, E3)6. Ensemble deep learning7,8 is an advanced technique that can improve brain parcellation by utilizing complementary features derived from multiple dMRI parameters. Uncertainty-aware learning9-13 is another advanced technique that explicitly models the confidence or uncertainty of the prediction to improve parcellation and enable abnormality detection.This study presents an uncertainty-aware deep learning method for accurate and reliable brain parcellation from dMRI data. We propose a novel evidence-based ensemble network that enables robust uncertainty estimation of the predicted parcellation at each voxel. Our experiments show not only improved parcellation accuracy but also the ability to detect abnormal brain regions.
Methods
Figure-1 shows our network architecture that comprises five subnetworks, where each learns a FreeSurfer parcellation from an individual dMRI parameter (i.e., FA, MD, E1, E2, E3, respectively) using a FastSurferCNN2. Rather than using the classifier module in FastSurferCNN to predict a region label for each voxel, we employ evidential deep learning11,14 to produce an evidence-based uncertainty map from each subnetwork and then ensemble the outputs from all subnetworks for the final parcellation prediction and uncertainty estimation.(1) Evidential deep learning is used to estimate prediction uncertainty derived from each subnetwork. Based on the subjective logic theory11,15, each subnetwork learns the evidence $$$e_{i, j, k}^n$$$ that assigns the voxel at $$$(i,j,k)$$$ to class $$$n$$$. By parameterizing a Dirichlet distribution with $$$\alpha_{i,j,k}^n=e_{i,j,k}^n+1$$$, the evidence-based uncertainty $$$u_{i, j, k}$$$ can be defined as
$$u_{i, j, k}=\frac{N}{S}, S=\sum_{n=1}^N \alpha_{i,j,k}^n$$
where $$$S$$$ is the Dirichlet strength, and $$$N$$$ is the total number of parcellated regions.
The loss function of our evidential learning is:
$$L=L_{Dice}+L_{CE}+\lambda_{rce} L_{rce}+\lambda_{kl} L_{KL}$$
Here, $$$L_{Dice}$$$ and $$$L_{CE}$$$ are the standard Dice and cross-entropy losses, respectively, to ensure stable parcellation of the overall region and the boundary area, as applied in FastSurferCNN2. $$$L_{rce}$$$ and $$$L_{KL}$$$ are newly added, where $$$L_{rce}$$$ is a revised cross-entropy loss that regulates the evidential learning process by Bayes risk, and $$$L_{KL}$$$ is the KL divergence to ensure incorrect labels produce less evidence, defined as follows:
$$L_{rce}=\int\left[\sum_{n=1}^N-y_m^n \log \left(p_m^n\right)\right] \frac{1}{B\left(\boldsymbol{\alpha}_m\right)} \prod_{n=1}^N p_m^{n \alpha_m^n-1} d \mathbf{p}_m=\sum_{n=1}^N y_m^n\left(\psi\left(\boldsymbol{S}_m\right)-\psi\left(\boldsymbol{\alpha}_m^n\right)\right)$$
$$L_{KL}=\log \left(\frac{\Gamma\left(\sum_{n=1}^N \widetilde{\boldsymbol{\alpha}}_m^n\right)}{\Gamma(N)\sum_{n=1}^N\Gamma\left(\widetilde{\alpha}_m^n\right)}\right)+\sum_{n=1}^N\left(\widetilde{\boldsymbol{\alpha}}_m^n-1\right)\left[\psi\left(\widetilde{\boldsymbol{\alpha}}_m^n\right)-\psi\left(\sum_{n=1}^N \widetilde{\boldsymbol{\alpha}}_m^n\right)\right]$$
where $$$\Gamma(\,)$$$ is the gamma function, $$$\psi(\,)$$$ is the digamma function, $$$B(\alpha)$$$ is the multinomial beta function, and $$$\widetilde{\boldsymbol{\alpha}}_m^n=y_m^n+\left(1-y_m^n\right) \odot \boldsymbol{\alpha}_m^n$$$.
In our experiment, for the overall loss, the regulatory factors $$$\lambda_{kl}$$$ and $$$\lambda_{rce}$$$ are set to 0.5 and 0.4, respectively.
(2) Evidence-based ensembling is performed based on the outputs of the five subnetworks. For each voxel, the subnetwork prediction with the minimum evidence-based uncertainty is adopted as the final prediction, and the entropy of the averaged evidence is used for the final uncertainty estimation.
Experiments
We used high-quality dMRI scans from 140 young healthy adults in Human Connectome Project (HCP)16,17 for model training, validation and testing (100, 20 and 20 subjects, respectively). The provided FreeSurfer parcellation in the dMRI space is utilized as the ground truth. Experimental evaluation includes a comparison with the SOTA methods2,18-21 and an ablation study.We also used dMRI scans from 11 patients diagnosed with cerebral small vessel disease (CSVD) to evaluate the generalization of our method. These dMRI datasets include white matter hyperintensity (WMH) regions that are not seen in the training data.
Results & Discussion
Table-1 demonstrates that our network outperformed the SOTA methods. Figure-2 illustrates that our method generates a visually smoother segmentation that is more consistent with the structure boundary appearing on the FA image. Table-2 shows that our evidence-based ensemble network yielded better parcellation results than using other ensemble criteria or backbone networks. Figure-3 demonstrates that our network is capable of producing reasonable parcellation for unseen patient scans, and the output uncertainty heatmap holds great potential for lesion detection.Conclusion
Our uncertainty-aware deep learning network delivers accurate brain parcellation by utilizing multiple dMRI parameters and a novel evidence-based ensemble learning approach. The method effectively addresses challenges posed by traditional methods, outperforms state-of-the-art models, and generates anatomically reasonable brain parcellation. Notably, its ability to produce reliable parcellations for unseen patient scans and yield uncertainty heatmaps demonstrates its potential use in lesion detection, further demonstrating its value within a real-world medical context.Acknowledgements
This work is in part supported by the National Natural Science Foundation of China (No. 62371107) and the National Institutes of Health (R01MH108574, P41EB015902, R01MH125860, R01MH119222, R01MH132610, R01NS125781).References
1. Fischl B. FreeSurfer. Neuroimage 62, 774–781 (2012).
2. Henschel L. et al. FastSurfer - A fast and accurate deep learning based neuroimaging pipeline. Neuroimage 219, 117012 (2020).
3. Billot B. et al. SynthSeg: Segmentation of brain MRI scans of any contrast and resolution without retraining. Med. Image Anal. 86, 102789 (2023).
4. Zhang F et al. Deep learning based segmentation of brain tissue from diffusion MRI. Neuroimage 233, 117934 (2021).
5. Basser P J, Pierpaoli C. Microstructural and physiological features of tissues elucidated by quantitative-diffusion-tensor MRI. J. Magn. Reson. B 111, 209–219 (1996).
6. Mori S, Zhang J. Principles of diffusion tensor imaging and its applications to basic neuroscience research. Neuron 51, 527–539 (2006).
7. Fort S, Hu H, Lakshminarayanan B. Deep Ensembles: A Loss Landscape Perspective. arXiv [stat.ML] (2019).
8. Ganaie M A, Hu M, Malik A K, Tanveer M, Suganthan P N. Ensemble deep learning: A review. Eng. Appl. Artif. Intell. 115, 105151 (2022).
9. Gawlikowski J et al. A survey of uncertainty in deep neural networks. Artificial Intelligence Review 56, 1513–1589 (2023).
10. Tanno R et al. Uncertainty modelling in deep learning for safer neuroimage enhancement: Demonstration in diffusion MRI. Neuroimage 225, 117366 (2021).
11. Sensoy M, Kaplan L, Kandemir M. Evidential deep learning to quantify classification uncertainty. Adv. Neural Inf. Process. Syst. 31, (2018).
12. Roy A G, Conjeti S, Navab N, Wachinger C, Alzheimer’s Disease Neuroimaging Initiative. Bayesian QuickNAT: Model uncertainty in deep whole-brain segmentation for structure-wise quality control. Neuroimage 195, 11–22 (2019).
13. DeVries T, Taylor G W. Leveraging Uncertainty Estimates for Predicting Segmentation Quality. arXiv [cs.CV] (2018).
14. Zou K, Yuan X, Shen X, Wang M, Fu H. TBraTS: Trusted Brain Tumor Segmentation. in Medical Image Computing and Computer Assisted Intervention – MICCAI 2022 503–513 (Springer Nature Switzerland, 2022). doi:10.1007/978-3-031-16452-1_48.
15. Jøsang A. Subjective Logic: A Formalism for Reasoning Under Uncertainty. (Springer, 2016).
16. Glasser M. F. et al. The minimal preprocessing pipelines for the Human Connectome Project. Neuroimage 80, 105–124 (2013).
17. Glasser M. F. et al. A multi-modal parcellation of human cerebral cortex. Nature 536, 171–178 (2016).
18. Hatamizadeh A. et al. Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images. in Brainlesion: Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries 272–284 (Springer International Publishing, 2022). doi:10.1007/978-3-031-08999-2_22.
19. Isensee F, Jaeger P F, Kohl S A A, Petersen J, Maier-Hein K H. nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nat. Methods 18, 203–211 (2021).
20. Lakshminarayanan B, Pritzel A, Blundell C. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30 (2017).
21. Zhang F, Cho K, Seitz-Holland J, Ning L, Legarreta J, Rathi Y, Westin C, O’Donnell L, Pasternak O. DDParcel: deep learning anatomical brain parcellation from diffusion MRI. IEEE Transactions on Medical Imaging, accepted, 2023.
Figures
