Uncertainty mapping and probabilistic tractography using Simulation-based Inference in diffusion MRI: A comparison with classical Bayes

Abstract

Simulation-Based Inference (SBI) has recently emerged as a powerful framework for Bayesian inference: Neural networks are trained on simulations from a forward model, and learn to rapidly estimate posterior distributions. We here present an SBI framework for parametric spherical deconvolution of diffusion MRI data of the brain. We demonstrate its utility for estimating white matter fibre orientations, mapping uncertainty of voxel-based estimates and performing probabilistic tractography by spatially propagating fibre orientation uncertainty. We conduct an extensive comparison against established Bayesian methods based on Markov-Chain Monte-Carlo (MCMC) and find that: a) in-silico training can lead to calibrated SBI networks with accurate parameter estimates and uncertainty mapping for both single- and multi-shell diffusion MRI, b) SBI allows amortised inference of the posterior distribution of model parameters given unseen observations, which is orders of magnitude faster than MCMC, c) SBI-based tractography yields reconstructions that have a high level of agreement with their MCMC-based counterparts, equal to or higher than scan-rescan reproducibility of estimates. We further demonstrate how SBI design considerations (such as dealing with noise, defining priors and handling model selection) can affect performance, allowing us to identify optimal practices. Taken together, our results show that SBI provides a powerful alternative to classical Bayesian inference approaches for fast and accurate model estimation and uncertainty mapping in MRI.

1. Introduction

Identifying the parameters of a model that best represent a set of observations is an omnipresent task in science. For Magnetic Resonance Imaging (MRI) in particular, this is a key task, as on many occasions biophysical quantities of interest are indirectly inferred from measurements. In diffusion MRI (dMRI) of the brain, measurements of the diffusion scatter pattern of water molecules within tissues can be mapped through models to neuronal tissue microstructure and brain connectivity [1, 2, 3].

Deterministic approaches, from non-linear optimisation to machine learning methods (see [4] for a review), can be used to fit such models to data, providing point estimates of model parameters. However, these estimates are affected by different sources of uncertainty, from noise in the observations (aleatoric uncertainty), to preprocessing errors, modelling assumptions, and incomplete sampling (epistemic uncertainty). Mapping uncertainty provides a principled manner for assessing confidence in the estimates [5], separating noise and signal, identifying covariance structures in the parameters and aiding experimental design [6]. In dMRI, uncertainty in fibre orientations has been the basis for probabilistic tractography [7, 8, 9, 10], which estimates spatial distributions of white matter bundles.

Classical approaches for mapping uncertainty in dMRI include bootstrapping and Bayesian inference. Bootstrapping [11] involves iteratively fitting the same model to randomly resampled (or perturbed) versions of the original dataset. This provides an empirical distribution for each model parameter, whose width denotes the respective estimation uncertainty [12, 13, 14, 15], but can be sensitive to how resampling is performed [16, 17]. On the other hand, Bayesian approaches offer a principled framework for estimating the posterior distribution of the model parameters given the data, providing measures of confidence and covariance in the estimates. For instance, Variational Bayes [18, 19, 20], or Markov Chain Monte Carlo (MCMC) sampling [21, 22, 23] provide flexibility in modelling both signal and noise components, and make it possible to explicitly incorporate prior knowledge. These approaches have a number of limitations: 1) they require the likelihood of the model to be evaluated analytically, which is not always available or it is too expensive to compute, 2) although great speed-ups have been achieved (e.g. utilising GPUs [24] and/or improving the algorithmic design [25, 26]), they can be highly iterative in nature, which can result in substantial computational complexity and limited scalability, 3) convergence to sampling the correct posterior can be a challenge.

Simulation-based inference (SBI) [27] has been developed to tackle such limitations in Bayesian inference. The basic idea is that artificial neural networks are trained on simulations from a forward model, in a way that allows them to learn Bayesian posterior distributions. Importantly, this approach can amortise inference: after initial training, performing inference on unseen data only requires a single forward pass through the trained neural network, without the need for costly iterative sampling and likelihood evaluations [28, 29]. This amortisation property is particularly effective in scenarios where the same problem needs to be solved many times, as in the voxel-wise inference done in microstructural models of dMRI.

In this work, we demonstrate how SBI can be used for mapping orientations and uncertainty in dMRI and, for the first time, all the way up to probabilistic tractography. We comprehensively compare SBI with a classical MCMC-based and heavily-used implementation (FSL-BedpostX) [7, 30, 31], with respect to both voxel-wise fibre orientation uncertainty estimates and whole-brain white matter tract reconstructions.

1.1. Related work

In Bayesian inference, the posterior distribution p(Ω |Y_obs) of model parameters Ω given the measured data Y_obs is proportional to the product of the prior p(Ω) and the likelihood function p(Y_obs|Ω), the conditional distribution of observations given a model. For problems where the likelihood is intractable or expensive to compute [32, 33], Approximate Bayesian Computation (ABC) or Likelihood-Free Inference (LFI) approaches have been proposed [34, 35]. In ABC, the posterior is approximated by simulating large amounts of data Y_sim from a forward stochastic simulator f_Y (Ω) (synthetic likelihood), with parameters Ω randomly drawn from the defined priors p(Ω). Samples for which a distance ρ(Y_sim, Y_obs) is below a pre-defined threshold are accepted as posterior samples (otherwise rejected). SBI recasts this as a density estimation problem, where neural networks are used to learn a mapping between observations and the posterior distribution [35, 36]. This avoids inefficient sampling-rejection schemes and allows problems with high-dimensional observations to be tackled. Such networks, known as Neural Density Estimators (NDEs) [37], have enabled inference in a purely data-driven fashion on different target (conditional) densities, such as the posterior distribution of model parameters (Neural Posterior Estimation or NPE) [38, 39, 40, 41], a surrogate likelihood (Neural Likelihood Estimation) [42, 43], or variants based on denoising diffusion models [44, 45, 46, 47].

A few recent studies have explored inferring parameters from dMRI models using such approaches, demonstrating how SBI can assist in fitting models that are typically challenging to invert. Jallais et al. [48] fitted the SANDI model of grey matter [49] to direction-averaged dMRI signal, using the width of the SBI-estimated posteriors to identify regions where the model presented indeterminacies or biases. To reduce dimensionality and avoid indeterminacies, they employed rotationally invariant dMRI summary statistics to train SBI networks on synthetic data (relying on stacked Masked Autoencoders (MADE)), demonstrating accurate parameter estimation given appropriate q-space sampling. A generalisation of this work was presented in [50], proposing a toolbox applicable to a broader range of models with examples of fitting to direction-averaged dMRI signals. In Eggl and De Santis [51], the authors used the mean of the SBI posterior to demonstrate that it can provide robust point-estimates with fewer observations than deterministic non-linear fitting. In another recent work, Karimi et al. [52] presented a likelihood-free framework to obtain the full posterior of parameters in a multi-tensor model. This employed a two-step strategy: a convolutional neural network (CNN) for model selection (classifying dMRI signals into single or multi-fibre), followed by Mixture Density Networks (MDNs) with Gaussian kernels trained for inference against simulations obtained from the multi-tensor model. This work used a spherical harmonics (SH) representation of the signals improving generalisability of inference across different acquisition schemes. An extension of this work for the standard model of white matter [53] was presented in [54], demonstrating the ability to obtain fiber-specific microstructural features.

However, despite the great potential revealed for SBI from these studies, a comprehensive comparison to classical Bayesian approaches, which have been used in dMRI for more than two decades, is still missing. In particular, Bayesian frameworks for mapping orientation uncertainty and performing uncertainty-based probabilistic tractography [21, 7, 30, 24] have been collectively used in thousands of studies before [10]. Therefore, they naturally provide a reference against which SBI approaches can be calibrated, optimised and validated.

1.2. Contributions

In this study we close this gap, expanding on our previous preliminary work [55] and complementing recent work by others [48, 50, 52, 54]. We devise an SBI framework for mapping fibre orientation uncertainty from diffusion MRI data using parametric spherical deconvolution models. We subsequently use the SBI-derived voxel-wise posterior distributions to perform probabilistic tractography for the first time. Given the inherent difficulty in validating uncertainty estimates, we comprehensively evaluate the accuracy and precision of the estimates against an established MCMC-based classical inference approach, crucially using the same models and priors, as implemented in FSL-BedpostX [7, 30, 31]. We evaluate SBI both for voxel-wise and whole-brain level inference, exploring how estimated uncertainty is spatially propagated.

