1232

Transfer learning for non-parametric prediction of joint distributions of g-ratios and axon diameters from MRI

Gustavo Chau Loo Kung^1,2, Emmanuelle M.M. Weber², Ankita Batra³, Lijun Ni³, Michael Zeineh², Juliet Knowles³, and Jennifer A. McNab²
¹Bioengineering Department, Stanford University, Stanford, CA, United States, ²Radiology Department, Stanford University, Stanford, CA, United States, ³Neurology Department, Stanford University, Stanford, CA, United States

Synopsis

Keywords: Analysis/Processing, Microstructure, Histology, Diffusion Imaging, g-ratio, axon diameter

Motivation: Machine learning approaches are an alternative to conventional biophysical model fitting used to generate MRI microstructural maps, but the lack of paired MRI-histology data complicates end-to-end training of these models.

Goal(s): Develop a nonparametric deep learning based prediction of joint distributions of g-ratios and axon diameters from multimodal MRI data.

Approach: Histology-based synthetic MRI data was used to pretrain a conditioned normalizing flow model. Transfer learning was then performed on limited paired MRI-histology data.

Results: The joint distribution shows good visual agreement with actual samples and the distances between the marginal probabilities and their respective samples exhibit a Jensen-Shannon distance smaller than 0.22.

Impact: We present an optimized model to obtain non-parametric joint distributions of g-ratios and axon diameters from multimodal MRI from limited experimental data. The approach can easily be adapted to other microstructural modeling tasks.

Introduction

Different MRI contrasts such as Diffusion MRI(dMRI), Magnetization Transfer(MT), or relaxometry can be used to probe tissue microstructure non-invasively in the brain. This is usually achieved by performing a voxel-wise fit of the MRI data to a biophysical model and then extracting parameters that represent specific microstructural features. However, the employed biophysical models remain incomplete and rely on strong assumptions about the underlying tissue microstructure that may not always hold true[1]. Machine learning(ML) techniques have been proposed as a potential approach to learn the relationship between MRI and microstructural features[2-5]. However, paired MRI and histological ground-truth on the same tissue is onerous to acquire and usually limited to small samples. Here we build on a previous approach[6], which first pretrains a ML model on histologically-derived MRI data before performing transfer learning on limited experimental data. We now combine a convolutional neural network(CNN) that exploits local correlations, with a conditional normalizing flow model(cNFM) for non-parametric prediction of joint distributions of axon diameters and g-ratios enabling to more explicitly model the inherent correlations in the data. G-ratio is defined as the ratio between the inner axon diameter and the outer diameter including myelin. Both g-ratio and axon diameter have important functional implications in brain connectivity[7].

Methods

Pretraining Data: Electron microscopy(EM) images were obtained from a publicly available dataset containing a single slice of a canine spinal cord[8]. The method for generating histological-derived synthetic MRI is depicted in Figure 1. Histological images were divided into 125x125um tiles representing the regions corresponding to the MRI voxels. AxonDeepSeg[9] was used to obtain lists of g-ratios and axon diameters of each micrograph. We utilized the extracted axonal features to generate dMRI[10] and MT[11] data at varying SNR levels, aligning the scan protocol with the paired MRI/EM acquisition previously described[6,12].
Paired MRI/EM data: MRI (dMRI and MT) acquired on a 7T Bruker animal scanner and EM of segmented genu, body and splenium of the corpus callosum of nine mouse brains (4 wild-types and 5 absence seizure Scn8a^+/mutmutants[13]), acquired as previously described[6,12]. The axons and g-ratios were manually annotated[12] from EM micrographs.
Model Pretraining: Inspired by [14], we used a cNFM[15] to achieve a non-parametric prediction of axon diameter and g-ratio distributions.to achieve a non-parametric prediction of axon diameter and g-ratio distributions. The cNFM consists of a series of invertible transformation blocks that map a known basic distribution (e.g. bivariate Gaussian) to the target transformation, conditioned by a context vector (Figure 2). We chose to use two Masked Autoregressive blocks[16,17]. The context vector is generated using a CNN-based architecture which encodes patches of 3x3 neighboring voxels extracted from MRI data. Training was performed in 80% of the simulated data (~2000 voxels) using Adam[18] optimizer and a negative log-likelihood loss. Multiple models were trained[19] and the optimal hyperparameters were chosen based on the Jensen-Shannon distance(JSD) between the predicted distribution and the histogram of the ground-truth samples on the validation set.
Transfer Learning: After pretraining, we fine-tuned the network using the mice data split into training (2 wild-types, 3 mutants), validation (1 wild-type, 1 mutant) and testing (1 wild-type, 1 mutant).

