Entrack: Probabilistic Spherical Regression with Entropy Regularization for Fiber Tractography

The structural connectivity between different cortical brain regions is established by white matter, that is composed of myelinated axons to distribute action potentials as messages between communicating neurons. The functional importance of connectivity for cognition has been undisputedly recognized by neuroscience research (Bargmann and Marder 2013; Filley and Fields 2016).

The advent of diffusion-weighted magnetic resonance imaging (DWI) (Chilla et al. 2015; Soares et al. 2013) has empowered neuroscientists and neurologists to monitor changes in the structural connectivity with potential relevance for diagnosis, prognosis and therapy of neurodegenerative diseases (Oishi et al. 2011). DWI is currently the only non-invasive, and non-radiative imaging modality, which enables neurologists to investigate the connective micro-architecture of the white matter in a minimally invasive way. Its image contrast encodes the anisotropic diffusion of water in tissue (Beaulieu 2002), making it a highly informative probe of the fibrous white matter (Bihan and Iima 2015). The axon bundles of white matter locally exhibit clear preferential directions, as shown in Fig. 1a.

However, fiber tracking algorithms are required to reconstruct consistent long-range tissue connectivity from local, voxel-centric¹ DWI measurements.

Long-range connections in the white matter are commonly referred to as streamlines, fibers or tracts. Algorithmic methods to computationally reconstruct such streamlines from DWI are known as tractography (Jeurissen et al. 2019; Nimsky et al. 2016). Schematically, tractography infers structural connections between voxels to answer questions like “Does there exist a structural connection between regions A and B?”. We show a prototypical tractography result, also referred to as tractogram, in Fig. 1b.

Tractography is clinically applied to gather health information for a number of neurological conditions, especially for preoperative planning of neurosurgery, and for research on stroke and dementia impact on brain function (Yamada et al. 2009). The data processing pipeline for tractography is composed of three stages, (i) DWI measurements of apparent diffusion, (ii) estimation of a local diffusion model per voxel, and (iii) inference of streamlines following the local diffusion model—as illustrated in Fig. 2.

Tractography poses a major challenge due to its ambiguity mostly caused by partial volume effects, since axon diameters rarely exceed few micrometers, while the DWI resolution is limited to the scale of millimeters. This lack of resolution severely complicates inference of streamlines since the superposition of diffusion information renders it difficult to disambiguate locations where fibers cross, touch, or fan apart (Jbabdi and Johansen-Berg 2011). These complex fiber configurations have been observed to be highly prevalent in the white matter of human brains (Jeurissen et al. 2013) which further impedes data analysis of white matter especially in neurology.

The majority of tractography algorithms reconstructs streamlines in a local manner, i.e. they proceed iteratively from a given seed point, and greedily determine the direction of the next step based only on the local diffusion features, and information from previous direction estimates. Tractography algorithms can be distinguished into deterministic and probabilistic methods depending on how they estimate the direction of the next step. While deterministic methods compute a point-estimate of the next direction in line with the most-likely direction, probabilistic methods infer a distribution over possible directions. Sampling from this distribution supports following multiple traces along different directions at every step. In particular, probabilistic methods are able to express the uncertainty of their predictions, which also renders them more robust in the presence of noise.

Local tractography algorithms, as opposed to global tractography (Reisert et al. 2011), reconstruct streamlines independently, in an iterative way, based on the DWI signal in a close spatial neighborhood. This design strategy has served as the core idea of many tractography algorithms since streamline tractography originally used Runge-Kutta methods to integrate the streamline progression (see Basser et al. (2000)).

Later, the works of Behrens et al. (2003), Friman et al. (2006) introduced local, probabilistic tractography models based on mathematical models for the posterior distribution of the streamline direction. While these probabilistic methods have proven to be robust to noise, they are computationally expensive, because they need to re-estimate the high-dimensional integrals involved in the posterior for every voxel.

Very recently, a new generation of models based on supervised machine learning (ML) entered the scene, and they promise to solve deficiencies of traditional models (Poulin et al. 2019). (i) ML-methods solve the parameter estimation problem only once over a representative set of examples during their training phase. Afterwards during inference, the algorithm only requires arithmetic evaluation of the model function at each voxel, which is efficiently achieved.

(ii) Moreover, ML-methods are better suited to capture complex patterns between DWI data and fiber direction in a non-parametric way than traditional approaches, which are limited by the richness of parametric statistical models.

However, ML-methods rely on fiber tracking examples to yield supervision information which is not required for traditional (“unsupervised”) methods. To circumvent this issue, supervised approaches have been trained on the output of the previous state-of-the-art algorithms in traditional tractography; furthermore, well-curated training sets are becoming available in increasing numbers (Wasserthal et al. 2018; Essen et al. 2013).

Depending on the estimation technique for local streamline directions, ML models for tractography have been formulated either as regression problems or as classification tasks. Classification models (Neher et al. 2017; Benou and Riklin-Raviv 2019) are probabilistic in nature, but require categorical classes to approximate continuous directions. In contrast, regression models (Wegmayr 2018; Poulin et al. 2017) provide the more appropriate representation for continuous directions, but we are not aware of probabilistic regression models in the context of tractography.

Uncertainty in statistical inference arises in two distinctly different flavors – epistemic and aleatoric uncertainty – as described by Kiureghian and Ditlevsen (2009) for general engineering.

Kendall and Gal (2017) discusses the estimation of epistemic and aleatoric uncertainty in the context of computer vision. The first one, epistemic uncertainty, refers to our uncertainty about the model parameters, and it decreases when more observed samples become available. The second one, aleatoric uncertainty, refers to input-dependent noise, which is inherent to the data distribution. As such, it is unaffected by the number of observed samples. Very recently, predictive models, which also incorporate estimation of aleatoric noise, have received increasing attention, e.g. for categorical classification (Sensoy et al. 2018).