We find that in-silico training can lead to SBI networks that yield accurate mean parameter estimates and uncertainty for unseen observations in both single- and multi-shell dMRI acquisitions, while offering orders of magnitude computational speed-ups for inference. We also show that SBI-based probabilistic tractography for a representative set of more than 40 white matter bundles exhibits a high level of agreement with MCMC-based counterparts; and that this level of agreement is within the range of individual scan-rescan reproducibility of estimates. We further identify optimal SBI design features (such as how to account for noise, how to restrict and constraint priors and how to deal with model selection) that improve training, increase performance and similarity to MCMC estimates. Our findings augment recent studies suggesting that SBI provides new and exciting opportunities for amortised inference in MRI models.

2. Methods

2.1. NPE theory and network architecture

In classical Bayesian inference (like MCMC, see Fig. 1), given a signal and noise model, a likelihood function p(Y_obs|Ω) can be defined. Any prior knowledge of the parameters Ω can be embedded in the prior p(Ω) and the two can be used to estimate the posterior distribution of the parameters given the observed data, i.e. p(Ω|Y_obs) ~ p(Y_obs|Ω)p(Ω). This typically necessitates highly iterative processes and potentially expensive likelihood calculations. Furthermore, parameter estimation given a dataset starts from scratch for any new observations Y_obs (e.g. data for different voxels).

Figure 1. Schematic overview of Markov-Chain Monte-Carlo (MCMC) (an example of Classical Bayesian Inference) and Simulation-Based Inference (SBI).

For NPE, an artificial neural network (ANN) is trained on a set {Ω_i, Y_i} obtained from the forward simulator, to learn the latent mapping q_Φ between observations and posterior samples.

Within SBI, this problem can be addressed using Neural Posterior Estimation (NPE) [38, 39, 41]. In NPE, the mapping between observations and posterior is learned once using a training dataset and can be applied on new unseen data without any expensive calculations, a property known as amortised inference. Being able to evaluate the likelihood function is not necessary as long as we have a forward process f_Y (Ω) to generate simulated data Y. As in all Bayesian inference schemes, we can code prior knowledge into our prior distribution p(Ω) to define what regions of the parameter space are sampled for such forward predictions (or simulations). NPE approximates the target posterior using a parameterised conditional density distribution q_Φ(Ω|Y) represented by neural networks with parameters Φ. Hence, having a training set of pairs (Ω_i, Y_i), a NPE is trained to learn the mapping function q_Φ(Ω|Y) ≈ p(Ω|Y). After training, we can use q_Φ(Ω|Y) to obtain samples from the (approximate) posterior p(Ω|Y_obs) with a single forward pass through the network.

More formally, we minimise the expected Kullback-Leibler divergence between true posterior distribution p(Ω|Y) and the approximation q_Φ(Ω|Y) [38, 56]:

$$\underset{\Phi }{\min }ℒ(\Phi )={\mathbb{E}}_{y~p(Y)}\left[D_{\text{KL}}\left(p(\Omega \mid Y)\Vert q_{\Phi }(\Omega \mid Y)\right)\right].$$ (1)

This is equivalent to minimising the negative log-likelihood of the posterior approximation [41, 56]:

$$\underset{\Phi }{\min }ℒ(\Phi )\approx -\frac{1}{L}{\displaystyle \sum _{i=1}^L\log }\left(q_{\Phi }\left({\Omega }_i\mid Y_i\right)\right),$$ (2)

where L denotes the number of pair-vectors (Ω_i, Y_i) used to train the networks, where Ω ~ p(Ω) and Y_i ~ f_Y (Ω_i). During training, the parameters Φ of the network are optimised to minimise the loss in eq. 2.

Different model choices have been made to parameterise the posterior distribution q_Φ(Ω|Y) in previous applications to dMRI, such as Mixture Density Networks (MDNs) [52, 54] and Masked Autoencoders (MADE) [48, 50]. In our work, we use Normalising Flows (NFs) as the basis structure for the density estimators given their flexibility and superior performance in previous studies [57, 58, 55]. NFs can produce complex probability distributions by transforming a simple base distribution through a series of invertible transformations f_k. A key feature of NFs is their use of differentiable closed-form functions that support efficient computation of both the transformations and their Jacobian determinants, ensuring rapid density estimation and sampling. In this paper, we use Neural Spline Flows (NSF) as the basis function [59], as implemented in the SBI package [60]. NSFs comprise K rational-quadratic splines, each parameterised by a neural network.

2.2. Forward model f_Y (Ω)

In this work, we use variants of the Ball & Sticks model [7, 30] as the forward simulator f_Y (Ω). Our goal is to infer the posterior of parameters Ω from dMRI data in each voxel, with Y_obs = {Y_1,obs, …, Y_j,obs} of j = 1, …, J signals, each corresponding to a different diffusion-sensitising gradient. This model has been heavily used for performing parametric spherical deconvolution and uncertainty-based probabilistic tractography over the past 20 years. An MCMC-based implementation for classical Bayesian inference on this model has been developed by members of our team and is part of FSL [7, 31, 24] and will be used for direct comparisons with SBI.

The forward model is multi-compartment, comprising of a weighted sum of N anisotropic compartments (sticks), representing distinct white matter fibre compartments, and an isotropic compartment (ball), representing partial volume and non-axonal contributions to the signal. The voxel-wise dMRI signal attenuation A in single-shell data is described as:

$$f_Y(\Omega )\to A_j=S_j/S_0=\left(1-{\displaystyle \sum _{n=1}^N{f}_n}\right)e^{-b_{j}d}+{\displaystyle \sum _{n=1}^N{f}_n}{e}^{-b_{j}d{\left(g_j·v_n^T\right)}^2},$$ (3)

where S_j and S₀ are the signals measured with and without diffusion-sensitising gradients, respectively, with b_j value and orientation g_j (protocol-defined parameters). The unknown model parameters to estimate are the mean diffusivity d, the vector v_n describing the n^th fibre orientation, represented in spherical coordinates by [θ_n, ϕ_n] (θ is the angle from the positive z-axis and ϕ is the angle of the projection on the x-y plane from the positive x-axis) and the volume fraction of the n^th fibre f_n (with Σ f_n ≤ 1). The maximum number of fibres N per voxel is typically set to N = 3, but the appropriate N per voxel needs to be determined (model selection challenge).

The above model can be extended for multi-shell data [30] as:

$$\begin{array}{l}{f}_Y(\Omega )\to A_j=S_j/S_0={\displaystyle \sum _{n=1}^N{f}_n}\exp \left(\frac{d^2}{d_{std}}\log \left(\frac{d}{d+b_j{\left(g_j·v_n^T\right)}^2{d}_{std}^2}\right)\right)\\ +\left(1-{\displaystyle \sum _{n=1}^N{f}_n}\right)\exp \left(\frac{d^2}{d_{std}}\log \left(\frac{d}{d+b_j{d}_{std}^2}\right)\right),\end{array}$$ (4)

where d and d_std denote the mode and spread of a Gamma distribution of diffusivities that can account for non-monoexponential signal decay. In summary, the vector of model parameters when the maximum number of fibres is n, 1 ≤ n ≤ N. is given by Ω = {d, d_std, {f₁, …, f_n, θ₁, …, θ_n, ϕ₁, …, ϕ_n}}.

2.3. SBI design and implementation

We designed and implemented a number of NPE networks for the forward models described in the previous section (setting N up to 3), using the Python SBI toolbox (v.0.22) [60]. We set the depth of the NPE to K = 10 transformations. For training and inference, we used the attenuation A_j as input for the networks, obtained after data pre-processing (in case of real data), by dividing the diffusion-weighted signal by the mean of the b = 0 volumes (without loss of generality, we set S₀=1 for the simulated training sets). Normalised signals generally provide more stable behaviour on neural networks, as the S₀ intensities are not quantitative and can be arbitrarily large or small, depending on the scaling factor applied by the scanner. Depending on model complexity and number of compartments considered, different training sizes were used for different SBI architectures (2M for single-fibre models, 4M for N = 2 fibres, and 6M when N = 3). Once the networks were trained, we drew M = 50 samples from q_Φ(Ω|Y_obs) to characterise the posterior given unseen observations, following the MCMC implementation that draws M=50 thinned samples from the posterior [7] (notice that none of the trends in the results changed with more samples). Code for the implemented architectures can be found on https://github.com/orgs/SPMIC-UoN/SBI_dMRI.