Results and Discussion

Figure 3a-b shows good agreement between the predicted joint distribution generated from the network and the kernel density estimation plots for one selected sample of the synthetic test set. Marginal distributions for axon diameters and g-ratios were obtained by integrating the predicted joint distribution along the corresponding axis. Comparisons of these predictions with the ground-truth histograms are shown in Figures 3c-d, along with the calculated JSD distance. On average, we obtained JSDs of 0.172 and 0.173 for axon diameter and g-ratio distributions, respectively. Figure 4 displays the joint distribution comparisons for the body of the mice in the experimental test set. The comparisons of marginal distributions for the different brain regions of the test set are shown in Figure 5. We observe a good agreement between the joint distributions and the histograms and, although the JSDs are higher, we obtained the relatively low values of 0.213 and 0.208 for axon diameter and g-ratio distributions, respectively.

Conclusion

We have presented an optimized approach to combine histologically-derived synthetic MRI data, transfer learning and a cNFM for nonparametric prediction of joint distributions of axon diameter and g-ratio from MRI data. The nonparametric approach reduces assumptions and constitutes a step towards a data-driven approach. Our experiments show the good performance that can be obtained on a limited amount of data. Future work will focus on a wider range of synthetic data and further experimental testing.

Acknowledgements

This work was partially supported by the Stanford Wu Tsai Neurosciences Institute and Stanford Enhancing Diversity in Graduate Education (EDGE) fellowship.

References

[1] D.C. Alexander DC, et al. "Imaging brain microstructure with diffusion MRI: practicality and applications." (2019) NMR Biomed. 32(4):e3841.

[2] Liang, Zifei, et al. "Virtual mouse brain histology from multi-contrast MRI via deep learning." Elife 11 (2022): e72331.

[3] Hédouin, Renaud, et al. "Decoding the microstructural properties of white matter using realistic models." NeuroImage 237 (2021): 118138.

[4] E. Weber, C. Leuze, D. Barbosa, G. Chau, K. Grill-Spector and J. McNab, “Prediction of human brain microstructure from raw diffusion MRI data using deep neural networks,” 2022 ISMRM Workshop on Ultra-High Field MR.

[5] E. Weber, C. Leuze, D. Barbosa, G. Chau, K. Grill-Spector and J. McNab, "Learning the relationship between human brain tissue microstructure and diffusion MRI data" (2021) in: Proceedings of the 30th Annual Meeting of ISMRM

[6] G. Chau Loo Kung*, E. M. M. Weber*, J. K. Knowles, A. Batra, L. Ni and J. A. McNab, "A transfer learning approach to predict Axon Diameter and g-ratio distributions from MRI Data" (2023) in: Proceedings of the 31st Annual Meeting of ISMRM

[7] Schmidt, Helmut, and Thomas R. Knösche. "Action potential propagation and synchronisation in myelinated axons." PLoS computational biology 15.10 (2019): e1007004.

[8] Vuong, M.-T., Duval, T., Cohen-Adad, J., Stikov, N., 2017. On the Precision of Myelin Imaging: Characterizing Ex Vivo Dog Spinal Cord with MRI and Histology, in: Proceedings of the 25th Annual Meeting of ISMRM. p. 3760.

[9] Zaimi, Aldo, et al. "AxonDeepSeg: automatic axon and myelin segmentation from microscopy data using convolutional neural networks." Scientific reports 8.1 (2018): 1-11.