Probabilistic regression has been addressed by Prokudin et al. (2018), who uses a mixture of 1-dimensional FvM distributions in the context of object-pose estimation, and by Kumar and Tsvetkov (2018) for sequence-to-sequence models for language generation. Similarly to this work, the latter proposes a probabilistic error function based on d-dimensional FvM distribution, however, their focus is rather on reducing model training time than on uncertainty quantification.

Lastly, we also mention the method of Hauberg et al. (2015) for shortest-path tractography, who uses probabilistic numerics to solve Gaussian-process ODEs.

The framework of expected log-posterior agreement defines a model selection method for algorithms and it was originally derived from information-theoretic principles (Buhmann 2010). More precisely, as described by Buhmann (2013), it measures the trade-off between informativeness, and stability of a cost minimizing algorithm in terms of the overlap that its posterior distribution exerts between repeated measurements. An algorithm’s posterior distribution is considered informative, if it narrows down the set of potential solutions for each measurement, and it is considered stable, if the posteriors obtained from repeated measurements agree with each other in spite of measurement noise.

The PA framework, sometimes also referred to as Approximation Set Coding, or Gibbs posterior agreement, has been applied in various settings such as singular value decomposition (Frank and Buhmann 2011), spectral clustering (Chehreghani et al. 2012), Gaussian process regression (Fischer et al. 2016), and combinatorial optimization problems (Buhmann et al. 2018).

The PA criterion has also been applied in the context of neuroscience, namely to determine the optimal number of clusters for cortex parcellation (Gorbach et al. 2018).

The FvM distribution is a unimodal, directional distribution defined on the d-sphere \(\mathbb {S}_{d-1}\). For random unit vectors \(\mathbf {s}~\in ~\mathbb {S}_{d-1}\), the FvM density is given bywith \(\langle \mathbf {s},\varvec{\mu }\rangle =\sum _{i=1}^d s_i \mu _i\), and the normalizing constantwhere \(I_n(.)\) denotes the modified Bessel function of the first kind. The FvM distribution is parameterized by the unit-length mean direction \(\varvec{\mu }\in \mathbb {S}_{d-1}\) and the scalar concentration \(\kappa \in \mathbb {R}^+\). We illustrate the \(d=3\) dimensional FvM density for three different concentration parameters \(\kappa \) in Fig. 3. The norm of the first moment \(W(\kappa )\), and the entropy \(H(\kappa )\) of the FvM distribution are given by We illustrate both functions for \(d=3\) in Fig. 4, and note that in contrast to the mean direction \(\varvec{\mu }\), the norm of the first moment, i.e. \(W(\kappa )\), can be smaller than 1. Indeed it vanishes in the limit of very small concentration \(\kappa \rightarrow 0\) when \(p^\text {FvM}\) approaches the uniform distribution proportional to the inverse surface of the d-sphere:For very large concentration, the FvM distribution contracts at the mean direction:where \(\delta \) denotes the Dirac measure.

In spherical regression, we want to estimate the regression function \({\mathbf {y}: \mathcal {X}\rightarrow \mathbb {S}_{d-1},\; x\mapsto \mathbf {y}(x)}\), which maps the feature space \({\mathcal {X}\subseteq \mathbb {R}^p}\) to the d-sphere \(\mathbb {S}_{d-1}\). As the feature vectors \({x\in \mathcal {X}}\) are drawn from a distribution p(x), the observations \(\mathbf {y}(x)\) are random variables. The estimated regression function is denoted as \(\varvec{\mu }\), and as the involved vectors have unit-length, the squared distance between a prediction \(\varvec{\mu }(x)\) and the corresponding observation \(\mathbf {y}(x)\) effectively reduces to the negative cosine loss, when disregarding constant terms:The loss in Eq. (7) is \(-1\), if the prediction points into the same direction as the observation, and 1 if they are anti-parallel. To obtain the corresponding probabilistic regression model, we additionally introduce a function \({\kappa : \mathcal {X} \rightarrow \mathbb {R}_+}\), which acts as the concentration parameter of a predicted FvM distribution \({p^\text {FvM}(.\mid x):=p^\text {FvM}\left( .\mid \varvec{\mu }(x),\kappa (x)\right) }\). The loss of the functions \(\varvec{\mu }(x),\kappa (x)\) is the negative log-likelihood of the direction \(\mathbf {y}(x)\) under the corresponding FvM distribution:

The functions \(\varvec{\mu },\kappa \) are typically parametrized, e.g. in terms of neural networks. Given inputs \(x\in \mathbb {R}^p\), the simplest such example would bewith the weights and biaseswhich define the parametrized FvM posteriorIn our experiments, the neural networks are more complex than Eq. (9) , but it effectively plays the same role. Given a training set \(\{(x_i,\mathbf {y}_i)\}_{i=1\dots n}\) (\(\forall i: \mathbf {y}_i:= \mathbf {y}(x_i)\)), the parameters \(\varphi \) are estimated by minimizing the empirical risk functionThe risk function inherits the property of loss attenuation from the loss function in Eq. (8), which means that the loss caused by a large deviation \(-\langle \mathbf {y}(x),\varvec{\mu }(x)\rangle \) is reduced by a low certainty \(\kappa (x)\). We illustrate the effect of loss attenuation for the FvM distribution \((d=3)\) in Fig. 5. The attenuation property ensures an increased robustness to outliers due to an adaptive sensitivity of the loss function. Furthermore, the concentration \(\kappa (x)\) is a function of the input and this fact enables us to assess the certainty of the predicted direction for any sample x.