We performed extensive comparisons between posterior estimates obtained from SBI and from MCMC (as available in FSL-BedpostX implementation (v.6.0.7) [24]) on both single and multi-shell models. These included comparisons using both synthetic and in-vivo data, at a voxel (e.g., fibre orientations) and whole brain scale (tractography reconstructions). The equations and parameters underlying the MCMC implementation are provided in Supplementary Material for completeness, as well as the metrics we used for assessing accuracy and precision of parameter estimates.

We also investigated how SBI design choices affect accuracy and precision, and agreement with MCMC. In the following subsections we present novel explorations on the training design for handling priors, treating noise and addressing model selection challenges. We finally present the data used and data processing aspects all the way to probabilistic tractography.

2.3.1. Uniform vs restricted priors p(Ω)

Non-informative priors have been used in recent dMRI SBI implementations [48, 50, 52, 54], as they provide the simplest choice for training and offer uniform coverage over the whole parameter space. However, certain combinations of parameters will never be found in real-world cases or would not be feasible to detect due to limitations in the measurement process. Furthermore, certain model parameters have non-linear boundary conditions that cannot be easily defined by parametric prior distributions, requiring ad-hoc reparametrisation that often are not trivial [50]. For instance, the sum of volume fractions in a multi-compartment model should be between 0 and 1. Using non-informative priors and forcing NPE to learn from implausible regions in parameter space could make training highly inefficient, potentially reduce the predictive performance, and increase estimation uncertainty.

We tested two SBI designs, one with uniform priors and another with restricted priors. While the former is straightforward to implement, the latter requires a more advanced design [61, 62]. For the uniform prior implementation we used: p(d) ~ 𝒰(10⁻⁵, 5 · 10⁻³), p(d_std) ~ 𝒰(0, 5 · 10⁻³),red p(f₁) ~ 𝒰(0, 1), and p(f_n) ~ 𝒰(0.05, 1). The fibre orientations were sampled in the positive hemisphere (i.e. θ ∈ [0, π/2], ϕ ∈ [0, π)) to avoid degenerate solutions due to antipodal symmetry. To ensure uniform sampling on the sphere, the following correction was applied: θ = arccos(1 − 2u) and ϕ = π·u, where u~U (0, 1) (see Supplementary Material and Suppl. Fig. S1).

Subsequently, we designed and trained a Restricted Prior distribution that learns feasible regions of the parameter space delimited by (non-linear) constraints on top of the parameter prior, as summarised in Table 1 (soft constraints based on empirical evidence and prior studies on the models considered here ([7, 30]) to avoid degeneracies). To achieve that, we trained a classifier to learn what is the signal stemming from “valid” parameter combinations, i.e. combinations for which restrictions and boundary conditions on individual parameters are respected (see Suppl. Fig. S2). To do so, we 1) sampled each parameter independently from uniform priors p(Ω) specified above, 2) used these samples to generate forward predictions using f_Y (Ω), 3) tagged the cases that did not meet the boundary conditions into a different class of “non-valid”, and 4) trained a classifier accordingly to separate valid vs non-valid cases (see Fig. S2). This classifier was wrapped into a density distribution object to produce samples only from the (“valid”) restricted regions of the parameter space. This trained Restricted Prior object is a joint prior distribution that produces samples for all model parameters at once, in a way that they fulfill all conditions imposed. We trained separate Restricted Priors for the cases of 1, 2 and 3 fibres, and compared the efficiency and performance of the inference against the network using the default uniform priors.

Table 1. Parameter constraints and restricted priors in the Ball&Sticks model.

These ensure that i) the sum of the volume fractions ranges within [0, 1], ii) there is some ordering in the volume fractions f₁ > f₂ > f₃, iii) detectable crossing angles cannot be lower than a minimum, reflecting limitations in angular contrast/resolution in the data.

	Parameter constraints
	1 Fibre	2 Fibres	3 Fibres
f 1	~ U[0, 1]	~ U[0.05, 1]
f 2	~ U[0, 0.001]	~ U[0.05, 1]
f 3	~ U[0, 0.001]		~ U[0.05, 1]
f sum	f1 ≤ 1	(f1 + f2) ≤ 1	(f1 + f2 + f3) ≤ 1
f n	None	f1 ≥ f2	f1 ≥ f2 ≥ f3
fn − fm (m > n)	None	> 0.05
angle(vn, vm)	None	> 10°	> 30°
dstd/d	≤ 2

2.3.2. Dealing with noise during training

Noisy realisations of the signal were considered during training. For consistency with the MCMC implementation, we assumed zero-mean Gaussian additive noise (i.e. f_Y (Ω) → f_Y (Ω) + ε) with ε ~ 𝒩 (0, σ), σ = S₀/SNR and SNR the signal-to-noise ratio of the b = 0 signal.

Matching the noise distribution between training and test data improves robustness and accuracy of SBI estimates [63]. However, as such matching cannot be performed a-priori for a generalisable implementation, we evaluated the performance of SBI under two strategies for considering noise during training. For a given set of parameter values Ω_i we explored two strategies: a) Single-level noise: A single noisy realisation Y_i was generated with an SNR value chosen randomly from a prior p(SNR) ~ 𝒰 (3, 80), b) Multi-level noise: The same noise-free signal was corrupted with different noise levels chosen from multiple predefined SNR intervals. Here we used 8 SNR intervals: p(SNR₁) ~ 𝒰 (3, 10), p(SNR₂) ~ 𝒰 (10, 20), … p(SNR₈) ~ 𝒰 (70, 80). I.e. we produce 8 noisy realisations of each Y_i, each at a different noise level. We kept the total number of training points equal for both strategies. Hence, the former approach provided for training more combinations of parameters sampled from the priors, each with a single noisy realisation; while the latter approach provided (8-times) fewer combinations of parameters for training, but with more examples of noise contamination for each combination.

2.3.3. Model selection using SBI

In the multi-compartment model considered, there is an inherent challenge of deciding how many fibre compartments (i.e. orientations) must be estimated in each voxel. This is a model selection problem that the MCMC implementation addresses using Automatic Relevance Determination (ARD) priors [7] (see Suppl. Material). These are sparsifying improper priors that cannot be sampled from, hence not well suited nor easily adaptable to simulations from the prior, as needed in SBI. Hence, to address the model selection problem with SBI, we have implemented and evaluated two strategies:

A) ClassiFiber + Individual NPEs: Similar to [52, 64], we subdivided the problem into 2 stages: an initial classification step to select how many fibre compartments correspond to the signal of each voxel, and a second step that fits the corresponding model based on the classifier output. For the first step, we designed the ClassiFiber, a feed-forward deep neural net that can take the dMRI signal (or a basis signal representation projection) as input and provide the number of fibres as output (with N = 1, 2, 3). The neural network is formed by 6 layers (input, 256, 128, 64, 32, output), with ReLu activation functions, 1% dropout, adaptive learning rate, Adam optimisation, and a Cross-Entropy loss function. We generated 6 million parameter combinations (2 million for N =1, 2 and 3 crossing-fibres configurations sampled from the restricted priors) and generated signal predictions using the multi-shell model. For the second step, we trained 3 independent NPEs for models with N =1, 2, or 3 fibre compartments, respectively, using the basis structure of sec.2.1. Hence, a network e.g. with N = 2 compartments was fed with only 2-way crossing patterns during training, while a network with N = 1 compartment was not trained with any fibre crossing examples. Note that training data with N = 1, as generated using p(f₁) ~𝒰 (0, 1), considered cases where f₁ was zero or very close to zero. Therefore, they implicitly included a set of N = 0 cases, where no fibres were present and the signal was dominated by the ball compartment. We refer to this setup as SBI_ClassiFiber.