[10] Ianuş, Andrada, Daniel C. Alexander, and Ivana Drobnjak. "Microstructure imaging sequence simulation toolbox." International workshop on simulation and synthesis in medical imaging. Springer, Cham, 2016.

[11] Karakuzu, Agah, et al. "qMRLab: Quantitative MRI analysis, under one umbrella." Journal of Open Source Software 5.53 (2020): 2343.

[12] G. Chau Loo Kung* and Knowles, J* et al. "Quantitative MRI reveals widespread, network-specific myelination change during generalized epilepsy progression." NeuroImage 280 (2023): 120312.

[13] Makinson, Christopher D., et al. "Regulation of thalamic and cortical network synchrony by Scn8a." Neuron 93.5 (2017): 1165-1179.

[14] M. Jallais and M. Palombo, “μGuide: a framework for microstructure imaging via generalized uncertainty-driven inference using deep learning,” in Proceedings of the 31st Annual Meeting of ISMRM, 2023

[15] Eike Cramer, Dirk Witthaut, Alexander Mitsos, Manuel Dahmen, "Multivariate probabilistic forecasting of intraday electricity prices using normalizing flows." Applied Energy 346 (2023), 121370

[16] Papamakarios, George, Theo Pavlakou, and Iain Murray. "Masked autoregressive flow for density estimation." Advances in neural information processing systems 30 (2017).

[17] Stimper, Vincent, et al. "normflows: A PyTorch Package for Normalizing Flows." arXiv preprint arXiv:2302.12014 (2023).

[18] Kingma, Diederik P., and Jimmy Ba. "Adam: A method for stochastic optimization." arXiv preprint arXiv:1412.6980 (2014).

[19] L. Biewald, "Experiment Tracking with Weights and Biases" (2020)

Figures

Figure 1. Schematic depicting the steps to generate histologically-derived synthetic MRI data. Microstructural features were extracted from electron microscopy images of a canine and rat spinal cords using AxonDeepSeg. For each voxel, diffusion signals were simulated using embedded cylinders with the MISST library and MT signals were simulated with the qMRLab package. The simulated protocol mimics the experimental one, consisting of 2 shells of diffusion at b=3000 and 6000 and 7 offsets and two equivalent flip angles for the MT.

Figure 2. Architecture of the network to predict joint axon diameter and g-ratio distributions. Neighborhoods of 3x3 voxels and 86 channels corresponding to the number of MRI contrast acquisitions are fed to a CNN with a kernel size of 2x2 to generate 8-dimensional context vectors. The latters are then used to condition a normalizing flow model which learns an invertible transformation from a bivariate Gaussian to the joint distribution of axon diameters and g-ratios.

Figure 3. Performance of the pretrained model on a selected sample from the test set of the histology-derived synthetic MRI data. (a) and (b) shows a kernel density estimate plot of ground-truth samples and the joint distribution predicted by the network, respectively. The marginal distributions obtained from integrating the joint distribution predicted by the network versus the actual ground-truth sample histogram are shown in (c) for axon-diameter and (d) for g-ratio.

Figure 4. Comparison of the joint distribution predicted by the network after transfer learning with a kernel density estimate plot of the samples on the wild-type mouse of the test set (b and c) and the mutant mouse (d and e) for the body of the corpus callosum. A reference structural image with the body highlighted is shown in (a).

Figure 5. Performance of the fine-tuned network on the test set of the mice data. The top six plots show the comparison of marginal distributions obtained from integrating the joint distribution predicted by the network versus the actual ground-truth sample histogram for axon diameters while the six bottom ones show the comparison for the g-ratios. For each plot, the Jensen-Shannon distance (JSD) is included. A reference structural image is shown on the right side with the genu, body and splenium of the corpus callosum highlighted.

Proc. Intl. Soc. Mag. Reson. Med. 32 (2024)

1232

DOI: https://doi.org/10.58530/2024/1232