However, in practice, these benefits of a probabilistic formulation will be severely reduced by overfitting. When the model complexity is large, e.g. for neural networks, the posterior \(p^\text {FvM}_\varphi \) can perfectly minimize the training risk, in particular its concentration estimates will be biased towards large values, as we can see from the gradient of the risk with respect to the parameters \(\varphi _\kappa \): which tends to zero asymptotically \({\forall i:\varvec{\mu }_{\varphi _{\mu }}(x_i)\rightarrow \mathbf {y}_i}\), and \({\forall i:\kappa _{\varphi _{\kappa }}(x_i)\rightarrow \infty }\), recalling that \(\lim _{\kappa \rightarrow \infty }W(\kappa )=1\). This trend is also documented in Fig. 5, where the concentration that minimizes the loss evolves towards large values for improved model fit. Eventually, the FvM distribution fitted to the training data will concentrate all its probability mass on the mean direction. As a consequence, the probabilistic model degenerates to the deterministic limit of Eq. (6), and this behavior is undesirable for down-stream information processing, where often access to typical samples is required rather than simply extracting the most-likely direction.

Moreover, at large concentrations the entropy of the output distributions is also minimal, which is detrimental for its robustness to noise, as we will discuss, together with a principled solution, in the next section.

In the previous section we argued and demonstrated that probabilistic models can treat experimental settings with noise more effectively than their deterministic counterparts. Still the essential question remains open: How robust are different parametric distributions to the fluctuations generated by a particular data source?

The Maximum Entropy Principle, well known in physics and information theory (Jaynes 1957), provides an answer to this question of model uncertainty. In contrast to maximum-likelihood estimation, which requires assumptions about the parametric form of the desired distribution, the maximum-entropy approach is based on the knowledge about moments of the desired distribution. The maximum-entropy distribution obtains its robustness from the fact that it is the least informative distribution, which still fulfills the known constraints.

To put it differently, the change of the maximum-entropy distribution with respect to perturbations of the constraints is the least possible, as it avoids any overspecification, which is not supported by the data. In the context of spherical regression, the constraints are represented by the observations x and the observed direction \(\mathbf {y}(x)\). Each observation provides a constraint in the entropy maximization over posterior distributions \(p(\mathbf {s}\mid x)\) in the following sense:where \(\mathbb {E}_{p(\mathbf {s}\mid x)}\) denotes the expectation with respect to the distribution \(p(\mathbf {s}\mid x)\). The parameter \(w\in [0,1]\) controls the width of the distribution \({p(\mathbf {s}\mid x)}\), i.e. \({p(\mathbf {s}\mid x)=\delta \big (\mathbf {s}-\mathbf {y}(x)\big )}\), if \({w=1}\), and \({p(\mathbf {s}\mid x)=C(0)}\), if \({w=0}\). The constraint can also be written as \({\mathbb {E}_{p(\mathbf {s}\mid x)} \big \langle \mathbf {s},\mathbf {y}(x)\big \rangle = w} \), which also shows that w should be interpreted as the amount of spread that \(p(\mathbf {s}\mid x)\) has around the observations \(\mathbf {y}(x)\). Using a Lagrange multiplier \(\beta \ge 0\), we can rewrite the constrained optimization problem in Eq. (14) as an unconstrained problem \(\min _{p(.\mid x)}g_\beta \big (\mathbf {y}(x),p(.\mid x)\big )\), withThe functional \(g_\beta \) represents the Gibbs free energy, which is the difference between the expected cost \(-\mathbb {E}_{p(\mathbf {s}\mid x)} \big \langle \mathbf {s},\mathbf {y}(x)\big \rangle \), and the entropy \(-\mathbb {E}_{p(\mathbf {s}\mid x)}\log {p(\mathbf {s}\mid x)}\) divided by the Lagrange factor \(\beta \), controlling the precision of the direction estimation. The precision is determined by the value w in the constraint of Eq. (14), and needs to be considered as a hyper-parameter if we only observe \(\mathbf {y}(x)\).

By variational calculus, we can derive the distribution that minimizes Eq. (15) for one particular x:which corresponds to the FvM distribution with a fixed concentration \(\beta \) around the mean direction \(\mathbf {y}(x)\). While it is well known that for the constraints in Eq. (14) the maximum-entropy distribution is given by the FvM (Mardia 1975), we are rather interested in learning the conditional distribution \(p(.\mid x)\) over all x, which minimizes the expected Gibbs free energyfor an arbitrary, but fixed data distribution p(x).

Based on the observation that the FvM functional in Eq. (16) minimizes \(g_\beta \), we make the ansatz \({p(.\mid x) = p^\text {FvM}}\)\({\big (.\mid \varvec{\mu }(x),\kappa (x)\big )}\), which replaces the constant \(\beta \) with the input-dependent concentration \(\kappa (x)\), and the unknown function \(\mathbf {y}(x)\) with the estimator \(\varvec{\mu }(x)\). If we plug this into Eq. (15), we obtain the proposed entropy-regularized loss function for the estimators \(\varvec{\mu }(x),\kappa (x)\):The entropy-regularized objective in Eq. (18) has a similar loss attenuating property as the loss based on the FvM log-likelihood from Eq. (8). As shown in Fig. 6 for \(d=3\), the free energy is minimized at lower concentration \(\kappa (x)\), when the deviation \(\langle \mathbf {y}(x),\varvec{\mu }(x)\rangle \) increases.

However, whereas the concentration diverges in the maximum-likelihood approach when \(\langle \mathbf {y}(x),\varvec{\mu }(x)\rangle \rightarrow 1\), it remains finite with the maximum-entropy method even in this case, because the precision parameter \(\beta \) limits the concentration, as we will discuss in more detail in the next section.