B) Joint NPE: Contrary to the independent NPEs trained above, we trained a network using a training super-set that included examples of 1,2 and 3 crossing-fibre cases (equal number of 1, 2, and 3 fibre cases). In order to keep a fixed parameter vector size, N = 3 compartments were set into the model. For cases where N < 3 fibre compartments were simulated (i.e. for single fibres and 2-way crossings), the respective volume fractions of the n^th degenerate compartments in the training data were set to f_n ~ 𝒰(0, 1e − 3), while the values for the orientations (θ_n, ϕ_n) were randomly sampled from the priors (e.g. when two fibres were simulated, the volume fraction of the third compartment was set to almost zero in the training data). This joint dataset was used to train a single NPE, which we evaluated on whether it can inherently learn the correct number of compartments, without a prior classification as above. Similar to SBI_ClassiFiber, N = 0 cases were implicitly considered in the N = 1 training set, with f₁ ~ 0.. We refer to this setup as SBI_joint.

Note that this joint NPE approach does not directly estimate the number of compartments N. Instead, a model with a maximum number of compartments N_max = 3 is fitted and the number of compartments supported by the data is found indirectly by a-posteriori setting a minimum threshold to the estimated volume fractions f_i > f_cutoff, i = 2, 3. We have used f_cutoff =5%, as similarly used in the MCMC implementation [7].

2.4. Datasets, processing and tractography

Data: We evaluated the different architectures using both synthetic and in-vivo data. Preprocessed in-vivo dMRI data were obtained from a previous study [65], consisting of 6 re-scans of the same subject, each acquired following a UKBiobank-like multi-shell protocol (2mm isotropic resolution, TR=3s, TE=92ms, MB=3, Siemens Prisma 3T scanner) [66]. Fifty gradient directions at b=1000 s/mm² and another fifty gradient directions at b=2000 s/mm² were obtained (i.e. a total of 100 datapoints per voxel), with a mean SNR of 30 for the b=0 s/mm² image. Diffusion-weighted volumes were normalised by the mean of the b₀ volumes to obtain the signal attenuation. Synthetic data allowed toy examples for testing the inference frameworks under scenarios where the true parameter values were known. We used the attenuation of the signal (with S₀ = 1) from the Ball&Sticks model to generate such synthetic datasets, following the same acquisition protocol (bvals and bvecs) as in the in-vivo datasets. We chose simulation parameters for the synthetic examples, where -based on prior sensitivity analyses studies [7, 30, 67]-reliable estimation would be expected given our acquisition protocol, and assessed the relative performance of SBI vs MCMC.

We considered both single-shell (Eq. 3) and multi-shell cases (Eq. 4) with up to N = 3 fibre compartments. The number of model parameters inferred in the multi-shell case was 2 + 3 ∗ N, where N is the number of fibre compartments, plus the noise variance. I.e., 6, 9 and 12 parameters for the 1-, 2-, and 3-fibres multi-shell model, respectively. For the single-shell models we had 1 parameter less (there is no d_std).

Acquisition-specific vs Acquisition-agnostic training

As the above approach returns trained networks for a specific set of bvals and bvecs, we also explored using an acquisition-agnostic basis set representation of the data as input to the networks, instead of using directly the dMRI signal attenuation. The signal attenuations were expressed using Spherical Harmonics (SH) [68] (for single-shell) and the MAP-MRI basis (for multi-shell) [69]. Given the number of volumes in the data, we retained the even SH coefficients up to order L = 6 (28 coefficients), and for the MAP-MRI up to radial order r = 6 (45 coefficients) with Laplacian regularization (no positivity constraints nor anisotropic scaling). Both signal representations were obtained using Dipy implementations [70].

Probabilistic tractography

Using the trained architectures, we drew 50 samples from the posterior distribution of parameters given the observed data. For the same data and model, we drew 50 (thinned) posterior samples using MCMC as well. Each provided voxel-wise orientation distributions that mapped uncertainty for the estimated parameters. We subsequently used the estimated posterior orientation samples in probabilistic tractography. We ran the standardised FSL-XTRACT tractography protocols [71] using the estimates of each method. These protocols included a range of white matter tracts (commissural, projection, association and limbic). The SBI-derived and MCMC-derived tracts were then directly compared. As both of them reflect stochastic estimation processes, we further explored how the agreement in tract reconstruction between SBI and MCMC compares with scan-rescan reproducibility in MCMC-based tract reconstruction, across the six available repeats.

3. Results

3.1. Single-fibre model

We performed initial evaluations in the simple case of a single-fibre model (N = 1 in Eq. 3 and Eq. 4). Fig. 2 shows example comparisons for the posterior mean and uncertainty estimated by MCMC and SBI, in both single-shell (B) and multi-shell (C) models. A default SBI training paradigm was used in these first comparisons, with the same priors as in MCMC. Representative examples of the marginal posteriors for f₁ in different areas of the brain are shown in Fig. 2A, with very good agreement between MCMC and SBI in the means of these posteriors, and some small variation in their coverage. Overall, as can be seen in the whole-brain maps, the agreement between these two very different inference approaches is extremely high, with correlations of the posterior means above 0.95 for both diffusivity and volume fraction, and median differences of 2o for fibre orientations across white matter. Similarly, uncertainty maps show correlations of 0.92, i.e. a similar uncertainty mapping across the brain (lower uncertainty in WM and higher in CSF and GM for the fibre orientation). A very similar picture is depicted in Fig. 2C for the multi-shell model. These first findings suggest high agreement in both accuracy and precision of SBI estimates against MCMC, in the case of a simple model with a default training paradigm. In the next sections, we increase the model complexity and evaluate how different SBI training and designs affect performance.

Figure 2. Comparison of SBI and MCMC estimates for a single-fibre model (N = 1) in case of single-shell (A,B) or multi-shell (C) data.

A) Example posterior density plots of the anisotropic volume fraction f₁, as estimated by MCMC (blue) and SBI (orange), in four representative voxels: Centrum Semiovale (where fibre crossings are expected), highly anisotropic Corpus Callosum, CSF-filled ventricles, and gray matter. B) Whole-brain maps of estimated parameters using single-shell data (from left to right): Mean of d posterior, mean of f₁ posterior, mean of v₁ posterior, orientation uncertainty (width of the v₁ posterior). C) Similar as in (B) for multi-shell data. For (B, C), density scatter plots represent the agreement of estimates throughout the brain; the correlation r is calculated for scalar parameters, and the angular difference (in degrees) is calculated for the mean fibre orientations. NPEs were trained with 1 million samples from uniform priors with random SNR values (within the range of 3 to 80).

3.2. Multi-fibre models - Synthetic data

3.2.1. Optimal training strategies

We subsequently increased the complexity of the models to include crossing fibres (2 ≤ N ≤ 3). First, we compared different training strategies using synthetic data (where ground truth values for model parameters are known). We investigated 1) the increase in estimation accuracy due to usage of restricted priors compared to uniform priors, 2) the optimal strategy for considering noise during training, 3) whether alternative signal representations can provide similar performance (e.g., using the raw signal vs spherical harmonics), and 4) how different SBI designs perform (independent vs joint training). We generated data with N = 2 fibre compartments, where the anisotropic volume fraction f₁ and the crossing angle were varied. Fig. 3 shows a collection of comparisons of the mean of the estimated (marginal) posterior distribution for volume fractions and crossing angles, demonstrating how accurately each inference approach captures the ground-truth patterns. For a representation of the error with respect to the ground-truth parameter values, see Suppl. Fig. S3; and Suppl. Fig. S4 for the precision/uncertainty in the estimates.

Figure 3. Evaluating performance of different SBI designs in estimating model parameters in synthetic fibre crossings (N = 2), where ground truth is known.

dMRI was synthetically generated for patterns with two crossing fibres, varying the anisotropic volume fraction f₁ and their crossing angle (rest of parameters were d = 0.0012mm²/s, f₂ = 0.2, SNR = 25 for all experiments). The ground-truth parameter values employed in these experiments are colour-coded (top-left corner), with each entry of these matrices corresponding to a different combination of f₁ and crossing angle. The ability to accurately estimate these values is explored with different SBI training designs. Estimates (median values over 100 noisy realisations) for the mean posterior of f_sum (top row) and crossing angle(v₁, v₂) (bottom row) are shown for the following cases: A) Different prior definition and noise strategies for SBI training (from left to right): (i) Uniform priors and single-level noise, (ii) Restricted priors and single-level noise, (iii) Restricted priors and multi-level noise. B) SBI trained with raw dMRI signal (acquisition-specific) vs a basis set signal representation (acquisition-agnostic). Spherical harmonics were used for single-shell data, and the MAP-MRI projection for multi-shell data. C) Comparing MCMC fitting a model with N = 2 fibres against NPEs trained only with 2-fibres examples (Independent) or with a joint superset that contains both 1- and 2-fibres cases (Joint). All SBI models were trained with 5 million samples, apart from Joint approach which was trained with 6 million samples including 1- and 2-fibre cases (3 million each). Results denoted with (*) correspond to the same SBI design across A, B, C (i.e. trained with only N = 2 examples, multi-shell signal, restricted priors and multi-level noise).

