2818
Enhancing Reliability in Model-based DL Reconstruction: A Systematic Study of MC Dropout for Uncertainty Quantification1Institute of Pure and Applied Sciences, University of Tsukuba, Tsukuba, Japan
Synopsis
Keywords: Machine Learning/Artificial Intelligence, Image Reconstruction
Motivation: Monte Carlo (MC) Dropout, a powerful uncertainty quantification (UQ) method for deep learning-based reconstruction, can impact reconstruction performance. Finding ways to enhance reliability assessment without compromising performance is essential.
Goal(s): This study aims to provide advisory information on how to incorporate MC Dropout into a model-based unrolled neural network, and to evaluate the reliability of UQ.
Approach: Different architectures with varying dropout rates are used to assess image quality. Images with visible structural aberrations and artificial perturbation are tested.
Results: Findings indicate that appropriate MC Dropout configurations improve reconstruction quality, and UQ maps effectively identify structural anomalies in images.
Impact: This research enhances the reliability of DL reconstructions by systematically investigating MC Dropout’s impact on reconstruction performance, particularly in scenarios lacking ground-truth references. The findings guide the incorporation of uncertainty quantification techniques, improving the overall quality of medical imaging applications.
Introduction
Deep learning (DL) has become increasingly popular for accelerated MR reconstruction, yet its black-box nature makes result interpretation challenging1,2. Evaluating model performance without ground-truth (GT) references in practical scenarios poses difficulties in ensuring image fidelity2. Reliability is paramount in diagnostic tasks. Uncertainty quantification (UQ) is instrumental in addressing the non-deterministic and non-transparent nature of DL reconstruction, allowing for case-by-case, pixel-wise assessment of network uncertainty3,4. Monte Carlo (MC) Dropout is commonly used for UQ. However, UQ itself can impact performance, therefore finding an appropriate approach to enhance reliability evaluation while avoiding performance degradation is essential4,5. Nevertheless, the different ways to incorporation of MC Dropout into network architecture has not been extensively studied. In this study, we systematically examined the influence of MC Dropout on a model-based unrolled neural network with conjugate gradient optimization. We provided a detailed evaluation of UQ reliability and dependability using MC Dropout, focusing on two different paradigms: supervised learning, where no GT reference is available in testing, and self-supervised learning, where no GT reference is available in both training and testing.Methods
Network ArchitectureFor image reconstruction, we considered a model adapted from the Self-Supervision via Data Undersampling (SSDU) architecture6. This is a model-based, unrolled iterative neural network with a regularizer (ResNet) and a conjugate gradient-optimized data consistency (DC) unit in each iteration (Figure 1-A). The original SSDU splits the acquired (uniformly undersampled) k-space data into two disjoint sets, with one used for DC, and the other one used to define the loss function in k-space6. In the supervised setting, we used the uniformly undersampled k-space for training, and the fully sampled k-space for the loss function (Figure 1-C).
Uncertainty Quantification
MC Dropout is a technique where certain neurons are randomly deactivated (“dropped out”) during both training and inference in neural networks to estimate epistemic uncertainty by obtaining multiple predictions7. At each training step, each neuron has some probability p of being dropped out (dropout rate). In order to investigate the effects of different dropout rates, p=0.1 and p=0.4 was tested. The ResNet module was modified to implement four different dropout patterns (Figure 1-B), with P0 being the baseline. Each inference was repeated for T=10 times. The reconstruction was calculated from the mean of all 10 repetitions and a UQ map was generated from the variance.
Training Conditions and Data Preparation
All experiments were trained under the same conditions, as shown in Table 1. Training, validation, and testing were done using coronal PD (knee) and axial FLAIR (brain) data from the fastMRI dataset8. The fastMRI+ database was used to isolate healthy subjects, and only cases without significant structural aberrations were used in training9. Clinical pathology annotations were also generated for one of the experiments.
Results and Discussions
Training and Validation LossTraining and validation loss for each model is shown in Figure 2. P1 demonstrated a higher loss across all training scenarios. This is likely due to the fact that in P1, MC Dropout was done on the convolution layer right before the DC layer, which introduces performance degradation.
Reconstruction Performance
Figure 3 showcases selected reconstruction results and Structural Similarity Index Measure (SSIM) evaluation for each case. P2 and P3, benefiting from a more appropriate MC Dropout configuration, outperformed both the baseline and P1 in terms of reconstruction quality. Similar trends can be observed with both supervised and self-supervised reconstruction. MC Dropout can help prevent overfitting in deep neural networks due to its inherent stochasticity and regularization effect, discouraging the network from memorizing noise in the training data and promoting more robust generalization.
UQ Variance Maps
Figure 4 shows the reconstruction of brain images with visible structural aberrations. Notably, the UQ map was able to qualitatively identify these anomalies, arguably better than the error map. Figure 5 demonstrates the reconstruction of images with artificially introduced perturbation. The added letters are fine details that are also out-of-distribution from the training data, which present a reasonable challenge for the network. Consistent with the previous case, the UQ maps proved effective in identifying the perturbed regions, with P2 and P3 showing more discernible patterns compared to P1. UQ variance maps in conjunction with AI-reconstructed MRI images provide critical information about the model's confidence in each pixel, enabling the identification of potential inaccuracies.
Conclusion
Systematic investigation revealed advisory information on the appropriate ways to incorporate MC Dropout into a model-based unrolled neural network. The findings of this study also underscore the significance of UQ techniques in enhancing model performance, as well as the reliability and interpretability of DL-based image reconstructions, especially when GT references are unavailable.Acknowledgements
No acknowledgement found.References
1. Ahishakiye E, Van Gijzen MB, Tumwiine J, Wario R, Obungoloch J. A survey on deep learning in medical image reconstruction. Intelligent Medicine. 2021;1(3):118-127. doi:10.1016/j.imed.2021.03.003
2. Antun V, Renna F, Poon C, Adcock B, Hansen AC. On instabilities of deep learning in image reconstruction and the potential costs of AI. Proceedings of the National Academy of Sciences. 2020;117(48):30088-30095. doi:10.1073/pnas.1907377117
3. Zou K, Chen Z, Yuan X, Shen X, Wang M, Fu H. A Review of Uncertainty Estimation and its Application in Medical Imaging. May 2023. http://arxiv.org/abs/2302.08119.
4. Shang R, O’Brien MA, Wang F, Situ G, Luke GP. Approximating the uncertainty of deep learning reconstruction predictions in single-pixel imaging. Commun Eng. 2023;2(1):1-12. doi:10.1038/s44172-023-00103-1
5. Schlemper J, Castro DC, Bai W, et al. Bayesian Deep Learning for Accelerated MR Image Reconstruction. In: Knoll F, Maier A, Rueckert D, eds. Machine Learning for Medical Image Reconstruction. Lecture Notes in Computer Science. Cham: Springer International Publishing; 2018:64-71. doi:10.1007/978-3-030-00129-2_8
6. Yaman B, Hosseini SAH, Moeller S, Ellermann J, Uğurbil K, Akçakaya M. Self-supervised learning of physics-guided reconstruction neural networks without fully sampled reference data. Magn Reson Med. 2020;84(6):3172-3191. doi:10.1002/mrm.28378
7. Gal Y, Ghahramani Z. Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning. October 2016. http://arxiv.org/abs/1506.02142.
8. Zbontar J, Knoll F, Sriram A, et al. fastMRI: An Open Dataset and Benchmarks for Accelerated MRI. December 2019. doi:10.48550/arXiv.1811.08839
9. Zhao R, Yaman B, Zhang Y, et al. fastMRI+: Clinical Pathology Annotations for Knee and Brain Fully Sampled Multi-Coil MRI Data. September 2021. doi:10.48550/arXiv.2109.03812
Figures
Figure 2. Training loss and validation loss for each case.
Table 1. Detailed network configurations, training conditions, and dataset information.