In analogy to Sect. 3.2, we denote the parameters of the parametrized FvM posterior as \(\varphi _\beta \). We propose to determine these parameters by minimizing the expected Gibbs free energy \(G_\beta \) (Eq. (17)). In a learning setting, we substitute the data distribution p(x) in \(G_\beta \) by the empirical distribution of the observed data \(x_i\) and \(\mathbf {y}_i := \mathbf {y}(x_i)\). The estimated parameters \(\hat{\varphi }_\beta \) of the FvM distribution areTo determine the optimal precision parameter \(\beta \), we need to obtain the posterior parameters \(\hat{\varphi }_\beta \) at different values of \(\beta \). While this sounds conceptually straightforward, more considerations are necessary in practice. To compare models at different precision values, we need to assert, that the optimization of the model parameters \(\varphi _\beta \) has indeed equilibrated at the given precision value \(\beta \). More formally, we consider the model parameters \(\varphi _\beta \) in equilibrium at a given precision \(\beta \), when the following condition holds:This condition can be motivated by the gradient of the risk with respect to the parameters of the concentration, i.e.which shows that the equilibrium condition in Eq. (20) is fulfilled, if the gradient vanishes. Additionally, we see again that the concentration \(\kappa \) remains finite, and is limited by the precision \(\beta \), in contrast to Eq. (13).

In practice, it will depend on the optimization parameters (learning rate, batch size, etc.), and the precision itself, if the condition Eq. (20) is approximately fulfilled. Besides the cumbersome tuning of optimization parameters, it is very time-consuming to re-run the optimization for each precision value with a new initialization of \(\varphi \).

Thus, we propose a robust, automatic annealing schedule to efficiently produce models in equilibrium at different precision values. The detailed annealing procedure is described by Algorithm 1, and its effect is illustrated in Fig. 7. Effectively, the progress of optimization is automatically paced by using Eq. (20) as a control criterion. As long as equilibrium is not established for a particular precision value, this value is held constant until the optimization of the model has equilibrated. This computational strategy renders model adaptation more robust to the choice of optimization parameters.

Moreover, we can efficiently extract models for different precision values during the course of a single training run, without re-initializing the parameters.

Figure 8 illustrates the joint distribution of \(\kappa (x_i)\), and \(\big \langle \mathbf {y}_i, \varvec{\mu }(x_i)\big \rangle \) at one point during a typical annealing process, which shows that the model is indeed approximately at equilibrium. To see the effect of \(\beta \) on the optimization of the parameters of \(\varvec{\mu }\) more clearly as well, we also consider the gradient of the loss in Eq. (18) with respect to \(\varvec{\mu }\). To maintain the unit-length constraint, we assume that \(\varvec{\mu }=\mathbf {z}/\Vert \mathbf {z}\Vert _2\). Thus, the gradient is given by

where we have written \(\mathbf {z},\kappa ,\mathbf {y}\) instead of \( \mathbf {z}(x),\kappa (x),\mathbf {y}(x)\) for brevity. The gradient vanishes when \(\mathbf {z}=\mathbf {y}\), and the optimum with respect to \(\varvec{\mu }\) does not depend on the precision \(\beta \). In practice, however, when we use gradient-descent to optimize the risk function, the magnitude of the gradient will be multiplied with \(W\big (\kappa \big )\), which clearly depends on \(\beta \). So while the optimum for \(\varvec{\mu }\) is still the same, we see that the precision influences the effective learning rate for the parameters of \(\varvec{\mu }\). At the beginning of the annealing schedule, the factor \(W(\kappa )\) is small, because the precision is small (see also Fig. 7), i.e. the parameters of \(\varvec{\mu }\) are less susceptible to the deviation from the target \(\mathbf {y}\). As the precision is gradually increased, the effective learning increases as well, and the gradient updates of \(\varvec{\mu }\) will push it stronger towards \(\mathbf {y}\).

To recapitulate, we have discussed how the precision parameter \(\beta \) constrains the average concentration of the FvM posterior during training, and how we can consistently, and efficiently obtain model parameters at different levels of precision. In the next section, we address the question of how to determine the optimal precision based on the noise in the data.

We have seen in the previous two sections, that the proposed entropy regularization effectively limits the concentration \(\kappa (x)\), however, it introduces the undetermined precision hyper-parameter \(\beta \).

A common strategy to determine hyper-parameters would be cross-validation with respect to the generalization error \(-\sum _i \langle \mathbf {y}_i,\varvec{\mu }(x_i)\rangle \) on a validation set. However, cross-validation of this kind does not provide a solution here, because we have shown in the last section that the optimum of the function \(\varvec{\mu }(x)\) is not affected by the precision (Eq. (22)).

One could object, that the generalization error does not entirely reflect the learned posterior, but only its mean direction, and that we should rather compute the expected generalization error \(\sum _i \rho _\beta (x_i, \mathbf {y}_i)\) withwhere we have defined \(p_\beta ^\text {FvM}:=p_{\hat{\varphi }_\beta }^\text {FvM}\), and \(\kappa _\beta \) analogously. Moreover, note again that \(\varvec{\mu }\) does not carry the precision subscript to indicate that it does not depend on \(\beta \). Even though the expected generalization error depends on the precision, it does not have an optimum for finite \(\beta \), as illustrated in Fig. 9. Instead, the minimum is always achieved at \(\beta \rightarrow \infty \), which corresponds to the well-known empirical risk minimizer. This argument shows that we can not determine the optimal width of the posterior by minimizing risk since it introduces a bias which underestimates uncertainty.

Instead, we need a criterion, which can assess the stability of the posterior distribution with respect to data fluctuations. For this purpose, we propose a method motivated by the information-theoretic framework of expected log-posterior agreement. It is applicable to any Gibbs posterior distribution, if we have access to repeated measurements, which are used to assess the noise level in the data, and to calibrate the precision \(\beta \) accordingly. Specifically, we require a validation set, which provides two independent realizations \(x_i^\prime ,x_i^{\prime \prime }\) for each measurement i, i.e. a set \(\{(x_i^\prime , x_i^{\prime \prime })\}_{i=1\dots n}\). In the context of spherical regression, we define the PA for one repeated measurement aswhere C(0) is the normalization of the uniform distribution as defined in Eq. (5). Performing the integration over all directions \(\mathbf {s}\), the PA readswhere we have written \(\kappa _\beta ^\prime =\kappa _\beta (x^\prime )\), etc. for brevity. The maximal agreement is realized between the following two limiting cases:

Indeed, if we increase \(\beta \) sufficiently high, the agreement integral drops below the value C(0) achieved in the uniform case, and the posterior agreement \(i_\beta (x^\prime ,x^{\prime \prime })\) becomes zero again. We note, that the PA also vanishes for \(\beta >0\), if either of the two measurements is uninformative, i.e. uniform by either \({\kappa ^\prime _\beta =0}\) or \({\kappa ^{\prime \prime }_\beta =0}\).

In Fig. 10 we illustrate the behavior of \(i_{\beta }\) when \(\kappa _\beta ^\prime ~=~\kappa _\beta ^{\prime \prime }~=~\beta \), showing clearly the discussed trade-off between low precision and high precision.

To determine the optimal precision based on a validation set, we compute the average over all n repeated measurementsand maximize it with respect to \(\beta \):Moreover, we can assign an interesting interpretation to the numeric values of \(i_{\hat{\beta }}\). Let us consider the solid angle \(\varOmega _p(\beta )\) centered on an arbitrary, but fixed \(\varvec{\mu }\), which contains p% of the probability mass of a general \(p^\text {FvM}(\mathbf {s}\mid \varvec{\mu }, \beta )\) distribution. We use it to partition the unit sphere into \(4\pi /\varOmega _p(\beta )\) effective, conic bins. In this sense, the precision \(\beta \) determines a quantization angle on the sphere. If we measure the number of effective, conic bits with the binary logarithm, we find thatas shown by the black line in Fig. 10.

This result is interesting, because it suggests that the value of \(i_{\beta }\) corresponds to the bit-rate of a noisy communication scenario, as described in the original, information-theoretic derivation by Buhmann (2010).

Besides, we can use it to define an intuitive scale for the directional precision by calibrating it with respect to \(\beta _\theta \), where \(\theta \) is the aperture angle of the solid angle \(\varOmega _{99.5}(\beta )\). More precisely, the relationship between \(\beta \) and \(\theta \) is given byUnless indicated otherwise, we scale all experimental precision values with respect to \(\beta _{15^\circ }=114.40\).

Given a measurement of DWI features \(\mathbf {X}\), and a corresponding reference tractogram \(\mathbf {T}\) as supervision information for training, our goal is to learn the posterior distribution of local streamline direction \({p^\text {trk}(\mathbf {y}\mid \varvec{x},\mathbf {y}^{in})}\) based on the entropy-regularized Gibbs free energy presented in Eq. (18). Note, that we have denoted \(\varvec{x}\in \mathbb {R}^p\) as the variable for the local diffusion data, and \(\mathbf {y}^{in}\in \mathbb {S}_2\) as the variable for the incoming fiber direction. Together, they constitute the input vector \(x=(\varvec{x}, \mathbf {y}^{in})\in \mathbb {R}^p\times \mathbb {S}_2\), which the Entrack posterior conditions on.

In the following, we discuss how to decompose the data set \((\mathbf {X},\mathbf {T})\) such that it can be used with the risk function introduced in Eq. (19). Specifically, we need to construct samples \(\big ((\varvec{x}_i,\mathbf {y}^{in}_i), \mathbf {y}_i\big )\) to capture the relationship between the target direction \(\mathbf {y}\) and the input (\(\varvec{x}\), \(\mathbf {y}^{in})\). We detail the corresponding sample generation process in Algorithm 2, and illustrate it in Fig. 11. With an accordingly generated training set \(\mathbb {X}\), we can use the entropy-regularized risk function from Eq. (19) to estimate the parameters of \(p_\beta ^\text {trk}:=p_{\hat{\varphi }_\beta }^\text {trk}\).

However, we also need to account for the inversion symmetry of DWI data, which makes it equivalent to traverse a streamline in forward, and backward direction. To ensure that the posterior can learn this symmetry from the data, i.e. \({p^\text {trk}(-\mathbf {y}\mid \varvec{x},-\mathbf {y}^{in}) = p^\text {trk}(\mathbf {y}\mid \varvec{x},\mathbf {y}^{in}) }\), we incorporate this invariance explicitly in the risk function by adding forwards \((u=+1)\), and backwards direction \((u=-1)\):where N refers to the number of sample streamline segments.

The trained FvM posterior \(p_\beta ^\text {trk}\) is employed by the iterative tracking algorithm as described by Algorithm 3.

To construct a streamline, we start from a seed point \({\mathbf {r}_1\in \mathcal {V}}\), and obtain the local DWI features \(\mathbf {X}(\mathbf {r}_1)\). Provided with a prior direction \(\mathbf {y}_1\in \mathbb {S}_{2}\), e.g. from a diffusion-tensor fit, we can establish the next point \(\mathbf {r}_2\) of the streamline by sampling a direction \(\mathbf {y}_2\) from \({ p_\beta ^\text {trk} \big ( \mathbf {y}\mid \mathbf {X}(\mathbf {r}_1),\mathbf {y}_1 \big ) }\), and setting \(\mathbf {r}_2=\mathbf {r}_1+\alpha \mathbf {y}_2\), with a step size \(\alpha \). This iteration repeats, until a termination criterion is met, such as a thresholds on fiber length, strength of the diffusion signal, fiber bending angle, or leaving a predefined region of interest (ROI). The corresponding streamline \(\mathbf {t}_i\) simply consists of the traversed points, i.e. \(\mathbf {t}_i=(\mathbf {r}_{i,1},\dots ,\mathbf {r}_{i,n_i})\).