Among the different strategies explored, using a multi-level noise scheme combined with restricted priors provided the best performance in retrieving the ground-truth. These improved both accuracy (see Fig. 3A) and precision (see Suppl. Fig. S4A). With restricted priors, the sampling efficiency was increased as all of the proposed parameter combinations met the conditions in Table 1. By sampling directly from “within-bounds” regions of the parameter space using the restricted prior, speed-ups of 62% were observed during the generation of the training sets, compared to sampling from uniform priors. This is because uniform sampling needed extra rejection steps to discard up to 39% of “out-of-bounds” parameter combinations (see Suppl. Fig. S2 for an example of how the parameter space gets constrained). Hence, multi-level noise and restricted-priors are used for any SBI model presented onwards.

Fig. 3B shows a comparison of the obtained estimates when SBI is trained against signal attenuation (for a given set of bvals and bvecs) or acquisition-agnostic signal representations, for both single-shell (using spherical harmonics) and multi-shell (using the MAP-MRI projections) scenarios. Some loss of accuracy can be observed in the signal representation cases, particularly at small crossing angles, suggesting that acquisition-agnostic training induces a small penalty in performance compared to training for data specific to a particular acquisition. However, overall good performance is preserved. For the rest of the paper we kept SBI trained against the full signal, but we will also present some in-vivo results using the MAP-MRI basis.

Finally, Fig. 3C shows the performance of MCMC (with N = 2 and no ARD priors) compared against NPEs trained only with 2-fibres examples (Independent) or with a joint superset that contains both 1- and 2-fibres cases (Joint). Both MCMC and Independent are always set to fit N = 2 fibre compartments (i.e. the right model for these data), while Joint inherently performs indirect model selection (it can fit either a N = 1 or N = 2 fibres model). Hence, the former corresponds to SBI_ClassiFiber in the ideal case of a perfect identification of the number of compartments by the classifier (N=2), while the latter corresponds to SBI_joint with N ≤ 2. Overall, all methods recovered appropriately the ground-truth values, with Joint having slightly worse performance than the other two, indicative of its higher complexity, but also its ability to indirectly perform model selection.

The above trends in terms of accuracy for the parameter estimates also hold in case of synthetic data where N = 3 fibres have been simulated (see Suppl. Fig. S5). Taken together, the findings suggest that SBI with appropriate training can perform very well in terms of posterior accuracy and precision (and similarly to MCMC) in synthetic data, when the model complexity (i.e. number of compartments) is a-priori known.

3.2.2. Dealing with model selection

We subsequently explored more comprehensively the performance of different inference approaches with respect to model selection in 3-way crossing patterns. One of the main challenges when performing inference for a multi-compartment model is the identification of the number of compartments that are relevant to the observed signal. We explored two alternatives to deal with this challenge, SBI_ClassiFiber and SBI_joint. SBI_ClassiFiber performs an initial classification step to decide the appropriate model complexity N and then fits the respective NPE with N = 1, N = 2 or N = 3 compartments, each of which are independently trained. SBI_joint aims to learn the inherent structure by being trained jointly against N = 1, 2, 3 cases, avoiding the classification step and the need of training different NPEs. We compared these architectures in synthetic data against MCMC, for which we fit the most complex model (N = 3) and used ARD shrinkage priors for the volume fractions of each compartment [7].

Fig. 4A shows the performance for similar synthetic examples as before, where 2-fibre crossings were simulated, but up to 3 compartments could be estimated by all approaches. This allowed us to explore whether underfitting/overfitting occurred, contrary to the previous tests, where the number of compartments in the model was pre-fixed to the actual number of compartments simulated. All methods correctly identified the correct number of fibre compartments and showed good overall performance, with MCMC providing the best mean estimates. SBI architectures were close in estimation accuracy, with SBI_joint, as the most complex NPE, having slightly lower performance.

Figure 4. Evaluating performance in detecting the right number of fibre compartments in synthetic experiments, of MCMC, SBI_ClassiFiber and SBI_joint, all fitting up to N = 3 fibre compartments.

A) Recovering a two-fibre pattern (same synthetic data as in Fig.3, where f₁, f₂ 0 and f₃ = 0), when fitting 3-fibre models to detect over/under-fitting. Estimates (median values over 100 noisy realisations) for the mean posterior of f₁, f₂ and f₃ and crossing angle(v₁, v₂) are shown. Both MCMC and SBI return zero f₃ in all cases, detecting two fibre compartments with the correct volume fractions. B) Performance in compartment identification as a function of SNR for the three different approaches and for 10,000 synthetic patterns containing N =0-3 fibres compartments (with minimum crossing angle of 40°). For each SNR, matrices show for each ground-truth number of fibres (matrix row), the percentage of cases that were estimated as supporting N = 0 − 3 patterns (matrix column), i.e. each row of each matrix sums up to 100%. The ideal detection rate corresponds to a diagonal matrix; off-diagonal elements indicate over/under-fitting (above and below diagonal, respectively). In all cases, mean posterior f_n >5% was used to indicate support for the existence of a compartment. Results for SBI_ClassiFiber refers to the number of fibres supported by the final NPE estimates (i.e. f_n >5%), and not to the initial output of the classifier, which can be refined by the NPE.

We then tested the model selection performance in a larger set of 10,000 synthetic scenarios sampled from the restricted priors, with N = 0 − 3 (i.e. from zero to three fibre compartments, with various combinations of diffusivities, volume fractions and crossing angles). Fig. 4B shows in matrix format, comparisons of the number of fibre compartments detected (matrix rows) when different number of fibres were simulated (matrix columns). These detection matrices are presented for different noise levels and for the different inference approaches; diagonal matrices correspond to ideal model selection, while elements above/below the diagonal correspond to over/under-fitting. Overall, the three approaches achieved comparable results, with better detection rates (diagonal of the true vs predicted matrices) at higher SNR and simpler fibre configurations. At mid- and low-SNR, MCMC achieved a better identification of the 2-way crossings while SBI approaches provided a better detection of 3-fibre patterns, especially for SBI_ClassiFiber. At low-SNR, MCMC had the tendency to underfit (lower diagonal of the matrices) rather than overfit, which was not the case for the SBI approaches. This could be explained by the usage of a shrinkage ARD prior in MCMC [7], while the model selection in SBI is learnt directly from data.

3.3. Multi-fibre models - In-vivo data

After evaluating and optimising the performance of SBI in terms of accuracy, precision, and model selection in synthetic data, we compared them against MCMC using in-vivo dMRI multi-shell brain data. As there is no ground-truth available in the case of real data, MCMC estimates can be only treated as a reference. Fig. 5A shows a comparison of whole-brain mean and uncertainty maps obtained by the different methods for a model with N = 2 fibre compartments. Both SBI models provide very high correlations with MCMC (above r = 0.9 for all scalar parameters, see Fig. 5B). This inherently implies also a good agreement in the model selection results. Estimated crossing fibres show a good agreement as well (Fig. 5C and Suppl. Fig. 6A), with the median angular difference between SBI and MCMC being 2o for v₁ and 6o for v₂ in SBI_ClassiFiber (3o and 8o, respectively, in SBI_joint). The example through the centrum semiovale shown in Fig. 5C, highlights all approaches resolving crossings between callosal and corona radiata projections, with the SBI architectures returning less of fibre fanning within the corpus callosum. Regarding orientation uncertainty, MCMC patterns (e.g. lower orientation uncertainty in WM) are overall reproduced through SBI although some broader posteriors were observed, especially in the second fibre orientation (see last two columns of Fig. 5A and B). Uncertainty maps for the secondary orientations have a bimodal distribution (low uncertainty for the voxels that support 2-way crossings and high for the rest in WM, GM and CSF). SBI approaches provided less “binary” uncertainty maps than the MCMC, with more regions having intermediate values between the two modes (i.e. between 0 and 0.6 - see Suppl. Materials for explanation of these values and Suppl. Fig. S7A for a different representation of these uncertainty plots).

