SBI Meets Tractography: A new approach for Bayesian inference in diffusion MRI models

SBI Meets Tractography: A new approach for Bayesian inference in diffusion MRI models J.P. Manzano-Patrón1*, Manuel Gloecker2*, Cornelius Schröder2, Jakob H. Macke2,3*, Stamatios N. Sotiropoulos1* 1Sir Peter Mansfield Imaging Centre, School of Medicine, University of Nottingham, UK. 2MPI for Intelligent Systems, University of Tübingen, Germany. 3Hertie Institute for AI in Brain Health, Tübingen, Germany. *Equal Contribution INTRODUCTION: Simulation-Based Inference (SBI) has emerged as a powerful framework for Bayesian inference. Neural networks are trained on in-silico simulations from a forward model, and subsequently applied to unseen data for rapidly estimating posterior distributions of the model parameters given the data (amortised inference, see Fig.1) [1]. We recently presented how SBI can be used for parametric spherical deconvolution from diffusion MRI (dMRI), for mapping uncertainty of fibre orientation estimates and performing probabilistic tractography [2]. However, inherent challenges remain: i) trained networks can be tied to the acquisition scheme and/or noise level/features used for training, ii) choosing the right model (including e.g. model complexity, number of modes, noise model, priors) can typically rely on heuristics/deterministic model selection [2,3]. Here, we propose a new transformer-based simulation-based model inference (SBMI) architecture [4,5] that addresses these two challenges, by allowing a) concurrent model selection and parameter inference over multiple model classes in a single framework, and b) flexible adaption to different acquisition schemes, priors, noise level and model configurations at inference time (i.e. post-training). We show results in the context of Bayesian inference for a multi-compartment model and probabilistic tractography, but the presented principles apply to any dMRI microstructure model. METHODS: We trained SBI networks on synthetic data using the multi-shell Ball&Sticks [6] as a forward model and acquisition schemes matched in length to the UKBiobank dMRI protocol (105 volumes), as in [2], but with varying b-values (up to b=4000 s/mm2) and gradient directions. We subsequently compared three approaches (Fig. 1) in their ability to resolve fibre crossings and the respective orientation uncertainty: random-walk MCMC, our previous SBI architecture [2] trained jointly on models with N=1,2,3 stick compartments, and the proposed SBMI_transformer, which was based on the simformer architecture [5]. Evaluations were first performed using UK Biobank-like data from [7]. For SBMI_transformer, model posterior probabilities were inferred jointly with parameter posteriors, preserving uncertainty across all plausible models. We assessed mean parameter estimates, uncertainty, and tractography by propagating posterior orientation samples to build spatial distributions of white matter (WM) pathways [8]. We finally tested the ability of SBMI_transformer in adapting to different acquisition schemes, by applying the network trained on UKB-like data to HCP-like data (i.e. considerably different b-vals, b-vecs and length of data compared to the training set). RESULTS: Compared to MCMC, both SBI approaches provide high agreement on mean estimates with MCMC (Fig.2A), demonstrating the feasibility of these frameworks. However, SBMI_transformer improves over previous SBI by achieving better estimation for the third orientations and their uncertainty, sharper contrast between WM and non-WM (Fig.2A), and greater coherence in the crossing fibres areas (Fig. 2B). These improvements are directly translated into a greater correlation with MCMC in the reconstructed tractography (Fig.2C) compared to SBI (0.92 vs 0.85 respectively). Importantly, we can use the same trained network for HCP-like data (Fig.3) showcasing the capability of SBMI_transformer to amortise now among any acquisition scheme, even when the number of volumes and q-space sampling differ substantially from those at training. CONCLUSION: We introduced a transformer-based SBI architecture for Bayesian inference in diffusion MRI, that can inherently handle model selection, as well as parameter estimation and can generalise to unseen data from different-to-training acquisition schemes on inference time. Demonstrated on the Ball&Sticks model and probabilistic tractography, the principles are applicable for estimation and uncertainty mapping to any dMRI microstructure or orientation model. References: [1] Cranmer, PNAS, 2020 [2] Manzano-Patron, Med. Imag. Analysis, 2025 [3] Karimi, Imag Neuro, 2024 [4] Schröder, ICML, 2024 [5] Gloeckler, ICML, 2024 [6] Jbabdi, MRM, 2012 [7] Manzano-Patron, Imag Neuro, 2024 [8] Warrington, NeuroImage, 2020