To obtain a dense tractogram \(\mathbf {T}=\{\mathbf {t}_i\}_{i=1\dots n}\), we place n seed points within a specified ROI, e.g. within a white matter mask for whole-brain tractography.

In Sect. 3.5 we introduced the PA for general directional regression with the FvM posterior, and now we describe how we implement it for tractography to determine the optimal \(\beta \) for \(p^\text {trk}_\beta \).

Given two independent DWI measurements \(\mathbf {X}^\prime ,\mathbf {X}^{\prime \prime }\) of the same subject, we denote the corresponding tractograms, obtained with Algorithm 3 in conjunction with \(p^\text {trk}_\beta \), as \(\mathbf {T}_\beta ^\prime ,\mathbf {T}_\beta ^{\prime \prime }\). The tractograms carry the subscript \(\beta \), because they implicitly depend on the precision via the Entrack posterior \(p^\text {trk}_\beta \). We recall that the PA \(i_\beta (x^\prime ,x^{\prime \prime })\) from Eq. (25) depends on two measurements \(x^\prime ,x^{\prime \prime }\) of the same input, and the input to the Entrack posterior consists of two components, i.e. \(x=\bigl (\mathbf {X}(\mathbf {r}), \mathbf {y}^{in}\bigr )\). Given proper image registration between \(\mathbf {X}^\prime \) and \(\mathbf {X}^{\prime \prime }\), it is straightforward to match the repeated measurements of the DWI data by considering the same location, i.e. \(\mathbf {X}^\prime (\mathbf {r}),\mathbf {X}^{\prime \prime }(\mathbf {r})\).

To obtain \(\mathbf {y}^{in}(\mathbf {r})^\prime ,\mathbf {y}^{in}(\mathbf {r})^{\prime \prime }\) from the discrete streamlines of the two tractograms \(\mathbf {T}_\beta ^\prime ,\mathbf {T}_\beta ^{\prime \prime }\) we consider a small volume around the location \(\mathbf {r}\), and compute \(\mathbf {y}^{in}(\mathbf {r})^\prime ,\mathbf {y}^{in}(\mathbf {r})^{\prime \prime }\) based on the corresponding streamlines which pass through this voxel. Thus, the continuous measurement volume \({\mathcal {V}}\) is decomposed into little cuboids with voxel size \(a>0\), that are indexed by their discrete location inside the measurement volume, i.e. \({\mathbf {z}\in \mathcal {V}_a=\mathcal {V}\cap \{z\cdot a\mid z\in \mathbb {Z}^3\}}\). We refer to the volume of a voxel as \(v_{\mathbf {z}}=\{\mathbf {r}: \Vert \mathbf {r}-\mathbf {z}\Vert _{\infty }\le a/2\}\). Even though we can now compute \(\mathbf {y}^{in}(\mathbf {z}\mid \mathbf {T}) \propto \sum _{i,j}\mathbb {I}\{r_{i,j}\in v_\mathbf {z}\}\mathbf {y}_{i,j}\), we still need to take into account that independent streamline bundles may cross at the same voxel, and they must not be confused with each other.

Instead we need to consider each bundle b separately, and condition the local direction on the bundle, too, i.e. \(\mathbf {y}^{in}(\mathbf {z},b\mid \mathbf {T})\). More precisely, we consider a bundle b as a set of coherent streamlines, which are similar in the sense that the average pointwise distance is small for each pair of streamlines in a bundle. This way, we can partition a tractogram into a set of bundles \(B(\mathbf {T})=\{b_1,\dots ,b_k\}\) such that \(\forall i,j:\ b_i\cap b_j=\emptyset \ \wedge \ \bigcup _i b_i=\mathbf {T}\). We also refer to Garyfallidis et al. (2012) for more details about the practical grouping of tractograms into bundles. Using the partitioned tractogram B, we can compute the local streamline direction per-bundle, up to normalization to unit-length, aswith \(\mathbf {y}_{i,j} = (\mathbf {r}_{i,j}-\mathbf {r}_{i,j-1}) / \Vert \mathbf {r}_{i,j}-\mathbf {r}_{i,j-1}\Vert _2\). Essentially, we have decomposed the tractogram \(\mathbf {T}\) into the directions of its individual fiber bundles at each voxel, which is also known as fixel representation (Raffelt et al. 2017), referring to “a specific fiber bundle within a specific voxel”. Consequently, we can think of each tuple \((\mathbf {z},b)\) as the coordinates of one fixel. We define the posterior mean direction of such a fixel aswhere \(\varvec{\mu }\) is the mean direction of the Entrack posterior \(p_\beta ^\text {trk}\).

Lastly, when we compute the posterior concentration of a fixel, i.e \(\kappa _\beta (\mathbf {z},b\mid \mathbf {X},\mathbf {T}_\beta )\), we also need to take into account the number of streamlines which represent the summary direction \(\mathbf {y}^{in}(\mathbf {z},b\mid \mathbf {T}_\beta )\). Intuitively, we should be more certain about the summary direction, if it is represented by many fibers, i.e. the concentration should be increased³. Formally, the fixel concentration is scaled by the streamline density, i.e.where the streamline density is defined asIn particular, the fixel concentration \(\kappa _\beta (\mathbf {z},b\mid \mathbf {X},\mathbf {T}_\beta )\) is zero, i.e. the posterior doesn’t contain any information about the fixel direction, when we don’t observe any streamline. Putting everything together, we obtain the posterior agreement for one fixel, based on Eq. (25), aswhere we have defined \(\kappa _\beta ^\prime (\mathbf {z},b) := \kappa _\beta (\mathbf {z},b\mid \mathbf {X}^\prime , \mathbf {T}_\beta ^\prime )\), etc. for brevity. Consequently, the average PA over all fixels is given bywith the set of bundles that intersect a particular voxel denoted as \(B_\mathbf {z}=\{b\in B(\mathbf {T}_\beta ^\prime \cup \mathbf {T}_\beta ^{\prime \prime }): \exists \mathbf {t}\in b: \exists \mathbf {r}\in \mathbf {t}: \mathbf {r}\in v_\mathbf {z}\}\).