Figure 5. Inference for in-vivo brain data using a N = 2 fibres multi-shell model and different inference approaches (MCMC, SBI_ClassiFiber, SBI_joint).

A) Exemplar coronal slice with the mean and uncertainty maps estimated by the different methods for different model parameters. B) Density plots representing the agreement between MCMC and SBI in the white matter; correlation r is calculated between scalar maps. C) RGB colour-coded crossing-fibres estimated by each method in the centrum semiovale. The mean posterior orientation is plotted for each compartment. The length of the orientation vectors are scaled by their relative voxelwise volume fraction.

Similar trends can be observed in Fig. 6 for a N = 3 fibres multi-shell model. Panel A) shows mean maps of the diffusivity and volume fractions. As in Fig. 4, SBI_joint showed a more similar number of complex fibre patterns detected to MCMC, with 22.5% and 25.4% of the voxels in the white matter identified as 3-way crossings by each method, respectively, while SBI_ClassiFiber returned 33.3%. High agreement with MCMC was found for the first two compartments by both SBI methods, with correlations of r_d = 0.97, $r_{d_{std}}=0.88,r_{f_1}=0.97$, and $r_{f_2}=0.88$ ($r_{f_2}=0.87$for SBI_joint). Larger differences were found in the third compartment; the correlation in the mean f₃ reduced to 0.7 in SBI_ClassiFiber, and 0.58 in the SBI_joint.

Figure 6. Inference for in-vivo brain data using a N = 3 fibres multi-shell model and different inference approaches (MCMC, SBI_ClassiFiber, and SBI_joint).

A) Exemplar coronal slice with the mean scalar maps estimated by the different methods for different model parameters (correlation with MCMC is indicated at the bottom of each map). B) RGB colour-coded crossing-fibres estimated by each method in the centrum semiovale. The mean posterior orientation is plotted for each compartment. The length of the orientation vectors are scaled by their relative voxelwise volume fraction. C) Orientation uncertainty maps for each fiber compartment and for each inference approach.

Panel B of Fig. 6 shows an example case of the orientations in the centrum semiovale, where 3-way crossings can be expected (corpus callosum fibres running left-to-right, corona radiata running superior-to-inferior and association fibres anterior-to-posterior). All methods depicted these, with SBI_joint having slightly lower coherence and crossings overall. Whole-brain orientation estimates are shown in Suppl. Figs. S6B.

Uncertainty maps for all three fibre orientations are shown in panel Fig. 6C, where SBI methods exhibited in general higher uncertainty compared to MCMC. Even if the general patterns were similar in all methods (i.e. highest uncertainty in CSF and GM, lower uncertainty in WM), SBI methods returned higher uncertainty in WM with increased model complexity.

Finally, as suggested in the previous section, using the MAP-MRI signal representations for a 3-fibres model led to similar parameter estimates (mean posterior) as using the raw dMRI signal (see Suppl. Fig. S8).

3.3.1. Sources of SBI vs MCMC uncertainty quantification differences

So far we have observed that MCMC and SBI return similar patterns in the parameter estimates, but some discrepancies in the width of the posterior can be observed, particularly evident in in-vivo data. We further explored these differences in uncertainty quantification and identified scenarios where these occur.

A first identified case corresponds to the potential of fitting fewer compartments than supported by the data (e.g. due to inaccurate model selection). Fig. 7A shows the marginal and pairwise joint distribution for all parameters in an example synthetic case where a 3-way crossing pattern is simulated (crossing-angles of 90 degrees, f₁ = 0.4, f₂ = 0.3, f₃ = 0.2), but a 2-fibre model is fitted, i.e. wrong model identification. MCMC returns unimodal marginal distributions for all parameters; it finds one of the 3 possible orientation modes and keeps sampling from it (as it is very expensive to jump to a different mode), missing the other modes and providing very low uncertainty (overconfidence). On the other hand, SBI captures all three orientation modes, assigning in this case their samples to (mainly) v₁ (see red arrows) which, therefore leads to increased uncertainty due to multiple modes arised from fitting a model with the wrong number of compartments (epistemic uncertainty).

Figure 7. Sources of differences in uncertainty mapped with MCMC and SBI.

A) Fitting fewer compartments than supported by data: Marginal and pairwise joint distributions of all parameters obtained from a 2-fibre model fitted to synthetic data from a 3-fibre pattern (ground-truth), using MCMC (left) and SBI (right). The marginal posterior distribution for every model parameter is shown first in each row, followed by the joint 2D posterior with every other parameter. When the number of compartments in the model is lower than the number of compartments supported by the data (i.e. effectively a wrong model is fitted), MCMC tends to return unimodal posteriors, missing the third compartment that is supported by the data. This leads to narrow (but wrong) distributions (overconfidence). SBI correctly captures all three modes within the compartments of the model (see red arrows, all captured in the v₁ posterior in this example). However, this leads to broader posteriors (i.e. uncertainty). B) Compartment swapping: Marginal and pairwise joint distributions of all parameters obtained from SBI for a voxel from in-vivo data (square), where f₁ ~ f₂ (blue). In such cases, SBI can include samples from the first and second fibre orientation indistinctively to the corresponding marginal posteriors, making them multimodal (red arrows). A simple re-ordering (based on similarity of orientation samples) can mitigate this (green), leading to unimodal posteriors. Notice that there is nothing wrong with either approaches, as the joint orientation posterior samples stay the same, and what changes is how they are distributed to compartments. For reference, orientation uncertainty maps of v₂ returned by SBI, before (left) and after the correction (right) are shown.

Another source of differences in uncertainty is compartment swapping, which can occur when volume fractions of fibre compartments are very close to each other e.g. (f_i − f_j) ≈ 0. In such cases, assigning a compartment to be first, second or third becomes arbitrary, as any combination could explain the signal in the same way. SBI is based on minimising the negative log-likelihood, so when the compartments have similar volume fractions, it will equally likely assign a sample from orientation v_i to parameter v_j and vice versa. Hence, drawing multiple samples can lead to a multi-modal marginal posterior for a specific fibre orientation parameter with P modes, with P the number of compartments with similar volume fraction. This compartment swapping is much more difficult to occur with a random-walk Metropolis Hastings MCMC algorithm, which needs very high energies to jump between multiple modes and therefore samples tend to inherently cluster together. Even if this issue in SBI is reduced by the prior constraint (f_i − f_j ≥ 0.05), it can still occur for relatively similar volume fractions that are just above this difference threshold. It can be further mitigated with a simple re-ordering of the posterior samples based on the similarity of orientations (rather than the default ordering based on the rankings of the volume fractions). Fig. 7B shows the marginal and pairwise joint distribution for all parameters in an example case from in-vivo brain data, with compartment swapping, before and after the re-ordering to compensate for the swapping. Suppl. Fig. S9 shows how reordering the SBI samples achieves very good correspondence between SBI and MCMC in an example voxel with 3-way crossings.

The above findings suggest that the differences observed in the uncertainty maps are potentially inflated by the way parameter samples are grouped into compartments and uncertainty is estimated. They also hint that individual samples in SBI are likely correct in all cases (Suppl. Fig. S9), even if summary uncertainty maps suggest otherwise when compared against summary maps from MCMC. We therefore performed probabilistic tractography, which uses the posterior orientation samples rather than summary uncertainty maps, to test this hypothesis.

3.4. Probabilistic tractography using SBI estimates

Probabilistic tractography spatially propagates orientation uncertainty (i.e. samples from the marginal posterior distribution of the fibre orientations) to build spatial distributions that estimate white matter pathways. We explored how the orientations estimated from the 3 different approaches (MCMC, SBI_ClassiFiber, SBI_joint) affect tractography performance. We used the posterior fibre orientation samples estimated by each method using a N = 3 fibres, multi-shell model. Fig. 8 shows example Maximum Intensity Projections (MIPs) of the path distributions for a range of 42 commissural, association, limbic and projection tracts [71]. For reference, the population-average UK-Biobank atlas is also shown, which was created using the same MCMC implementation for estimating fibre orientations and the same FSL-XTRACT implementation for tractography [71], as used here. All methods reconstructed all tracts successfully and showed a high qualititative agreement between the corresponding paths. SBI-reconstructed tracts had a higher correlation with the atlas (up to 15% more in SBI_ClassiFiber), compared to MCMC-based reconstructions.