In the following we summarize the most important details about the DWI data and its preprocessing, i.e. how we obtain the DWI features \(\mathbf {X}\).

In this section, we discuss our implementation of the Entrack posterior \({p^\text {trk} \Big ( \mathbf {y} \mid \varvec{\mu }\big (\mathbf {X}(\mathbf {r}), \mathbf {y}^{in}\big ), \kappa \big (\mathbf {X}(\mathbf {r}), \mathbf {y}^{in}\big ) \Big ) }\), in particular the implementation of the functions \(\varvec{\mu },\kappa \).

While the general formulation supports a wide range of possible functions, we chose a deep neural network model due to its superior ability to extract patterns automatically (Goodfellow et al. 2015). Moreover, neural networks (NN) naturally support modular architectures, which allows us to readily formulate \(\varvec{\mu },\kappa \) in terms of two output modules \(\text {NN}_\mu ,\text {NN}_\kappa \), based on a shared NN module \(\text {NN}_z\):Specifically, each NN module is a series of fully-connected layers, as shown in Fig. 12, together with the detailed parameters. The inputs \(\mathbf {X}(\mathbf {r}), \mathbf {y}^{in}\) are both flattened, and concatenated to form a 408-dimensional input vector, i.e. 3 dimensions for \(\mathbf {y}^{in}\), and \(405=(3\times 3\times 3)\times 15\) dimensions for \(\mathbf {X}(\mathbf {r})\), which represents the 15 FOD coefficients (\(l=4\)) for each voxel in a \(3\times 3\times 3\) cube centered on the location \(\mathbf {z}=a[\mathbf {r}/a]\), where a is the voxel size.

To better understand which patterns the Entrack model has recognized in the training data, we perform a series of experiments on prototypical inputs.

In Fig. 15, we present the TM metrics as a function of the precision \(\beta \). As a major observation, the TM scores do not seem to suggest a consistent optimal precision. The VC-score increasingly saturates to a maximum value of 0.52 (\(\beta /\beta _{15^\circ }=1.58\)), while the maximum F1-score of 0.54 (\(\beta /\beta _{15^\circ }=0.42\)) marginally decreases by \(2.5\%\) over the same range. The VB-score toggles between 23 and 24, and remains stable otherwise. The IB-score is the least consistent, with a local minimum of 123 at \(\beta /\beta _{15^\circ }=0.66\), however its total variation over the range of \(\beta \) is also very small \(\approx 5\%\). This observation indicates, that besides the lack of a clear optimum with respect to the precision, the TM scores are also not very sensitive to \(\beta \), except the VC-score, which increases by 20%. The behavior of the VC score rather suggests a comparison with the generalization error, described in Eq. (23), which also saturates for \(\beta \rightarrow \infty \), and does not have an optimum at finite precision.

A common quantitative measure of how much the diffusion signal is confined along a single direction is the fractional anisotropy \(\text {FA}\in [0,1]\):⁵ It relates to the eccentricity of the diffusion ellipsoid, and it is 0 for an isotropic sphere, which is characteristic for voxels with ambiguous DWI measurements, whereas it is 1 for voxels with a diffusion ellipsoid clearly elongated in one direction. In Fig. 16, we show that the posterior concentration \(\kappa \big (\varvec{x}, \mathbf {y}^{in}\big )\) has indeed a strong correlation with \(\text {FA}(\varvec{x})\) \((r=0.42)\), which means that there is a strong link between model certainty, and the articularity of the diffusion signal.

To provide an absolute reference point for the presented TM-scores, we include an overview of the scores achieved by current tractography solutions based on supervised machine learning in table 1. Besides the row “ISMRM15 \(\varnothing \)”, which represents the average over all teams of the ISMRM15 challenge (ML and non-ML), we have divided the results into two groups. The first group, marked with an asterisk, represents results where the model was trained on the synthetic DWI phantom, and these data are also used for evaluation. The work of Neher et al. (2017) refers to this setting as in silico\(\rightarrow \)in silico, meaning that training fibers for these algorithms were obtained on the phantom data by another state-of-the-art algorithm. Even though the training fibers do not exactly correspond to the evaluation fibers, there exists a strong statistical dependence and we can not consider this setting as a valid generalization test. Instead, the model should be trained, and evaluated on two independent instances of (synthetic) DWI data. But as the ISMRM challenge provides only one instance, models should be trained on real DWI data, e.g. from the HCP. As expected, the results in this setting fall behind compared to the results in the in silico\(\rightarrow \)in silico setting, which should be considered as (overly) optimistic estimates of the generalization performance. Aside from this criticism, it is apparent that all algorithms fail to consistently outperform the competitors in the realistic in vivo\(\rightarrow \)in silico setting.

Instead, we can observe a strong trade-off between OL and VC/IB. On one hand, the work of Wegmayr (2018) achieves very good VC/IB (72%/57), but poor OL (16%), on the other hand Entrack (sample) and Neher et al. (2017) achieve much better OL (58% and 59%, respectively), but poorer VC (52% and 52%, respectively). The results for the Entrack model were obtained at \(\beta /\beta _{15^\circ }=1.58\). A similar trade-off is seen between the (sample)/(mean) variants, which refer to how the fiber directions are obtained from the posterior during streamline progression. At each tracking step, the (sample) variant draws a random direction from the posterior, while the (mean) variant always chooses the most likely direction. The Entrack (sample) method achieves superior bundle coverage compared to the (mean) method (OL 58% vs. 45%), but at the cost of more false-positives (IB 123 vs. 116). The same is true for the FvM model, which has the same architecture as the Entrack model, but is trained without entropy regularization, i.e. using the probabilistic regression objective in Eq. (12). The bundle coverage of the FvM (sample) model is also better than its (mean) variant (OL 53% vs. 48%), but also at the cost of more false-negative bundles (IB 154 vs. 112).

Moreover, we highlight that the TM scores support a model ranking Entrack (sample) \(\succ \) FvM (sample) \(\succ \) Detrack \(\succ \) Classifier (mean), which denotes the respective benefits of entropy regularization, a probabilistic loss, and a regression model. The Detrack model has the same neural network architecture as FvM and Entrack, but without the output for \(\kappa \), and it is trained with the standard negative cosine loss of Eq. (7). The Classifier model also shares the same neural network architecture as all the other models, but it has a softmax output over directions, and is trained by the usual cross-entropy loss for classification. We note, that each of the listed models should be understood as a module in the complete pipeline described by Algorithm 3. Such a pipeline is controlled by various other significant influence factors (training data, seed points, etc.) that are different in each case, thereby limiting comparability. We can only assert that the results of Classifier, Detrack, FvM, and Entrack have been conditioned on the same pipeline, so that their differences can be attributed indeed to the respective choices of the objective functions.

In addition to the evaluation metrics provided by the Tractometer, we present qualitative tractogram visualizations. In Fig. 17 we show an overview-section of the whole-brain tractogram obtained with the Entrack model on the ISMRM data, which should be compared to the ground-truth fibers in Fig. 1b.

Additionally, to facilitate a more detailed analysis, we have computed the voxel mask of the predicted corticospinal tract (CST), and visualize its overlap, overreach, and underreach with respect to the ground-truth bundle in Appendix G. Lastly, we demonstrate visualizations of the heteroscedastic uncertainty estimated by the Entrack posterior in Fig. 18. On one hand, we can visualize the spatial dependence of \(\kappa \big (\mathbf {X}(\mathbf {r}), \mathbf {y}^{in}\big )\) (Figs. 18a and b), and on the other hand, we can compute per-fiber statistics, such as the log-probability per streamline, i.e.shown in Figs. 18c and d. As we have discussed before, the concentration parameter measures the degree of certainty that the model encodes on the fiber direction at a given location. In Fig. 18a we can observe that the concentration/certainty is larger at the core of bundles than in the periphery, which agrees with the fact that the diffusion data is less ambiguous at the bundle cores than at the boundaries. Areas that are closer to the white matter boundary, such as bundle outlines, have lower concentration, because the diffusion signal is weaker in those areas. In particular, fiber end points exhibit the lowest concentrations, as they are located right at the white matter boundary, as shown in Fig. 18b.

In addition to per-point statistics, the per-fiber statistic \(\overline{\log p}\) can be used to automatically detect fiber outliers, as shown in Fig. 18c and d. Clearly, without marking fibers in comparison to the average log-likelihood, it is highly error-prone to discover such outliers visually in a tangled whole-brain tractogram. The illustrated implausible loop was found in the ISMRM ground-truth fibers, which are otherwise well prepared. This finding also underlines the difficulty of preparing high-quality reference standards for tractography. Lastly, we note that in contrast to e.g. curvature based outlier-detection, the Entrack model acts as a data-informed filter, i.e. it can recognize fibers, which are strongly bending, but supported by the diffusion data, whereas these fibers would be discarded by a curvature dominated filter.

In this section, we show the results for the optimal precision obtained with the PA criterion from Sect. 4.3, based on two independent DWI measurements of the HCP subject 917255. In Fig. 19, we show the measured values of the expected posterior agreement \({\mathcal {I}_\beta }\) from Eq. (40), and the expected generalization error \(\rho _\beta \) from Eq. (23). In contrast to the Tractometer scores, we observe a clear optimum of \({\mathcal {I}_\beta }\) at \(\hat{\beta }/\beta _{15^\circ }=0.41\). As anticipated by our discussion in Sect. 3.5, the generalization error \(\rho _\beta \) suggests \(\hat{\beta }\rightarrow \infty \) and thus fails to provide a finite estimate for the precision.

Furthermore, we are interested to explain the empirical PA with a phenomenological model of the formwhich depends only on the summary statistics \(\bar{\theta }, \bar{n}\). These are the average number of fibers per fixelwhere W(.) was introduced in Eq. (4a), and the average deviation between the fixel direction on the two instances:This description has one free parameter \(\lambda \), which essentially captures how fast the local variance of a fiber bundle increases when the precision is decreased. We refer to appendix B for more details about the origin of the parameter \(\lambda \). It is a joint property of the iterative tracking together with the local posterior, and the DWI data distribution. In particular, it can be considered as a measure of how fast bundles produced by a particular tracking algorithm, on a particular DWI source, tend to disintegrate when the precision is lowered, as illustrated by Fig. 20. In our case, using \(\lambda =10\) provides a good match between the measured PA, and the phenomenological model, as shown by the orange curve in Fig. 19. Its maximum value is \(i_{\hat{\beta }}=4.14\), which means, at the given noise level, we can contract the Entrack posterior up to a concentration, which is equivalent to the partition of the sphere into \(\approx 16=2^{i_{\hat{\beta }}}\) equally sized cones. A higher resolution can not be argued for, since it would increasingly reduce the agreement between the posteriors on the two measurements. This effect is not captured by the expected generalization error \(\rho _\beta \) (blue curve in Fig. 19), showing again that \(\rho _\beta \) is not an appropriate measure to determine a finite optimal precision, which is necessary to maintain the benefits of probabilistic models.