Figure 8. Probabilistic tractography using MCMC and SBI posterior orientation distributions.

White matter tract reconstruction for 42 bilateral bundles were obtained using the orientation estimates of each method. The spatial distribution of each of the paths is shown as a maximum intensity projection on axial (top row) and saggital (bottom row) planes (thresholded at probabilities>0.1%). Tract reconstructions from a population-average atlas (using data from UK-Biobank and MCMC-based orientation inference) is shown as a reference. Each tract was spatially correlated against the corresponding reconstruction from the atlas. The median of these correlation values across tracts is indicated under each method.

We then looked into quantitative comparisons between the reconstructed tracts obtained using MCMC vs SBI estimates. Fig. 9A shows scan-rescan reproducibility of each method among 6 consecutive repeats of the same MRI acquisition protocol on the same participant. SBI methods reported higher reproducibility in tract reconstructions, with an overall median correlation r of r = 0.96 for SBI_joint and r = 0.95 for SBI_ClassiFiber (among tracts and repeats), which is consistently higher than in MCMC (r = 0.87). This provides evidence for the hypothesis that posterior orientation samples from SBI are appropriately precise and support the sources of the seemingly higher orientation uncertainty shown in the previous section. In addition, Fig. 9B demonstrates the median (among 6 repeats) correlation between MCMC-based and the corresponding set of SBI-based reconstructed tracts. This is overlayed on a shaded zone (in light blue), which depicts the interquantile range (across the 6 repeats) of the correlations between MCMC-based tract reconstructions themselves. In general, both SBI_ClassiFiber and SBI_joint perform similarly well to MCMC; certain tracts that have multiple 3-way crossings along their route are an exception (AF, AR, SLF1, SFL2, SLF3) as for these SBI_ClassiFiber performs better than SBI_joint. Nevertheless, these plots suggest that differences between MCMC and SBI in tract reconstructions are within the range of (and most of the time smaller than) scan-rescan variability (i.e., the variability of MCMC-based tract reconstructions across the 6 repeats). For the exceptions observed, where correlations between MCMC and SBI methods were lower, SBI methods showed a more precise reconstruction among repeats compared to MCMC (see 9A), as also indicated in Fig. 9C, where the averaged reconstructed tracts over the 6 repeats are shown.

Figure 9. Scan-rescan reproducibility of probabilistic tractography reconstructions for a range of white matter tracts, across 6 MRI scan repeats of the same participant.

Correlations of the spatial path distributions reconstructed for each tract and using orientation samples from MCMC, SBI_ClassiFiber and SBI_joint are shown. A) Median (across repeats) pairwise tract correlation for each method (i.e. median across rep1 vs rep2, rep1 vs rep3, rep1 vs rep4,… rep5 vs rep6). B) Correlation between SBI and MCMC-based reconstructed tracts (median across 6 repeats SB1 vs MCMC1, SB2 vs MCMC2,…) overlaid on the interquartile range of correlations of MCMC reconstructions across the 6 repeats. C) Averaged path distributions (across the 6 repeats) for a subselection of tracts that had lower agreement between MCMC and SBI.

Taken together, these results indicate that for a range of representative white matter tracts, including those with complex fiber architecture and 3-way crossings along them, SBI-based reconstructions i) agreed better with a population-average white matter atlas, ii) were more precise and reproducible across MRI repeats, iii) had differences with their MCMC counterparts that were within or lower than the magnitude of differences expected due to scan-rescan variability.

3.5. Computational performance

We measured single-CPU time required for simulation, training and parameter estimation in cases of N =1, 2, and 3 fibre compartments, and compared MCMC to SBI architectures. SBI has one-off overheads for the generation of the simulated training set and the training of the NPEs, but MCMC is considerably slower during inference. For the multi-shell models considered, the generation of synthetic signals with 3-way crossings took less than 30 minutes for 3 million samples. Time for SBI training increased with the number of fibres in the model and training samples, and both approaches required similar times. For instance, SBI_ClassiFiber with N up to 3 fibres required a total of ~94 hours in a single CPU (~2h for the ClassiFiber with 1M samples, ~5h for the NPE with N = 1 and 2M samples, ~33h for the NPE with N = 2 and 4M samples, and ~54h for the NPE with N = 3 and 6M samples). Training SBI_joint with a single NPE but the same total number of 12M training samples with N up to 3 fibres required ~92 hours in a single CPU. However, the above can be significantly parallelised.

Inference times are compared in Table 2). Inference using the trained SBI networks was up to 70x faster than the MCMC, amortising the costs of training. We should point out that in these comparisons SBI is a CPU-based Python implementation, while the CPU-implementation of FSL-BedpostX used for these computational performance comparisons is highly optimised in C++. In summary, with a simple data-level parallelisation across e.g. 8 CPU-cores, inference on a standard in-vivo dMRI brain dataset, as the one used here, with the current SBI implementation takes ~ 4 minutes, while MCMC takes more than 4 hours.

Table 2. Computational performance - Inference time (in milliseconds) for each method for obtaining 50 samples from the posterior (reported times averaged across 1,000 synthetic cases).

Times are all for a single CPU. Time for classification in SBI_ClassiFiber is orders of magnitude smaller and not included. MCMC includes 1,000 burn-in iterations and a thinning period of 25 (i.e. 2,250 iterations in total).

	MCMC	SBI ClassiFiber	SBI joint
1 fibre	110.65 ms	2.86 ms	2.86 ms
2 fibres	200.87 ms	3.74 ms	3.42 ms
3 fibres	311.27 ms	4.82 ms	4.38 ms

4. Discussion

In this work, we have presented a Simulation-Based Inference framework for Bayesian fitting of multi-compartment fibre orientation models in diffusion MRI data, mapping parameter uncertainty and utilising the posterior distributions of model parameters given the data to perform probabilistic tractography. Importantly, we performed a direct and comprehensive comparison against classical inference approaches, based on MCMC, to identify optimal training configurations and SBI architectures, and confirm the similarity of voxel-wise estimates and whole-brain tractography reconstructions to an established and heavily used approach. Our findings suggest that SBI with in-silico training can match the performance of classical inference approaches, opening new avenues for amortised inference in challenging likelihood-free scenarios [72]; while offering orders of magnitude of computational speed-up over MCMC and avoiding chain convergence challenges that MCMC can face.

We performed novel explorations for optimal SBI designs, in the context of dMRI modelling. Firstly, we explored how to incorporate noise into training. Previous studies have either not considered or added noise at random values [50, 52]. We have found that following a multi-level noise scheme improves accuracy and precision of SBI, even if 8-times fewer parameter combinations are considered for training (Fig. 3). From the perspective of neural networks applied to conditional density estimation, this noise-regularisation strategy can be seen as a form of data augmentation and is in line with recent findings [73].

Secondly, classical sampling-based, such as MCMC, can be very flexible in handling parameter constraints as prior knowledge. This is more challenging in an SBI framework, where parametric prior distributions are typically expected. We explored how to incorporate boundaries and non-linear constraints in the parameters. We showed that it is possible to train restricted prior objects to learn highly complex and non-linear constraints for the model parameters that translates into better performance (Fig. 3) and sampling efficiency (Suppl. Fig. S2), which can provide important speed-ups in problems where simulations are expensive or the parameter space is highly-dimensional. Density estimators can struggle to generalize across the entire parameter space even for simple problems [74], so a trade-off between enough diversity to capture the high-dimensional parameter space (restricted priors) and having enough, slightly-perturbed, samples (noise strategy) can be beneficial during training, as shown by our results.

Thirdly, we explored two strategies for handling model selection in cases of nested multi-compartment models. Recent studies have proposed a classification step that chooses what model complexity to fit [52, 54] followed by an NPE trained independently on that model. We took a step further, inspired by recent studies [75, 64], and evaluated our implementation of this approach (SBI_ClassiFiber) vs another one where the NPE learns the appropriate model complexity by being trained on multiple potential models (SBI_joint). Results in both synthetic and in-vivo brain data suggest a comparable performance of both approaches, with SBI_joint providing slightly worse crossing fibre rate detection, accuracy and precision than SBI_ClassiFiber, but tractography end results suggesting similar or better performance than MCMC (which relied on ARD shrinkage priors for indirect model selection [76]). We showed that training an NPE with enough examples from different models can be sufficient to learn the model selection, at the cost of some precision and accuracy loss. This is important, as in SBI_ClassiFiber (as well as in similar alternatives proposed, e.g. [64, 4]), the classifier and NPEs are trained independently, so a misclassification by the former will directly affect the latter without any feedback. However, by learning the complexity indirectly as part of training, model selection and parameter fitting can interact during inference. Exploring the joint-training of the ClassiFiber and the different NPEs to allow for such interactions could shed more light on what are the advantages and disadvantages of each architecture.

We evaluated the inference performance of SBI against an MCMC-based framework for the same models, which we have developed in the past [7, 31, 30, 24], and is widely and heavily used as part of FSL [77]. For both synthetic and in-vivo brain data, we found that SBI closely reproduces the patterns of MCMC. In terms of agreement in the means of the posterior distributions (as a proxy for accuracy), correlations above 0.9 between SBI and MCMC architectures were found across the brain for most of the parameters in the different models. In terms of uncertainty mapping, correlations above 0.9 were found for simple models (with a single fibre compartment, Fig. 2). As the number of compartments increased, the whole-brain patterns remained (i.e. maximum orientation uncertainty in GM and CSF, lower in WM) but orientation uncertainty estimated by SBI in the WM was generally higher than by MCMC (see Figs. 6 and Suppl. Fig. S7). Longer training did not change these trends considerably (results not shown). Despite these seemingly large differences in uncertainty mapping, probabilistic tractography (that spatially propagates uncertainty in orientation estimates) using MCMC and SBI approaches returned very similar results. In fact, SBI-based tract reconstructions agreed better with a population-level white matter atlas, were more reproducible across repeats of the same MRI experiment, and their differences with MCMC-based reconstructions had a magnitude smaller than expected from scan-rescan variability (Figs. 8, 9).

We explored what could drive this discrepancy between uncertainty mapping at the voxel vs whole-brain level. One clear difference is how uncertainty is depicted at the two levels: a) At the voxel-level, we used the approach implemented in FSL [7, 31] both for MCMC and SBI. This discretizes the joint posterior orientation distribution by grouping orientation samples into a finite set of N fibre compartments, and subsequently calculates their mean orientation and spread around it using vector dyadic products [5]. This provides a mean orientation and uncertainty around it for each fibre compartment supported per voxel. b) At the whole-brain level probabilistic tractography propagates using directly the orientation samples, i.e. from the joint posterior of orientations. Hence, if the grouping of orientation samples into finite sets is not consistent between MCMC and SBI, large differences can be observed at the voxel level and not at the tractography level. In Fig. 7 and Suppl. Fig. S9, we presented scenarios and confirmed with examples from the data where this occurred.

Fitting the wrong model to data can be one such scenario (see Fig. 7A). This adds epistemic uncertainty and increases the challenge of what the mapped uncertainty represents. Even if SBI and MCMC account explicitly for aleatoric uncertainty by considering the effects of noise in the training dataset (see Suppl. Fig. S10), epistemic uncertainty will be also indirectly captured in the estimated posteriors. However, this can be reflected differently in SBI and MCMC. Specifically, in cases of multi-modal orientation distributions, we observed how SBI could capture all the modes (even with a model with fewer compartments than modes), spreading the samples across modes and adding into orientation uncertainty. While the MCMC implementation sampled only from certain modes, providing samples with high precision, but completely ignoring an existing mode of the target distribution (Fig. 7A). Multiple modes can also arise from non-injective models, where different combinations of parameter values produce similarly likely solutions (i.e., degeneracies). This is a known challenge for certain microstructural models and the multi-modal posterior can be used to evaluate the different likely solutions, as illustrated in Jallais et al [50]. Here, by similar mechanisms, we have identified other type of degeneracies using SBI due to compartment swapping (Section3.3.1).

4.1. Limitations and Future work

Even if we aimed to match as closely as possible MCMC to SBI implementations in order to compare their performance, a perfect matching is not feasible as they represent quite different approaches. A number of algorithmic features can be contributing to observed differences. Firstly, the MCMC implementation relies on ARD shrinkage priors, which introduce large a-priori penalties on model complexity and have considerable effects on fibre crossing detection and on uncertainty mapping. ARD priors lead to more “binary” values on the uncertainty in the fibre orientations (either maximum or null uncertainty, see uncertainty maps in Figs.5 and corresponding histogram of Suppl. Fig. S7C). Such improper priors cannot be implemented for SBI in a straightforward manner. Secondly, SBI networks use a Kullback-Leibler divergence loss function to train NPEs. Due to the mode-covering property of the KL divergence, the approximate posterior will always cover the true posterior [78] with deviations leading to broader posteriors. In density estimators, this is known as leakage and is particularly problematic in neural spline flows because they do not have analytical corrections [38, 39, 79, 41]. Longer training or refinements, such as importance sampling have been proposed to compensate to some extent to converge to the true posterior [80, 81] but it is still an open problem that requires more research. Thirdly, the Metropolis-Hastings MCMC implementation used for comparisons is based on local sampling techniques. Hence, the energy needed to jump from one mode of a posterior to another can be very high. In this way, MCMC is more likely to converge to one of the modes and “stick to it” than SBI, missing other likely modes. This also suggests that SBI can perform a better global search than MCMC methods, which highly depends on the initialisation and the progression of the chain. Despite all these factors, the level of similarity between MCMC and SBI in local and global estimates was remarkable and considerably higher than the differences.

Another challenge is disentangling aleatoric from epistemic uncertainty in Bayesian approaches, and appropriate metrics are also needed for this purpose. We explored the ability of SBI to estimate the noise level in the data (i.e. noise variance), the main source of aleatoric uncertainty. Suppl. Fig. 10 demonstrates that estimation of noise levels is reasonable both in synthetic data and real data. It shows, however, how both MCMC and SBI can deviate from estimates that FSL-eddyqc [82] provides, which are obtained from the repeated b = 0 measurements directly in a model-independent way. As SNR estimates can become dependent on the model and noise model used (epistemic uncertainty), it is debatable on whether such estimates should be preferred upon measurements that can be directly obtained from the data. This is an interesting topic and future work is needed to identify what would be the optimal way forward.

In terms of performance, scalability with parameter space dimensionality can be an issue for all inference approaches. In this work, we explored multi-compartment models with different number of compartments and SBI performed as well as MCMC for tractography with the highest number of compartments typically considered. Recent studies have considered models with very high dimensionality and demonstrated that SBI can provide accurate solutions even for large parameter spaces [83]. However, orthogonality in parameter space may interact with scalability and further work is needed to explore these aspects comprehensively.

A limitation specific to ML-based frameworks such as SBI is that training against the dMRI signal involves model predictions for a particular acquisition scheme, i.e. a particular set of b values and diffusion-sensitising gradient orientations in the case of dMRI. Even if re-training for new data and acquisition schemes can be done once and is relatively cheap, another approach would be an model-agnostic representation of the the dMRI signal, e.g. by a number of orthonormal sets of basis functions or any other latent representation learnt by the networks [50, 52, 54]. This would allow acquisition-agnostic inference, as NPEs are trained against generalisable basis sets rather than signal sampled with a specific scheme. We explored training against such representations, including spherical harmonics (for single-shell) and MAP-MRI (for multi-shell data). Promising results in both synthetic and in-vivo brain data suggested that NPEs trained against both representations could provide similar accuracy performance compared to NPEs trained against the signal, at the cost of some precision loss (Fig. 3B and Suppl. Fig. S4). This would be interesting to explore further, where spatial neighbourhood information could be also explored for this purpose [84].

In summary, we argue that SBI is a powerful approach for performing Bayesian inference and has shown comparable performance to an established and heavily-used MCMC-based framework. Although further research and validation is needed, SBI could pave the way for tackling unaddressed problems in neuroimaging, given its reliance on in-silico, likelihood-free forward predictions and its scalability to multi-dimensional problems, the two main limitations of classical Bayesian inference approaches. The advances presented here on fibre orientation mapping are directly applicable to usage scenarios of tractography towards clinical translation, for instance in the development of structural scaffolds for whole-brain personalised models of neuronal dynamics in health [85] and disease [86] or in neurosurgical planning [87].