Generative diffeomorphic modelling of large MRI data sets for probabilistic template construction

Abstract

In this paper we present a hierarchical generative model of medical image data, which can capture simultaneously the variability of both signal intensity and anatomical shapes across large populations. Such a model has a direct application for learning average-shaped probabilistic tissue templates in a fully automated manner. While in principle the generality of the proposed Bayesian approach makes it suitable to address a wide range of medical image computing problems, our work focuses primarily on neuroimaging applications. In particular we validate the proposed method on both real and synthetic brain MR scans including the cervical cord and demonstrate that it yields accurate alignment of brain and spinal cord structures, as compared to state-of-the-art tools for medical image registration. At the same time we illustrate how the resulting tissue probability maps can readily be used to segment, bias correct and spatially normalise unseen data, which are all crucial pre-processing steps for MR imaging studies.

Highlights

• •We present a generative modelling framework to process large MRI data sets.
• •The proposed framework can serve to learn average-shaped tissue probability maps and empirical intensity priors.
• •We explore semi-supervised learning and variational inference schemes.
• •The method is validated against state-of-the-art tools using publicly available data.

Introduction

Over the past two decades neuroimaging techniques have been widely recognised has powerful non-invasive tools to answer neuroscientific questions regarding, for instance, the variability of neuroanatomy across different populations (Good et al., 2001, Ashburner and Friston, 2000), the role of structural changes in determining disease onset and progression (Fox et al., 2000, Chetelat et al., 2005, Brex et al., 2002), as well as the relationships between neural activations and mental functions (Friston et al., 1996) during common cognitive, somatosensory or motor tasks.

One of the main challenges, which is encountered in all neuroimaging studies, originates from the difficulty of mapping between different anatomical shapes. In particular, a fundamental problem arises from having to ensure that this mapping operation preserves topological properties and that it provides, not only anatomical, but also functional overlap between distinct instances of the same anatomical object (Brett et al., 2002).

This explains the rapid development of the discipline known as computational anatomy (Grenander and Miller, 1998), which aims to provide mathematically sound tools and algorithmic solutions to model high-dimensional anatomical shapes, with the ultimate goal of encoding, or accounting for, their variability. In this framework, the problem of anatomical mapping is most commonly addressed by assuming that a population of subjects can be fully described by means of an average-shaped template, together with a set of high-dimensional transformations, which, once applied to the template, generate the observed anatomical configurations (Thompson et al., 2000). In this setting, an ideal template is one that can accurately be matched to the observed data while minimising the geodesic distance between the template itself and each observation (Avants and Gee, 2004).

One of the practical advantages of this approach is that, once a template has been constructed, it can be used to model unseen data drawn form the same population, thus providing an ideal reference coordinate system for neuroimaging studies and for statistical testing of neuroscientific hypotheses.

In this paper we propose a general modelling scheme and a training algorithm, which, given a large cross-sectional data set of MR scans, can learn a set of average-shaped tissue probability maps, either in an unsupervised or semi-supervised manner. This is achieved by building a hierarchical generative model of MR data, where image intensities are captured using multivariate Gaussian mixture models, after diffeomorphic warping (Ashburner and Friston, 2011, Joshi et al., 2004) of a set of unknown probabilistic templates, which act as anatomical priors. In addition, intensity inhomogeneity artefacts are explicitly represented in our model, meaning that the input data does not need to be bias corrected prior to model fitting.

Our work builds on a number of state-of-the-art techniques, some of which were already explored and validated individually as part of our previous work (Ashburner and Friston, 2005, Ashburner and Friston, 2011, Blaiotta et al., 2016). In particular, we rely on the variational Bayes expectation-maximisation (VBEM) algorithm for image segmentation proposed in Blaiotta et al. (2016), which is a robust generalisation of the segmentation method implemented in the SPM12 software (Ashburner and Friston, 2005). Additionally, we exploit the diffeomorphic image registration framework described in Ashburner and Friston (2011), which is an accurate, fast converging and memory efficient strategy to align complex anatomical shapes using diffeomorphic transformations.

To the best of our knowledge, the particular mathematical formulation that we adopt to combine such modelling techniques has never been adopted before. The resulting approach enables processing simultaneously a large number of MR scans in a groupwise fashion and particularly it allows the tasks of image segmentation, image registration, bias correction and atlas construction to be solved by optimising a single objective function, with one iterative algorithm. This is in contrast to a commonly adopted approach to mathematical modelling, which involves a pipeline of multiple model fitting strategies that solve sub-problems sequentially, without taking into account their circular dependencies. Our strategy instead, addresses the fact that learning anatomical priors requires collecting accurate segmentations from a training population, while such segmentations can only be obtained if suitable priors are available. Thus, the two modelling problems are solved more effectively within an integrated framework, rather than independently.

Related work was done by Bhatia et al. (2007), however their method relies on classical point estimation techniques to perform image segmentation, as opposed to our variational approach, which instead allows estimating full posteriors on the intensity distribution parameters. Additionally we incorporate bias correction and also explore semi-supervised learning, as opposed to the fully unsupervised scheme adopted in their work. The work presented in Lorenzen et al. (2004) is also along a similar line but it only address the problem of Bayesian diffeomorphic template construction given a set of pre-computed segmentations, without embedding image segmentation and atlas construction in a single mathematical model of the data.

The main aim of this paper is to demonstrate the methodological validity of the proposed approach, which, for this purpose, has been thoroughly tested on both synthetic and real MR neuroimaging data sets. However, the method is a general one and, since it was not tuned or optimised for a particular type of imaging data, we believe that it holds a potential for application in a wider range of medical image computing problems, some of which we aim to explore as part of our future work.

Methods

Let us consider a population of M subjects belonging to a homogeneous group, from an anatomical point of view, and let us assume that D image volumes of different contrast are available for each subject.

From a generative perspective, the image intensities $\mathbf{X}={\left\{{\mathbf{X}}_i\right\}}_{i=1,\dots ,M}$, which constitute the observed data, can be thought of as being generated by sampling from D-dimensional Gaussian mixture probability distributions, after non-linear warping of a probabilistic anatomical atlas (Evans et al., 1994). The use of Gaussian mixture models to capture the probability density function of MR data is a well established approach (Ashburner and Friston, 2005, Greenspan et al., 2006, Zhang et al., 2001). For an extensive review of such methods see Balafar (2014).

The probabilistic atlas carries prior anatomical knowledge, in the form of average shaped tissue probability maps. From a mathematical modelling point of view, the atlas encodes local mixing proportions ${\text{Θ}}_{\pi }={\left\{{\mathit{\pi }}_j\right\}}_{j=1,\dots ,N_{\pi }}$ of the mixture model, with j being an index set over the $N_{\pi }$ template voxels (for a list of all the mathematical symbols used in this section see Table 1). Each vector ${\mathit{\pi }}_j$ has K elements, which indicate the prior probability of voxel j belonging to one of K tissue classes. These spatially varying mixing proportions can also be thought of as the coefficients used to parametrise K continuous scalar functions of space ${\left\{{\pi }_k\right\}}_{k=1,\dots ,K}$.

Table 1.

List of mathematical symbols used in this paper.

Symbol	Meaning
${\mathbf{x}}_{ij}$	Observed image intensity at voxel j for subject i.
${\mathbf{z}}_{ij}$	Vector of latent class membership probabilities.
${\mathit{\pi }}_j$	Tissue priors at voxel j.
${\mathit{\mu }}_{ik}$	Mean intensity of class k for subject i.
${\mathbf{\Sigma }}_{ik}$	Covariance of intensities for class k and subject i.
${\mathit{W}}_{0k}$	Scale matrix of Wishart prior distribution on ${\mathbf{\Lambda }}_k={\left({\mathbf{\Sigma }}_k\right)}^{-1}\text{.}$
${\nu }_{0k}$	Degrees of freedom of Wishart prior distribution on ${\mathbf{\Lambda }}_k$.
${\mathit{m}}_{0k}$	Mean of Gaussian prior distribution over ${\mathit{\mu }}_k$
${\beta }_{0k}$	Scaling hyperparameter of Gaussian prior distribution over ${\mathit{\mu }}_k$
${\alpha }_0$	Hyperparameter governing the Dirichlet prior on $\mathit{\pi }$.
${\text{Θ}}_{\beta }$	Bias field parameters.
${\mathit{\mu }}_{\beta }$	Prior mean of bias parameters.
${\mathbf{\Sigma }}_{\beta }$	Prior covariance matrix of bias parameters.
${\text{Θ}}_a$	Affine transformation parameters.
${\mathit{\mu }}_a$	Prior mean of affine transformation parameters.
${\mathbf{\Sigma }}_a$	Prior covariance matrix of affine transformation parameters.
${\mathit{w}}_i$	Weights for rescaling the tissue priors.
${\mathbf{u}}_{ij}$	Initial velocity at voxel j for subject i.
${\mathbf{L}}_u$	Differential operator to compute penalty on ${\mathbf{u}}_i$.

Tissue priors

Each image voxel, $j\in \left\{1,\dots ,N_i\right\}$¹ for each subject $i\in \left\{1,\dots ,M\right\}$ is considered as being drawn from K possible tissue classes. The following prior latent variable model defines the probability of finding tissue type k, at a specific location j (i.e. centre of voxel j), in image i, prior to observing the corresponding image intensity signal

$$p(z_{ijk}=1|{\text{Θ}}_{\pi },{\text{Θ}}_w,{\text{Θ}}_u)=\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_i\left({\mathbf{y}}_j\right)\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_i\left({\mathbf{y}}_j\right)\right)},$$ (1)

or equivalently

$$p\left({\mathbf{z}}_{ij}|{\text{Θ}}_{\pi },{\text{Θ}}_w,{\text{Θ}}_u\right)=\prod _{k=1}^K{\left(\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_i\left({\mathbf{y}}_j\right)\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_i\left({\mathbf{y}}_j\right)\right)}\right)}^{z_{ijk}}\text{.}$$ (2)

Class memberships, for each subject and each voxel, are encoded in the latent variable ${\mathbf{z}}_{ij}$, using a one-of-K scheme (i.e. ${\mathbf{z}}_{ij}$ is K-dimensional vector with all elements equal to zero except for one, which is equal to one). In particular, $z_{ijk}\in \left\{\mathrm{0,1}\right\}$ is equal to one if voxel j of image i belongs to tissue class k and zero otherwise, which also explains the equivalence of (1), (2). ${\left\{{\pi }_k\right\}}_{k=1,\dots ,K}$ are continuous scalar functions ${\pi }_k:{\text{Ω}}_{\pi }\to \mathrm{ℝ}$, defined on the template domain ${\text{Ω}}_{\pi }$. Such functions are common across the entire population, which satisfy the constraint

$$\sum _{k=1}^K{\pi }_k\left(\mathit{y}\right)=1,\forall \mathit{y}\in {\text{Ω}}_{\pi }\subset {\mathrm{ℝ}}^3,$$ (3)

with $\mathit{y}$ being a continuous coordinate vector field. Global weights ${\text{Θ}}_w={\left\{{\mathit{w}}_i\right\}}_{i=1,\dots ,M}$, with ${\mathit{w}}_i\in {\mathrm{ℝ}}^K$, are introduced to further compensate for individual differences in tissue composition.

In equation (1), ${\mathit{\xi }}_i$ denotes a generic spatial transformation, parametrised by ${\text{Θ}}_u$, which allows projecting prior anatomical information onto individual data, with ${\mathit{\xi }}_i:{\text{Ω}}_i\to {\text{Ω}}_{\pi }$ being a continuous mapping from the domain ${\text{Ω}}_i\subset {\mathrm{ℝ}}^3$ of image i, into the space of the tissue priors ${\text{Ω}}_{\pi }\subset {\mathrm{ℝ}}^3$. In this work we adopt a large deformation diffeomorphic model (Ashburner and Friston, 2011, Beg et al., 2005, Trouvé, 1998), where, as explained in the following subsection, the diffeomorphisms are parametrised by means of an initial velocity field, denoted by $\mathit{u}$.

Since digital image data is a discrete signal, defined on a tridimensional voxel grid, each mapping ${\mathit{\xi }}_i$ needs to be discretised as well, via sampling at the centre of every voxel $j\in \left\{1,\dots ,N_i\right\}$, to give the discrete mapping ${\left\{{\mathit{\xi }}_i\left({\mathbf{y}}_j\right)\right\}}_{j=1,\dots ,N_i}$ that appears in (1).

As opposed to the modelling approach described in Ashburner and Friston (2005) and Blaiotta et al. (2016), where the tissue priors were considered as fixed and known a priori quantities, here the tissue probability maps are treated as random variables, whose point estimates or full posteriors can be inferred via model fitting (Bhatia et al., 2007, Ribbens et al., 2014).

For this purpose, a finite dimensional parametrisation of the continuous functions ${\left\{{\pi }_k\right\}}_{k=1,\dots ,K}$ needs to be defined. Typically, whenever a continuous function needs to be reconstructed from a finite discrete sequence, it is possible to formulate the problem as an interpolation that makes use of a finite set of coefficients and continuous basis functions. Since the priors ${\left\{{\pi }_k\right\}}_{k=1,\dots ,K}$ are bounded to take values in the interval $\left[\mathrm{0,1}\right]$ on the entire domain ${\text{Ω}}_{\pi }$ (see equation (3)), not all basis functions are well suited here. For this reason we use first degree B-splines, which ensure that the tissue priors satisfy the above mentioned constraint, while being quite a computationally efficient choice compared to higher order interpolation methods. The coefficients used to parametrise the tissue priors belong to the discrete set ${\text{Θ}}_{\pi }={\left\{{\mathit{\pi }}_j\right\}}_{j=1,\dots ,N_{\pi }}$ of K-dimensional vectors, with

$$\sum _{k=1}^K{\pi }_{jk}=1,\forall j\in \left\{1,\dots ,N_{\pi }\right\}\text{.}$$ (4)

and $N_{\pi }$ being the number of template voxels. Such coefficients can be learned directly from the data, as it will be shown in the following section.

Additionally, prior distributions on the parameters ${\left\{{\mathit{\pi }}_j\right\}}_{j=1,\dots ,N}$ can be introduced (Bishop, 2006). Dirichlet priors are the most convenient choice here, since they are conjugate to multinomial forms of the type in (2), and they can be expressed as

$$p\left({\mathit{\pi }}_j\right)=\text{Dir}\left({\mathit{\pi }}_j|{\mathit{\alpha }}_0\right)=C\left({\mathit{\alpha }}_0\right)\prod _{k=1}^K{\pi }_{jk}^{{\alpha }_k-1},$$ (5)

where the normalising constant is given by

$$C\left({\mathit{\alpha }}_0\right)=\frac{\text{Γ}\left(\overline{\alpha }\right)}{\text{Γ}\left({\alpha }_1\right)\dots \text{Γ}\left({\alpha }_k\right)},$$ (6)

with $\text{Γ}(\cdot )$ being the gamma function and

$$\overline{\alpha }=\sum _{k=1}^K{\alpha }_k\text{.}$$ (7)

Diffeomorphic image registration

As anticipated in the previous sections, the generative interpretation of imaging data that this work relies on involves warping an unknown, average-shaped atlas to match a series of individual scans.

Such a problem, that is to say template matching via non-rigid registration, has been largely explored in medical imaging, mainly for solving image segmentation or structural labeling problems, in an automated fashion (Ashburner and Friston, 2005, Shen and Davatzikos, 2004, Christensen, 1999, Chui et al., 2001, Bajcsy et al., 1983, Iglesias et al., 2012, Pluta et al., 2009, Warfield et al., 1999, Khan et al., 2008, Bowden et al., 1998).

Indeed, the modelling of spatial mappings between different anatomies can be approached in a variety of manners, depending on the adopted model of shape and on the objective function (i.e. similarity metric and regularisation) that the optimisation is based on, thus leading to a variety of algorithms with remarkably different properties (Penney et al., 1998, Denton et al., 1999, Klein et al., 2009).

The work presented here is formulated according to the Large Deformation Diffeomorphic Metric Mapping (LDDMM) framework (Younes, 2010), where the transformations mapping between the source images and the target image are assumed to belong to a Riemannian manifold ² of diffeomorphisms (Ashburner, 2007). A diffeomorphism $\varphi :\text{Ω}\to \text{Ω}$ is a smooth differentiable map (with a smooth differentiable inverse ${\varphi }^{-1}$) defined on a compact, simply connected domain $\text{Ω}\subset {\mathrm{ℝ}}^3$.

One way of constructing transformations belonging to the diffeomorphic group $\text{Diff}\left(\text{Ω}\right)$ is to solve the following non-stationary transport equation (Joshi and Miller, 2000)

$$\frac{\text{d}}{\text{d}t}\varphi \left(\mathit{y},t\right)=\mathit{u}\left(\varphi \left(\mathit{y},t\right),t\right),\varphi \left(\mathit{y},0\right)=\mathit{y},t\in \left[\mathrm{0,1}\right],$$ (8)

where $\mathit{u}\left(\varphi \left(\mathit{y},t\right),t\right)\in \mathrm{\mathscr{H}}$ is a time dependent, smooth velocity vector field, in the Hilbert space³ $\mathrm{\mathscr{H}}$.

The initial map, at $t=0$, is equal to the identity transform $\varphi \left(\mathit{y},0\right)=\mathit{y}$, while the final map, endpoint of the flow of the velocity field $\mathit{u}$, can be computed by integration on the unitary time interval $t\in \left[\mathrm{0,1}\right]$ (Beg et al., 2005).

$$\varphi \left(\mathit{y},1\right)={\int }_0^1\mathit{u}\left(\varphi \left(\mathit{y},t\right),t\right)\text{d}t+\varphi \left(\mathit{y},0\right)\text{.}$$ (12)

A procedure known as geodesic shooting (Miller et al., 2006, Ashburner and Friston, 2011, Allassonnière et al., 2005, Vialard et al., 2012, Beg and Khan, 2006) is applied, within the work presented here, to compute the final diffeomorphism given the initial map (i.e. the identity transform) and the velocity field at $t=0$. Such a procedure exploits the principle of conservation of momentum (Younes et al., 2009), which is given by ${\mathbf{L}}_{\mathbf{u}}^{†}{\mathbf{L}}_{\mathbf{u}}{\mathbf{u}}_t$, with ${\mathbf{L}}_{\mathbf{u}}^{†}$ being the adjoint of the differential operator ${\mathbf{L}}_{\mathbf{u}}$, to integrate the dynamical system governed by (8) without having to compute and store an entire time series of velocity fields. The implementation adopted here relies on the work presented in Ashburner and Friston (2011).

A diffeomorphic path $\varphi$ is not only differentiable, but also guaranteed to be a one-to-one mapping. Such a quality is highly desirable for finding morphological and functional correspondences between different anatomies without introducing tears or foldings, which would violate the conditions for topology preservation (Christensen, 1999). Additionally, the diffeomorphic framework provides metrics to quantitatively evaluate distances between anatomies or shapes. It should also be noted that diffeomorphisms are locally analogous to affine transformations (Avants et al., 2006).

In practice, finding an optimal diffeomorphic transformation to equation a pair, or a group, of images involves optimising an objective function (e.g. minimising a cost function), in the space $\mathrm{\mathscr{H}}$ of smooth velocity vector fields defined on the domain $\text{Ω}$. The required smoothness is enforced by constructing the norm on the space $\mathrm{\mathscr{H}}$ through a differential operator ${\mathbf{L}}_{\mathbf{u}}$ (Beg et al., 2005), such that a quantitative measure of smoothness can be obtained via

$$\mathrm{ℛ}\left(\mathbf{u}\right)=\left|\right|{\mathbf{L}}_{\mathbf{u}}\mathbf{u}|{|}_{L^2}^2,$$ (13)

where $\mathbf{u}$ is a discretised version of $\mathit{u}$.

The form of the cost function will depend on how the observed data is modelled. For the work presented here, groupwise equationment is achieved via maximisation of the following variational objective function

$$\begin{array}{l}\mathcal{E}\left({\text{Θ}}_u\right)={\mathbb{E}}_{\mathbf{Z}}\left[\text{log}p\left(\overline{\mathbf{Z}}|{\text{Θ}}_{\pi },{\text{Θ}}_w,{\text{Θ}}_u\right)\right]+\text{log}p\left({\text{Θ}}_u\right)+\text{const}\\ =\sum _{i=1}^M\sum _{j=1}^{N_i}\sum _{k=1}^K{\gamma }_{ijk}\text{log}\left(\frac{w_{ik}{\pi }_k\left({\varphi }_i\left({\mathbf{y}}_j\right)\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\varphi }_i\left({\mathbf{y}}_j\right)\right)}\right)\\ -\frac{1}{2}\sum _{i=1}^M\left|\right|{\mathbf{L}}_{\mathbf{u}}{\mathbf{u}}_i|{|}_{L^2}^2+\text{const},\end{array}$$ (14)

where $\mathbf{Z}={\left\{{\mathbf{Z}}_i\right\}}_{i=1,\dots ,M}$ is the set of latent variables across the entire population, ${\left\{{\mathit{\gamma }}_{ij}\right\}}_{i,j}={\{\mathbb{E}\left[{\mathbf{z}}_{ij}\right]\}}_{i,j}$ are K-dimensional vectors encoding the posterior probabilities of each voxel belonging to the K tissue types. The coordinate mappings ${\left\{{\varphi }_i\right\}}_{i=1,\dots ,M}$ are encoded in the parameter set ${\text{Θ}}_u$, which consists of M vectors of coefficients ${\left\{{\mathbf{u}}_i\right\}}_{i=1,\dots ,M}$, containing $3\times N_i$ elements each. Such coefficients can be used to construct continuous initial velocity fields via trilinear, or higher order, interpolation.

The posterior membership probabilities ${\left\{{\mathit{\gamma }}_{ij}\right\}}_{i,j}$ that appear in (14) can be computed by combining the prior latent variable model introduced in 2.1 with a class conditional likelihood model of image intensities, which will be described in subsection. In such a case, learning posterior label probabilities can be addressed as a standard mixture distribution inference problem, which can be conveniently solved using the expectation-maximisation algorithm or its variational extensions (Bishop, 2006, Blaiotta et al., 2016), thus leading to a fully unsupervised learning scheme.

Alternatively, when manual labels are available, binary posterior class probabilities can be derived directly from such categorical annotations, without performing inference from the observed image intensity data. In particular, if all input data has been manually labelled, then the resulting algorithm would implement a fully supervised learning strategy. Instead, if only some of the data has associated training labels, a hybrid approach can be adopted, which would fall into the category of semisupervised learning (Chapelle et al., 2006, Filipovych et al., 2011).

Essentially these three approaches, namely unsupervised, semisupervised and fully supervised, differ in the relative ratio between the number of labelled and unlabelled voxels. Even so, they can all be framed within an expectation-maximisation setting, where depending on whether a voxel is labelled or not, the E-step is performed differently. Specifically, if ${\mathbf{x}}_{ij}$ is an unlabelled observation

$$\begin{array}{l}{\gamma }_{ijk}=p(z_{ijk}=1|{\mathbf{x}}_{ij},\text{Θ})\\ =\frac{p\left({\mathbf{x}}_{ij}\right|z_{ijk}=1,\text{Θ})p(z_{ijk}=1|\text{Θ})}{{\sum }_{c=1}^{K}p\left({\mathbf{x}}_{ij}\right|z_{ijc}=1,\text{Θ})p(z_{ijc}=1|\text{Θ})},\end{array}$$ (15)

with $\text{Θ}$ being the current estimate of model parameters at each iteration of the algorithm. Instead, if ${\mathbf{x}}_{ij}$ has been manually labelled

$${\gamma }_{ijk}=p(z_{ijk}=1|l_{ij})=\{\begin{array}{cc}1,& \text{if }{l}_{ij}=k\\ 0,& \text{if }{l}_{ij}\ne k\end{array}$$ (16)

where $l_{ij}$ is a categorical manual label assigned to voxel j in image i.

Finally, it is also possible to take into account the uncertainty inherent in the process of manual rating. In such a case, the actual posterior probabilities, for labelled observations, can be computed by making use of the categorical output of manual labelling together with an estimate of the rater sensitivity and with a generative intensity model.

Making use of Bayes rule, this gives

$$\begin{array}{l}{\gamma }_{ijk}=p(z_{ijk}=1|{\mathbf{x}}_{ij},\text{Θ},l_{ij})\\ =\frac{p\left({\mathbf{x}}_{ij}\right|z_{ijk}=1,\text{Θ})p(z_{ijk}=1|\text{Θ})p(z_{ijk}=1|l_{ij})}{{\sum }_{c=1}^{K}p\left({\mathbf{x}}_{ij}\right|z_{ijc}=1,\text{Θ})p(z_{ijc}=1|\text{Θ})p(z_{ijc}=1|l_{ij})},\end{array}$$ (17)

where $p(z_{ijk}=1|l_{ij})$ indicates the probability of voxel j in image i belonging to class k, given the manual label attributed to the same voxel.

A simple model for this, is

$$p(z_{ijk}=1|l_{ij})=\{\begin{array}{cc}{\zeta }_l,& \text{if }{l}_{ij}=k\\ \frac{1-{\zeta }_l}{K-1},& \text{if }{l}_{ij}\ne k\end{array}$$ (18)

where ${\zeta }_l$ is the sensitivity of the rater that generated the set of labels ${\left\{l_{ij}\right\}}_{j=1,\dots ,N}$ for image i. The problem of how to evaluate the performance of a manual or automated rater is not addressed here. For instance, a probabilistic scheme, which has been widely used to assess segmentation performance in medical imaging, is presented in Warfield et al. (2004).

In the remainder of this paper we will be focusing on the semi-supervised approach and we will consider the general case where only some voxels in few images have been manually labelled by adopting the model in (17) with a fixed sensitivity parameter ${\zeta }_l$ equal to 0.9.

Combining diffeomorphic with affine registration

Anatomical shapes are very high dimensional objects. The diffeomorphic model described in the previous section, which is encoded using $3\times N_i$ free parameters, can account for a significant amount of shape variability in the observed data. Nevertheless, it is still convenient, mainly for computational reasons, to combine such a local, high dimensional shape model with global, lower dimensional transformations, such as rigid body or affine transforms. In fact, by beginning to solve the registration problem from the coarsest deformation components (e.g. rigid body or affine), it is possible to ensure that the subsequent diffeomorphic registration starts from a good initial estimate of image alignment, that is to say closer to the desired global optimum (Lester and Arridge, 1999). This makes the optimisation problem faster to solve and at the same time it reduces significantly the chance of registration failure (Modersitzki, 2004). It is relatively common for non-linear registration algorithms to perform poorly in the presence of a large translational or size mismatch between the reference and the target images (Jenkinson and Smith, 2001).

A possible parametrisation that combines affine and diffeomorphic transformations is

$${\mathit{\xi }}_i\left(\mathit{y}\right)={\mathbf{T}}_i{\varphi }_i\left(\mathit{y}\right)+{\mathbf{t}}_i,\forall \mathit{y}\in {\text{Ω}}_i,$$ (19)

where ${\mathit{\xi }}_i\left(\mathit{y}\right)$ is the resulting mapping from image of subject i into the template space. Such a mapping is obtained by affine transforming the diffeomorphic deformation field ${\varphi }_i$. The transformation matrix ${\mathbf{T}}_i$ encodes nine degrees of freedom (rotation, zooming and shearing) and is computed via an exponential map ${\mathbf{T}}_i=\text{exp}\left({\mathbf{Q}}_i\left({\mathbf{a}}_i\right)\right)$ with ${\mathbf{Q}}_i\left({\mathbf{a}}_i\right)\in \mathfrak{g}\mathfrak{a}\left(3\right)$, where $\mathfrak{g}\mathfrak{a}\left(3\right)$ is the Lie algebra for the affine group in three dimension $GA\left(3\right)$ and ${\mathbf{a}}_i$ is a vector of nine parameters (Ashburner and Ridgway, 2013). Translations are modelled by the vector ${\mathbf{t}}_i\in {\mathrm{ℝ}}^3$. The entire set of affine parameters is denoted as ${\text{Θ}}_a={\left\{{\mathbf{a}}_i,{\mathbf{t}}_i\right\}}_{i=1,\dots ,M}$.

Intensity model

From a general probabilistic perspective, classification of tissue types based on MR signal intensities requires a model of the observed data that is capable of capturing the probability of occurrence of each signal sample value ${\mathbf{x}}_{ij}$, provided that the true labels are known. In other words, the problem breaks down into defining suitable conditional probabilities $p\left({\mathbf{x}}_{ij}\right|z_{ijk}=1)$, for each $k=\left\{1,\dots ,K\right\}$ and then applying Bayes rule to infer the posterior class probabilities.

In the model adopted here, image intensity distributions are represented as Gaussian mixtures, with the unknown mean ${\mathit{\mu }}_{ik}$ and covariance matrix ${\mathbf{\Sigma }}_{ik}$ of each Gaussian component k, for subject i, being governed by Gaussian-Wishart priors (Bishop, 2006, Blaiotta et al., 2016).

Correction of intensity inhomogeneities is also performed within the same modelling framework and it involves multiplying the uncorrected intensities of each image volume by a bias field, which is modelled as the exponential of a weighted sum of discrete cosine transform basis functions (Styner et al., 2000, Ashburner and Friston, 2005). Such an approach is conceptually equivalent to scaling the probability distributions of all Gaussian components by a local scale parameter, which is the bias itself, such that

$$p\left({\mathbf{x}}_{ij}|z_{ijk},{\mathit{\mu }}_{ik},{\mathbf{\Sigma }}_{ik},{\text{Θ}}_{\beta }\right)=\mathcal{N}\left({\mathbf{x}}_{ij}|{\hat{\mathit{\mu }}}_{ik},{\hat{\mathbf{\Sigma }}}_{ik}\right),$$ (20)

with

$$\begin{array}{l}{\hat{\mathit{\mu }}}_{ik}={\left(\text{diag}\left({\mathbf{b}}_{ij}\right)\right)}^{-1}{\mathit{\mu }}_{ik},\\ {\hat{\mathbf{\Sigma }}}_{ik}={\left(\text{diag}\left({\mathbf{b}}_{ij}\right)\right)}^{-1}{\mathbf{\Sigma }}_{ik}{\left(\text{diag}\left({\mathbf{b}}_{ij}\right)\right)}^{-1},\end{array}$$ (21)

where $\mathcal{N}$ indicates a Normal distribution, ${\text{Θ}}_{\beta }$ denotes the set of bias field parameters and ${\mathbf{b}}_{ij}$ is a D-dimensional vector representing the bias at voxel j of subject i for each of the D imaging modalities. The result in (20) is independent from the particular parametrisation of ${\mathbf{b}}_{ij}$, as long as the bias is multiplicative, and it can be easily proven by applying the change of variable ${\mathbf{x}}_{ij}={\mathbf{c}}_{ij}/{\mathbf{b}}_{ij}$, with ${\mathbf{c}}_{ij}$ being the corrected image intensity

$$\begin{array}{l}p\left({\mathbf{x}}_{ij}|z_{ijk},{\mathit{\mu }}_{ik},{\mathbf{\Sigma }}_{ik},{\text{Θ}}_{\beta }\right)=\\ \text{det}\left(\text{diag}\left({\mathbf{b}}_{ij}\right)\right)p\left({\mathbf{c}}_{ij}|z_{ijk},{\mathit{\mu }}_{ik},{\mathbf{\Sigma }}_{ik},{\text{Θ}}_{\beta }\right)=\\ \text{det}\left(\text{diag}\left({\mathbf{b}}_{ij}\right)\right)\mathcal{N}\left({\mathbf{c}}_{ij}|{\mathit{\mu }}_{ik},{\mathbf{\Sigma }}_{ik}\right)=\mathcal{N}\left({\mathbf{x}}_{ij}|{\hat{\mathit{\mu }}}_{ik},{\hat{\mathbf{\Sigma }}}_{ik}\right)\text{.}\end{array}$$ (22)

Graphical model

A graphical representation of the model adopted in this paper is depicted in Fig. 1, while a legend of the symbols used to indicate the different variables can be found in Table 1.

Fig. 1.

Graphical representation of the model adopted in this paper. Observed variables $\left\{{\mathbf{x}}_{ij}\right\}$ are represented by a filled circle. Latent variables $\left\{{\mathbf{z}}_{ij}\right\}$ as well as model parameters are depicted as unfilled circles. Blue solid dots correspond to fixed hyperparameters. The so called plate notation is adopted to indicated repeated variables. Symbols referring to all variables and parameters are listed in Table 1.

Given such a model, it is possible to define the following variational objective function $\mathrm{ℒ}$, which constitutes a lower bound on the logarithm of the marginal joint probability $p\left(\overline{\mathbf{X}},{\text{Θ}}_{\beta },{\text{Θ}}_a,{\text{Θ}}_u,{\text{Θ}}_{\pi }|{\text{Θ}}_w\right)$, such that

$$\text{log}p\left(\overline{\mathbf{X}},{\text{Θ}}_{\beta },{\text{Θ}}_a,{\text{Θ}}_u,{\text{Θ}}_{\pi }|{\text{Θ}}_w\right)\ge \mathrm{ℒ}$$ (23)

and

$$\mathrm{ℒ}={\mathbb{E}}_{\mathbf{Z},{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[\text{log}p\left(\mathbf{X}|\mathbf{Z},{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}},{\text{Θ}}_{\beta }\right)\right]+{\mathbb{E}}_{\mathbf{Z}}\left[\text{log}p\left(\mathbf{Z}|{\text{Θ}}_{\pi },{\text{Θ}}_w,{\text{Θ}}_u,{\text{Θ}}_a\right)\right]+{\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[\text{log}p\left({\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)\right]+\text{log}p\left({\text{Θ}}_{\pi }\right)+\text{log}p\left({\text{Θ}}_{\beta }\right)+\text{log}p\left({\text{Θ}}_a\right)+\text{log}p\left({\text{Θ}}_u\right)-{\mathbb{E}}_{\mathbf{Z}}\left[\text{log}q\left(\mathbf{Z}\right)\right]-{\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[\text{log}q\left({\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)\right],$$ (24)

where the expectations indicated as ${\mathbb{E}}_{\mathbf{Z}}$ and ${\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}$ are computed with respect to variational posterior distributions (Blaiotta et al., 2016, Bishop, 2006), denoted by $q(\cdot )$, on the latent variables $\overline{\mathbf{Z}}$ and on the Gaussian means and covariances $\left\{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right\}$, respectively. Optimisation of $\mathrm{ℒ}$, which provides optimal parameter and hyperparameter estimates, will be discussed in the following section.

Model fitting

The model described in the previous section can be fit to data sets of MR images by combining a variational expectation-maximisation (VBEM) algorithm with gradient based numerical optimisation techniques.

Indeed, the VBEM algorithm described in Blaiotta et al. (2016) is well-suited for addressing the model estimation problem discussed here, since it allows learning posterior distributions on the Gaussian mixture parameters, under the assumption that $q\left(\overline{\mathbf{Z}},{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)$ factorises as $q\left(\overline{\mathbf{Z}}\right)q\left({\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)$ (Bishop, 2006), and at the same time it is able to transfer the information encoded in such posteriors by estimating empirical intensity priors for each tissue type. As shown in our previous work (Blaiotta et al., 2016), this approach has several advantages over maximum likelihood estimation, including lower vulnerability to overfitting, faster convergence and higher robustness against misregistration, which inevitably occurs in the early iterations of any image registration algorithm.

Additionally, the algorithm proposed in this paper loops over all subjects in the population and, for each subject, it iterates over estimating the Gaussian-Wishart posteriors, the bias field parameters, the affine parameters and the initial velocities, which are all treated as conditional optimisations. Subsequently the tissue probability maps and the intensity priors are updated and the whole cycle is repeated until convergence.

Estimation of the bias field parameters ${\text{Θ}}_{\beta }$ can be conveniently performed via non-linear optimisation techniques. Here the problem is solved using the Gauss-Newton method (Bertsekas, 1999) with a backtracking line search, so as to maximise the objective function in (24) with respect to ${\text{Θ}}_{\beta }$. The resulting implementation is very similar to the one described in Ashburner and Friston (2005), therefore further details are omitted here. Optimisation of the affine parameters ${\text{Θ}}_a={\left\{{\mathbf{a}}_i,{\mathbf{t}}_i\right\}}_{i=1,\dots ,M}$ can also be carried out by means of a Gauss-Newton scheme and a brief description of the required computations can be found Appendix A. For the update of the weight parameters ${\text{Θ}}_w$ we adopt the same strategy outlined in Ashburner and Friston, 2005, Blaiotta et al., 2016.

The following sections instead will present in detail the algorithmic scheme used to learn the average shaped tissue templates ${\text{Θ}}_{\pi }={\left\{{\mathit{\pi }}_j\right\}}_{j=1,\dots ,N_{\pi }}$ and to estimate the set of initial velocity fields ${\text{Θ}}_u={\left\{{\mathbf{u}}_i\right\}}_{i=1,\dots ,M}$.

The pseudocode reported in alg:opt illustrates the different optimisation stages of the proposed algorithm. Essentially, the purpose of each stage within the inner loop is to optimise a subset of parameters for a particular subject, while keeping all the other parameters fixed at their current estimates. This strategy leverages the conditional dependencies among different variables in the model (such conditional dependencies are illustrated in Fig. 1), with the aim of learning more accurate parameters compared to those estimated by solving the different optimisation problems independently (Mahapatra and Sun, 2012, Ashburner and Friston, 2005, Meng and Rubin, 1993). The outer loop serves to regenerate group priors both in the form of tissue probability maps and parametric intensity priors.

Algorithm 1

Optimisation algorithm for generating average-shaped tissue probability maps using the approach presented in this paper.

Updating the tissue priors

At each outer iteration of the algorithm, the tissue priors ${\text{Θ}}_{\pi }={\left\{{\mathit{\pi }}_j\right\}}_{j=1,\dots ,N_{\pi }}$ need to be updated, given the current estimates of all the other parameters, which are kept fixed for each individual in the population.

Considering only the terms in (24) that depend on ${\text{Θ}}_{\pi }$ gives the following objective function, which has to be maximised with respect to ${\text{Θ}}_{\pi }$

$$\begin{array}{l}{\mathrm{ℒ}}_{\pi }={\mathbb{E}}_{\mathbf{Z}}\left[\text{log}p\left(\mathbf{Z}|{\text{Θ}}_{\pi },{\text{Θ}}_w,{\text{Θ}}_u,{\text{Θ}}_a\right)\right]+\text{log}p\left({\text{Θ}}_{\pi }\right)+\text{const}\\ =\sum _{i=1}^M\sum _{k=1}^K{\int }_{{\text{Ω}}_i}{\gamma }_{ik}\left(\mathit{y}\right)\text{log}\left(\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_i\left(\mathit{y}\right)\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_i\left(\mathit{y}\right)\right)}\right)\text{d}\mathit{y}\\ +\text{log}p\left({\text{Θ}}_{\pi }\right)+\text{const}\text{.}\end{array}$$ (25)

It should be noted that the parameters ${\text{Θ}}_{\pi }$ that need to be estimated are defined on the domain of the template ${\text{Ω}}_{\pi }$, rather than on the individual spaces ${\left\{{\text{Ω}}_i\right\}}_{i=1,\dots ,M}$. For this reason equation (25), which is a sum of integrals on the native domains, needs to be mapped to ${\text{Ω}}_{\pi }$, by inverting the warps ${\left\{{\mathit{\xi }}_i\right\}}_{i=1,\dots ,M}$, to give

$$\begin{array}{l}{\mathrm{ℒ}}_{\pi }^{\prime }=\sum _{i=1}^M\sum _{k=1}^K\underset{{\text{Ω}}_{\pi }}{\int }\text{det}\left(\frac{\partial {\mathit{\xi }}_i^{-1}}{\partial \mathit{y}}\right){\gamma }_{ik}\left({\mathit{\xi }}_i^{-1}\left(\mathit{y}\right)\right)\text{log}\left(\frac{w_{ik}{\pi }_k\left(\mathit{y}\right)}{{\sum }_c^K{w}_{ic}{\pi }_c\left(\mathit{y}\right)}\right)\text{d}\mathit{y}\\ +\text{log}p\left({\text{Θ}}_{\pi }\right)+\text{const},\end{array}$$ (26)

where ${\mathit{\xi }}_i^{-1}$ is the transformation mapping from image i into the template space and the determinants of the Jacobian matrices of the deformations are included to preserve volumes after the change of variables.

Finally equation (26) is discretised on a regular voxel grid, whose centres have coordinates ${\left\{{\mathbf{y}}_j\right\}}_{j=1,\dots ,N_{\pi }}$, to give

$$\begin{array}{l}{\mathrm{ℒ}}_{\pi }^{\prime }=\sum _{i=1}^M\sum _{j=1}^{N_{\pi }}\sum _{k=1}^K\text{det}\left({\mathbf{J}}_{ij}^{{\mathit{\xi }}^{-1}}\right){\gamma }_{ik}\left({\mathit{\xi }}_{ij}^{-1}\right)\text{log}\left(\frac{w_{ik}{\pi }_{jk}}{{\sum }_{c=1}^K{w}_{ic}{\pi }_{jc}}\right)\\ +\text{log}p\left({\text{Θ}}_{\pi }\right)+\text{const},\end{array}$$ (27)

where

$${\mathit{\xi }}_{ij}^{-1}={\mathit{\xi }}_i^{-1}\left(\mathit{y}\right){|}_{\mathit{y}={\mathbf{y}}_j},$$ (28)

$$\text{det}\left({\mathbf{J}}_{ij}^{{\mathit{\xi }}^{-1}}\right)=\text{det}\left(\frac{\partial {\mathit{\xi }}_i^{-1}\left(\mathit{y}\right)}{\partial \mathit{y}}\right){|}_{\mathit{y}={\mathbf{y}}_j},$$ (29)

$${\pi }_{jk}={\pi }_k\left(\mathit{y}\right)|$$ (30)

The prior term $p\left({\text{Θ}}_{\pi }\right)$ is given by the following Dirichlet distribution

$$p\left({\text{Θ}}_{\pi }\right)=\prod _{j=1}^{N_{\pi }}\text{Dir}\left({\mathit{\pi }}_j|{\mathit{\alpha }}_0\right)=C\left({\mathit{\alpha }}_0\right)\prod _{j=1}^{N_{\pi }}\prod _{k=1}^K{\pi }_{jk}^{{\alpha }_{0k}-1}\text{.}$$ (31)

Maximising equation (27) is a constrained optimisation problem, subject to

$$\sum _{k=1}^K{\pi }_{jk}=1,\forall j\in \left\{1,\dots ,N_{\pi }\right\}$$ (32)

A closed form solution could be easily found if the rescaling weights ${\left\{{\mathit{w}}_i\right\}}_{i=1,\dots ,M}$ were all equal to one. In such a case

$$\begin{array}{l}{\mathrm{ℒ}}_{\pi }^{\prime }=\sum _{i=1}^M\sum _{j=1}^{N_{\pi }}\sum _{k=1}^K\text{det}\left({\mathit{J}}_{ij}\right){\gamma }_{ik}\left({\mathit{\xi }}_{ij}^{-1}\right)\text{log}\left({\pi }_{jk}\right)\\ +\sum _{j=1}^{N_{\pi }}\sum _{k=1}^K\left({\alpha }_{0k}-1\right)\text{log}p\left({\pi }_{jk}\right)+\text{const},\end{array}$$ (33)

which could be maximised under the constraint (32), by making use of Lagrange multipliers (Falk, 1967), to give

$${\pi }_{jk}=\frac{N_{jk}+{\alpha }_{0k}-1}{{\sum }_{k=1}^K\left(N_{jk}+{\alpha }_{0k}\right)-K},$$ (34)

with $N_{jk}={\sum }_{i=1}^M\text{det}\left({\mathit{J}}_{ij}\right){\gamma }_{ik}\left({\mathit{\xi }}_{ij}^{-1}\right)$.

This solution would provide maximum a posteriori point estimates of ${\text{Θ}}_{\pi }={\left\{{\mathit{\pi }}_j\right\}}_{j=1,\dots ,N_{\pi }}$. However for this problem, it would also be possible to derive a full variational posterior distribution, which, like its prior, would take a Dirichlet form, with parameters ${\mathit{\alpha }}_j={\mathit{\alpha }}_0+{\mathit{N}}_j$ .

When rescaling of the tissue priors is allowed the optimisation problem becomes more complex. The strategy adopted here consists in finding an approximate solution to the unconstrained optimisation problem by setting the derivatives of the objective function in (26) to zero

$$\frac{{\alpha }_{0k}-1}{{\pi }_{jk}}+\sum _{i=1}^M\text{det}\left({\mathbf{J}}_{ij}^{{\mathit{\xi }}^{-1}}\right){\gamma }_{ik}\left({\mathit{\xi }}_{ij}^{-1}\right)\left(\frac{1}{{\pi }_{jk}}-\frac{w_{ik}}{{\sum }_{c=1}^K{w}_{ic}{\pi }_{jc}}\right)=0\text{.}$$ (35)

Solving with respect to ${\pi }_{jk}$, under the simplifying assumption that the term ${\sum }_{c=1}^K{w}_{ic}{\pi }_{jc}$ can be treated as a constant, gives

$${\overline{\pi }}_{jk}=\frac{N_{jk}+{\alpha }_{0k}-1}{{\sum }_{i=1}^M\frac{\text{det}\left({\mathbf{J}}_{ij}^{{\mathit{\xi }}^{-1}}\right){\gamma }_{ik}\left({\varphi }_{ij}^{-1}\right)w_{ik}}{{\sum }_{c=1}^K{w}_{ic}{\pi }_{jc}}}\text{.}$$ (36)

Such a solution is then projected onto the constraining hyperplane, by preserving tissue proportions at each voxel

$${\pi }_{jk}=\frac{{\overline{\pi }}_{jk}}{{\sum }_{c=1}^K{\overline{\pi }}_{jc}}\text{.}$$ (37)

Experimental testing of this strategy indicated that it gave a constant improvement of the objective function at a relatively cheap computational cost. Alternatively, iterative constrained non-linear optimisation techniques (Powell, 1978) could have been exploited to solve the template update problem.

Computing the deformation fields

Groupwise image alignment is achieved by optimisation of the variational objective function defined in (24), with respect to the parameters used to compute the deformations. This is equivalent to adopting the following image matching or similarity term

$$\begin{array}{l}\mathcal{D}={\mathbb{E}}_{\mathbf{Z}}\left[\text{log}p\left(\mathbf{Z}|{\text{Θ}}_{\pi },{\text{Θ}}_w,{\text{Θ}}_u,{\text{Θ}}_a\right)\right]\\ =\sum _{i=1}^M{\int }_{\mathit{y}\in {\text{Ω}}_i}\sum _{k=1}^K{\gamma }_{ik}\left(\mathit{y}\right)\text{log}\left(\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_i\left(\mathit{y}\right)\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_i\left(\mathit{y}\right)\right)}\right)\text{d}\mathit{y}\text{.}\end{array}$$ (38)

Additionally, working on discretised image grids, with associated voxel centres ${\left\{{\mathbf{y}}_{ij}\right\}}_{j=1,\dots ,N_i}$, requires reformulating $\mathcal{D}$ as

$$\mathcal{D}=\sum _{i=1}^M\sum _{j=1}^{N_i}\sum _{k=1}^K{\gamma }_{ijk}\text{log}\frac{w_{ik}\pi {\text{'}}_{jk}}{{\sum }_{c=1}^K{w}_{ic}\pi {\text{'}}_{jc}},$$ (39)

with

$${{\pi }^{\prime }}_{jk}={\pi }_k\left({\mathit{\xi }}_i\left(\mathit{y}\right)\right){|}_{\mathit{y}={\mathbf{y}}_{ij}}\text{.}$$ (40)

The penalty term for this groupwise image registration problem is given by

$$\begin{array}{l}\mathrm{ℛ}={\mathrm{ℛ}}_{dif}+{\mathrm{ℛ}}_{af}=\text{log}p\left({\text{Θ}}_u\right)+\text{log}p\left({\text{Θ}}_a\right)\\ =-\frac{1}{2}\sum _{i=1}^M\left(\left|\right|{\mathbf{L}}_{\mathbf{u}}{\mathbf{u}}_i|{|}_{L^2}^2+{\mathbf{a}}_i^T{\mathbf{\Sigma }}_{\mathit{a}}^{-1}{\mathbf{a}}_i\right)+\text{const},\end{array}$$ (41)

with ${\mathbf{u}}_i$ being a $3\times N_i$ dimensional vector of parameters used for representing the initial velocity field of image i and ${\mathbf{a}}_i$ encoding affine deformation parameters used to compute the transformation in (19).

For each image i in the data set, updating the corresponding initial velocity field, given the current estimates of the templates and all the other model parameters, involves optimising the following objective function

$$\begin{array}{l}{\mathcal{E}}_{dif}^{\left(i\right)}={\mathcal{D}}^{\left(i\right)}+{\mathrm{ℛ}}_{dif}^{\left(i\right)}\\ =\sum _{j=1}^{N_i}\sum _{k=1}^K{\gamma }_{ijk}\text{log}\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_{ij}\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_{ij}\right)}-\frac{1}{2}\left|\right|{\mathbf{L}}_{\mathbf{u}}{\mathbf{u}}_i|{|}_{L^2}^2,\end{array}$$ (42)

with respect to ${\mathbf{u}}_i$, under the following deformation model

$${\mathit{\xi }}_{ij}={\mathit{\xi }}_i\left({\mathbf{y}}_{ij}\right)={\mathbf{T}}_i{\varphi }_i\left({\mathbf{y}}_{ij}\right)+{\mathbf{t}}_i,$$ (43)

where ${\varphi }_i$ is a diffeomorphism computed via geodesic shooting (Ashburner and Friston, 2011) from the corresponding initial velocity field ${\mathbf{u}}_i$.

Here image registration is solved via Gauss-Newton optimisation, with a backtracking line search to ensure convergence. This requires computing both the first and second derivatives of the objective function (Hernandez and Olmos, 2008). Such derivatives can be found in Appendix B. This leads to a very high dimensional inverse problem, which unfortunately cannot be solved via numerical matrix inversion, since this would be prohibitively expensive from a computational point of view. The approach adopted in this work consists in treating this optimisation as a partial differential equation problem, which can efficiently be solved using multigrid methods (Modersitzki, 2004). In particular, we adopt the same full multigrid implementation as in Ashburner (2007).

Validation and discussion

In this section we present results obtained by applying the presented modelling framework to real MR scans acquired with different imaging systems and protocols, as well as to synthetic MR volumes. Both qualitative and quantitative measures will be provided to assess the behaviour of the proposed approach.

All validation analyses have been performed on the full model presented in this paper, since the validity of the individual components (i.e. bias correction, diffeomorphic registration, variational Gaussian mixture modelling) was already assessed in our previous work (Ashburner and Friston, 2005, Ashburner and Friston, 2011, Blaiotta et al., 2016), primarily at a single subject level. Additionally, numerous studies have demonstrated that solving image segmentation and registration tasks in a coupled and iterative manner, as opposed to sequential approaches, ensures more accurate results while reducing the chances of getting prematurely trapped in a local optimum (Ashburner and Friston, 2005, Pohl et al., 2006, Yezzi et al., 2001, Bhatia et al., 2007). This results from the conditional dependencies between the parameters controlling the deformation fields and the intensity distribution parameters, which are used to infer tissue labels. In other words, closely matching the template to each scan is a necessary condition in order to compute accurate segmentations while, in turn, accurate segmentations of the data help to improve alignment with the templates.

Template construction

We begin validating our approach by illustrating a set of templates obtained by fitting the proposed model to a large neuroimaging data set. Indeed, in remainder of this paper we will be using neuroimaging data for both training and testing, even if, given the generality of the presented approach, applications to many more types of data could be explored.

Data

The input data for training the model was obtained from three different databases, two of which are freely accessible for download, thus ensuring that the results presented here could readily be compared to those produced by competing algorithms for medical image registration or segmentation.

OASIS data set. The first data set consists of thirty five T1-weighted MR scans from the OASIS (Open Access Series of Imaging Studies) database (Marcus et al., 2007). The data is freely available from the web site http://www.oasis-brains.org, where details on the population demographics and acquisition protocols are also reported. Additionally, the selected thirty five subjects are the same ones that were used within the 2012 MICCAI Multi-Atlas Labeling Challenge (Landman and Warfield, 2012).

Balgrist data set. The second data set consists of brain and cervical cord scans of twenty healthy adults, acquired at University Hospital Balgrist with a 3T scanner (Siemens Magnetom Verio). Magnetisation-prepared rapid acquisition gradient echo (MPRAGE) sequences, at 1 $mm$ isotropic resolution, were used to obtain T1-weighted data, while PD-weighted images of the same subjects were acquired with a multi-echo 3D fast low-angle shot (FLASH) sequence, within a whole-brain multi-parameter mapping protocol (Weiskopf et al., 2013, Helms et al., 2008).

IXI data set. The third and last data set comprises twenty five T1-, T2-and PD-weighted scans of healthy adults from the freely available IXI brain database, which were acquired at Guy's Hospital, in London, on a 1.5T system (Philips Medical Systems Gyroscan Intera). Additional information regarding the demographics of the population, as well as the acquisition protocols, can be found at http://brain-development.org/ixi-dataset.

The complete data set therefore consists of eighty multispectral scans of healthy adults, obtained with fairly diverse acquisition protocols and using scanning systems produced by different vendors.

Unfortunately, not all the three modalities of interest (T1-, T2-and PD-weighted) are available for all of the subjects. To circumvent the difficulties arising from the presence of missing imaging modalities, without neglecting any of the available data (indeed deletion of entries with missing data is still, in spite of its crudity, a common statistical practice), the Gaussian mixture modelling approach discussed in Blaiotta et al. (2016) was generalised by introducing an additional variational posterior distribution over the missing data points.

In practice, the resulting variational EM scheme iterates over first estimating an approximated posterior distribution on the unknown image intensities, secondly updating the sufficient statistics of the complete (observed and missing) data and finally computing variational posteriors on the Gaussian mixture parameters. Additional computational details relative to this strategy are provided in Appendix C.

In synthesis, it was possible to fit the generative groupwise model described in this paper to the entire data set, in spite of having different imaging modalities available from the different acquisition sites. This is indeed a very common scenario in real life medical imaging problems, therefore it should be actively addressed by processing or modelling solutions that claim to be applicable to large population data (van Tulder and de Bruijne, 2015).

Manual brain labels are freely available for the selected subset of the OASIS data set. Such labels have been generated and made public by Neuromorphometrics, Inc. (http://Neuromorphometrics.com) under academic subscription and they provide a fine parcellation of cortical and non cortical structures, for a total of 139 labels across the brain.

Part of this label data was used for training of the model while the remainder was left out for testing and validation. In particular, brain labels of twenty out of the thirty five OASIS subjects were used to create gray and white matter ground truth segmentations, which were provided as training input for semisupervised model fitting. In spite of having defined only one gray matter training label, two distinct gray matter classes were introduced in the mixture model (top two rows in Fig. 2), to best capture the corresponding distribution of image intensities, which is poorly represented by a single Gaussian component, as opposed to the distribution of white matter intensities. In this case, posterior membership probabilities were computed making use of equation (17).

Fig. 2.

Tissue probability maps obtained by applying the presented groupwise generative model to a multispectral data set comprising head and neck scans of eighty healthy adults, from three different databases.

Spinal cord manual labels were created for forty subjects (twenty from Balgrist data set and twenty from the IXI data set), by manually delineating the contour of the cord in each transverse slice of the data beginning from the lower extremity of the medulla oblongata. Such labels were randomly split in half for training and half for subsequent test analyses. Due to the limited resolution of the data it was not possible to manually delineate gray and white matter within the spinal cord. For this reason, each voxel classified as spinal cord in the training data was allowed to be assigned either to the gray or to the white matter tissue classes, based on the fit of its intensity value to the underlying Gaussian mixture model, as outlined in equation (17).

Tissue templates and intensity priors

The tissue probability maps obtained by applying the modelling framework presented in this paper to the data set described above are depicted in Fig. 2. The total number of tissue classes used for this experiment is equal to twelve but three classes, representing air in the background, are not shown.

The total number of classes was selected based on empirical evidence to obtain a reasonable trade-off between goodness of fit and computational cost. In principle the proposed algorithm would be able to automatically determine the optimal number of classes, as demonstrated in Blaiotta et al. (2016), however this would require setting an initial number of components higher than the unknown optimal one, which, for the size of the data set considered here, would have been computationally very expensive.

Apart from the number of Gaussians, the only hyperparameters that are not estimated from the data but need to be fixed a priori are those controlling the bias and registration regularisation. Indeed, one of the strengths of the proposed approach is that, in spite of its complexity most parameters are automatically inferred from the data, thus requiring minimal parameter tuning. For the experiments described in this section we used the settings described below.

The prior distribution on the bias field parameters was assumed to be Gaussian with zero mean and a precision matrix corresponding to a penalty on the Laplacian of the resulting non-uniformity field. Similarly, the affine model was regularised using a Gaussian prior with zero mean and a diagonal covariance matrix, where the magnitude of parameters modelling translations and rotations was less heavily regularised (the prior variance was set equal to 1), compared to that of parameters controlling scaling and shearing, whose prior variance was instead set equal to 0.001. Finally the regularisation settings for diffeomorphic registration were borrowed from the default settings of the Geodesic Shooting toolbox (Ashburner and Friston, 2011) provided with SPM12, which correspond to penalising a linear combination of absolute displacement, membrane energy, bending energy and linear elasticity (more details can be found in the documentation of the SPM12 software).

Fig. 2 shows how one of the two gray matter classes (first row) best fits the subcortical nuclei and also includes voxels affected by partial volume effects at the interface between gray and white matter, while the second one (second row) is more representative of cortical structures, with partial volume effects generated by the mixed presence of gray matter and CSF. The third row in Fig. 2 shows the white matter class, which also includes most of the brainstem and the spinal cord.

The remaining tissue classes were estimated in a purely unsupervised way. Therefore a non ambiguous anatomical interpretation is not straightforward.

Tissue class four (fourth row) mainly contains CSF, even if other tissues are also present, especially in the neck area. This should be attributed to the lack of CSF training labels as well as to a poor multivariate coverage of the cervical region in the available data. In fact, data from the OASIS set is truncated around the first cervical vertebra. The T1-weighted scans of the IXI data set cover up to the C2/C3 vertebral level, but the corresponding T2-and PD-weighted scans do not extend beyond the brainstem. Indeed, only the data from the second database (Balgrist hospital) provides more than one modality covering up to around the fourth cervical vertebra. In this case though, additional difficulties arose from poor inter-modality alignment of the data, a problem that turned out to be particularly severe in the cervical region and that, given its non-linearity, cannot be fully compensated for by affine inter-modality coregistration. This result confirms the importance and usefulness of non-linear image coregistration tools, particularly when modelling highly deformable anatomical structures (Stroman et al., 2008, Fonov et al., 2014). Such a problem was not addressed in this paper, which limits the applicability of the proposed approach when scans of the same subject are not in good alignment. However it would in principle be possible to introduce inter-modality deformation fields within the same generative model adopted here and we aim to address this issue as part of our future work.

Bone tissue is also not easily identifiable from the data available for this experiment, but it could have potentially been much better extracted by incorporating some CT scans into the training data.

Fat and soft tissues are mainly represented in the last two classes (bottom two rows in Fig. 2).

While quantitative evaluation analyses of the proposed modelling approach will be reported in the following sections, the results presented in Fig. 2 provide a qualitative insight into the performance of the algorithm discussed in this paper. In particular, the sharp appearance of the tissue probability maps suggests that the proposed model can capture a significant amount of anatomical variability thus ensuring fine alignment of complex anatomical structures, whose shape varies significantly across individuals. In turn, this represents a very valuable property for the purpose of performing statistical group analyses on neuroimaging data, since most of these studies rely on the assumption that accurate anatomical and functional correspondences can be estimated through some form shape mapping procedure, which is most commonly implemented using non-linear image registration techniques.

Fig. 3 illustrates zoomed views of the gray and white matter tissue probability maps at the brainstem and spinal cord levels.

Fig. 3.

Zoomed views of the gray (top two rows) and white matter (bottom row) tissue probability maps at the brainstem and spinal cord levels. Each transverse slice covers an area of $45\times 53{\text{ mm}}^2$, while the axial distance between adjacent slices along each row equals $7.5\text{ mm}$. The order of tissue classes from top to bottom is the same as in the top three rows of Fig. 2, however more densely spaced transverse slices are depicted here.

The empirical Bayes learning procedure, introduced in Blaiotta et al. (2016) to estimate suitable prior distributions for the parameters of the Gaussian mixture model, was applied here to the same data used to construct the templates. Some of the results are summarised in Fig. 4, where the estimated empirical prior distributions on the mean intensity of gray and white matter are depicted, with overlaid contour plots showing some of the individual posteriors (randomly selected across the entire population).

Fig. 4.

Prior distribution over the mean intensity of gray and white matter, in T1-and PD-weighted data. Contour plots illustrate a number of randomly selected individual posteriors.

Such results indicate that the proposed empirical Bayes learning scheme can serve to capture, not only the variability of mean tissue intensity across subjects for each of the modalities of interest, but also the amount of covariance between such modalities. Information of this sort can potentially be used in a number of different frameworks, for solving problems such as tissue segmentation, pathology detection or image synthesis.

In Fig. 5 we report a plot of the lower bound during model fitting as a qualitative demonstration of convergence.

Fig. 5.

Lower bound values as a function of iteration number during model training.

Run time for the data set presented here was around 50 h on an 8 Core PC at 3 GHz with 32 GB RAM.

Validity of groupwise registration

The performance of groupwise registration achieved by the presented algorithm was assessed by computing pairwise overlap measures for all possible couples of spatially normalised test label maps (i.e. label maps that were not used for training the model), which were obtained by applying the inverse of the estimated deformation fields to the maps in their native space. The Dice⁴ score coefficient was chosen as a metric of similarity.

Results are summarised in Fig. 6, where the accuracy of the algorithm presented here is compared to that achieved by the groupwise image registration method described in Avants et al. (2010), whose implementation is publicly available, as part of the Advanced normalisation Tools (ANTs) package, through the web site http://stnava.github.io/ANTs/. Indeed, the symmetric diffeomorphic registration framework implemented in ANTs has established itself as the state-of-the-art of medical image nonlinear spatial normalisation (Klein et al., 2009).

Fig. 6.

Registration accuracy achieved by the presented method, compared to the performance of the diffeomorphic groupwise template construction method distributed with ANTs. Boxplots, in green for the proposed method and blue for ANTs, illustrate the distributions of overlap measures (Dice score coefficients) obtained by considering all possible pairs of spatially normalised test images, whose labels were not used for model fitting. Stars indicate statistically significant differences between the two methods, assessed by means of paired t-tests without correcting for multiple comparisons. Star colours encode the method that achieved higher performance (black for the proposed algorithm and red for ANTs).

A number of options can be customised within the template construction framework distributed with ANTs. The experiments, whose results are reported here, were performed with the settings recommended in the package documentation for brain MR data, which are also reported in Table 2. The same strategy outlined above was applied in order to compute Dice score coefficients for all structural labels and all possible pairs of test images.

Table 2.

Options selected to perform groupwise registration with ANTs, using the antsMultivariateTemplateConstruction script provided with the ANTS package.

Option	Value
Similarity Metric	Cross-correlation (CC)
Transformation model	Greedy SyN (GR)
Initial rigid body	yes
N4 Bias Correction	yes
Number of resolution levels	4
Number of iterations	$100\times 70\times 50\times 10$
Gradient step	0.2
Number of template updates	4

Results of this validation analyses indicate that the method presented here, in spite of not being as accurate as ANTs for aligning some subcortical brain structures (e.g. thalamus, putamen, pallidum and brainstem), provided significantly better overlap when registering cortical regions, as assessed by means of paired t-tests with a significance threshold of 0.05 and without correcting for multiple comparisons.

Accuracy of tissue classification

The accuracy of tissue classification achieved by the method presented in this paper was first evaluated on test data that was used to create the templates but without providing manual labels for training the model. The aim in this case is to determine to which extent the proposed method can capture significant features of the training data, by learning from few annotated examples.

Dice scores were computed to compare the automated segmentations produced via semisupervised groupwise model fitting, with the ground truth, obtained by merging all the gray and white matter brain structures (labels) into two tissue classes respectively. The probabilistic gray and white matter segmentations produced by the proposed algorithm were thresholded at 0.5, in order to obtain binary label maps, directly comparable to the ground truth. Results are summarised in seg, which shows the distributions of Dice scores obtained for gray and white matter.

Such results were then compared to those produced by the brain segmentation algorithm implemented in SPM12, using the tissue probability maps illustrated in Fig. 2. Indeed the proposed approach extends and generalises the model underlying SPM12 segmentation method. This is achieved by replacing a small deformation approach to image registration with a diffeomorphic representation of shapes, as well as by introducing empirical intensity priors (within a variational Bayes framework) and by allowing learning of the tissue probability maps directly from the data.

Therefore, results of these analyses, which are also shown in Fig. 7, provide a direct insight into the performance gain achieved by incrementing the model available in SPM12 as indicated above. In particular, our experiments indicate that the proposed method enables higher tissue classification accuracy, for both gray and white matter (assessed by means of paired t-tests with a significance threshold of 0.05).

Fig. 7.

Brain segmentation accuracy of the presented method in comparison to SPM12 image segmentation algorithm. Boxplots indicate the distributions of Dice score coefficients, with overlaid scatter plots of the estimated scores. Red stars denote outliers.

The experiments described in this section however did not test the generalisation capability of the proposed method, which would have required a k-fold cross-validation design in order to exploit as much as possible of the available data during training. Unfortunately this was highly unpractical in this case due to the expensive computational cost of groupwise model fitting. The generalisation performance of the proposed model is instead assessed in the next section, using a hold-out approach and making use of both real and synthetic unseen brain MR data, for which tissue labels are known.

Modelling unseen data

Further validation experiments were performed to quantify the accuracy of the framework described in this paper to model unseen data, that is to say data that was not included in the atlas generation process.

In particular, we evaluated registration accuracy using data from the Internet Brain Segmentation Repository (IBSR), which is provided by the Centre for Morphometric Analysis at Massachusetts General Hospital (http://www.cma.mgh.harvard.edu/ibsr/). Experiments to assess bias correction and segmentation accuracy were instead performed on synthetic T1-weighted brain MR scans from the Brainweb database (http://brainweb.bic.mni.mcgill.ca/), which were simulated using a healthy anatomical model under different noise and bias conditions.

Accuracy of registration

To assess the performance of the presented method for spatially normalising unseen test data we made use of the ISBR data set, which consists of 18 T1-weighted brain images with manual labels of 84 anatomical structures. Such a data set was also used by Klein et al. (2009) for their evaluation of 14 brain image registration algorithms. Therefore the results presented in the remainder of this section allow comparing the proposed approach to a number of image registration methods available to the neuroimaging community.

For this experiment we registered each of the 18 unseen scans to the tissue probability maps presented in Fig. 2 by adopting the algorithmic framework presented in this paper. We then computed target overlap measures for all pairs of spatially normalised images. This was achieved by composing the inverse of each estimated deformation field with every other direct transformation and applying the resulting warps to the labels in their native space, so as to map between each couple of subjects. Our method, after averaging across different brain regions, achieved a median overlap score of 0.54 with the 25th and 75th percentiles equal to 0.51 and 0.57 respectively. Instead, the best performing method in the experiments performed by Klein et al. (2009) obtained a median score of 0.55 on the same data set.

Such results indicate that our approach, in spite of being intrinsically best suited for groupwise analyses, where the reference image is constructed iteratively from the same scans that are being spatially normalised (Avants et al., 2010), can provide accurate results also when mapping test data to a fixed reference that encodes the average shape of a different population. Indeed, this experiment was intended as a means to assess the generalisation capability of the templates constructed with the proposed approach, while ensuring a fairer comparison to the results of Klein et al. (2009), which were produced via pairwise non-linear registration rather than by estimating group averages.

Accuracy of bias correction

A healthy adult brain MR model was processed by means of the algorithm discussed here, using the brain and spinal cord templates previously constructed as tissue priors. Different levels of noise and bias field were added to the uncorrupted synthetic data, to test the behaviour of the proposed modelling scheme in different noise (1%, 3%, 7%) and bias conditions (20% and 40%).

The noise in these simulated images has Rayleigh statistics in the background and Rician statistics in the signal regions and its level is computed as a percent standard deviation ratio, relative to the MR signal, for a reference tissue (Cocosco et al., 1997).

Regarding the bias field instead, 20% bias is modelled as a smooth field in the range [0.9, 1.1] while 40% bias is obtained by rescaling of the 20% field, so as to range between 0.8 and 1.2.

Table 3 reports the Pearson product-moment correlation coefficients between the ground truth and the estimated bias fields, for the different bias ranges and noise levels. Results indicate that the similarity between the estimated and true bias decreases for more intense non-uniformity fields and higher noise levels.

Table 3.

Pearson's correlation coefficients between the ground truth bias fields and those estimated by the presented algorithm, for simulated T1-weighted data.

Bias	Noise
	1%	3%	7%
20%	0.86	0.86	0.70
40%	0.72	0.72	0.51

Indeed this is not surprising, as the penalty term, which enforces smoothness of the bias field, has a greater impact in determining the shape of the estimated bias when the non-uniformity fields have a larger dynamic range. Nevertheless, results reported in the following section will show how this increased mismatch between the estimated and true bias, for higher non-uniformities, does not seem to affect the accuracy of tissue segmentation. On the other hand, the accuracy of bias correction is directly related to the amount of noise corrupting the data, mainly due to how this affects the precision associated with estimation of the Gaussian mixture parameters. For a comparison of these results with the performance of SPM12 bias correction on simulated T1-weighted scans from the Brainweb database see Blaiotta et al. (2016).

Accuracy of tissue classification

For the same synthetic data the accuracy of tissue classification was also evaluated, by comparing the similarity between the estimated gray and white matter segmentations and the underlying anatomical model.

Results are reported in Fig. 8, which shows the Dice score coefficients obtained under different bias and noise conditions.

Fig. 8.

Dice scores between the estimated and ground truth segmentations for brain white matter and brain gray matter, under different noise and bias conditions, for synthetic T1-weighted data.

The Brainweb database has been extensively used in the neuroimaging community to validate MR image processing algorithms. Therefore the results reported here should be directly comparable to the performance of many brain segmentation techniques present in the literature.

Conclusions

This paper presented a comprehensive generative modelling framework, suitable to capture anatomical variability from large cross-sectional MR data sets. From a theoretical perspective, such a framework relies on variational probability density estimation techniques to model the observed data (i.e. MR image intensities). Additionally, a hierarchical modelling perspective is proposed, where observations from a population of subjects are used to construct empirical intensity priors, which can then serve to inform models of new data. Shape modelling is performed via groupwise diffeomorphic registration, thus ensuring bijective (i.e. one-to-one) differentiable mappings between anatomical configurations (Miller, 2004). Such an approach enables a rigorous mathematical encoding of anatomical shapes via deformable template matching (Christensen et al., 1996), therefore providing a quantitative framework for the analysis of shape variation and covariation.

Data for validating the method was collected from many different brain databases, most of which are publicly accessible to the research community. Results of our experiments, performed both on training and unseen test data, indicate that the proposed approach can accurately align complex anatomical shapes, such as the brain and spinal cord, and segment data into tissue types, while being robust to inter-scanner signal variations. Therefore, the proposed algorithm defines a convenient framework to extract volumetric and morphometric information from large structural neuroimaging data sets, in a fully automated manner. At the same time it provides outputs that can be readily interpreted, for instance via statistical hypothesis testing, with the ultimate goal of comparing different populations, treatment effects etc. (Ashburner and Friston, 2000).

Additionally, our results suggest that the proposed approach could be useful to construct templates capable of capturing the peculiar anatomical features of populations poorly represented by standard anatomical atlases (such as young or elderly populations, diseased populations, or individuals belonging to different ethnic groups (Tang et al., 2010, Fillmore et al., 2015)). This would not only lead to more accurate segmentation results, but as a direct consequence, also increase the reliability of subsequent data analyses, which build models of the segmented data to infer or predict clinically meaningful information.

It should be noted that, in spite of having been tested on neuroimaging data sets, our method was intentionally formulated as a general approach, which makes it potentially suitable to solve a wide range of imaging problems, for instance in the context of animal imaging studies or for the analysis of different human organs using multispectral MR data sets. We aim to explore some of this applications as part of our future work.

Appendix A. Derivatives of the lower bound with respect to the affine parameters

The affine parameters, for each subject i, can be estimated (i.e. optimised) in a Gauss-Newton fashion, so as to maximise of the following objective function

$$\begin{array}{l}{\mathcal{E}}_{af}^{\left(i\right)}={\mathcal{D}}^{\left(i\right)}+{\mathrm{ℛ}}_{af}^{\left(i\right)}\\ =\stackrel{N_i}{\sum _{j=1}}\stackrel{K}{\sum _{k=1}}{\gamma }_{ijk}\text{log}\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_{ij}\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_{ij}\right)}-\frac{1}{2}{\mathbf{a}}_i^T{\mathbf{\Sigma }}_{\mathit{a}}^{-1}{\mathbf{a}}_i,\end{array}$$ (A.1)

with respect to ${\mathbf{a}}_i$.

The gradients and Hessians, which are useful to solve this problem are reported below. In particular, for the matching term, the following derivatives need to be computed

$$\frac{\partial {\mathcal{D}}^{\left(i\right)}}{\partial {\mathbf{a}}_i}=\stackrel{N_i}{\sum _{j=1}}\stackrel{K}{\sum _{k=1}}\left({\gamma }_{ijk}-\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_{ij}\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_{ij}\right)}\right){\mathit{g}}_{jk}^{\pi },$$ (A.2)

where ${\mathit{g}}_{jk}^{\pi }$ is defined as

$${\mathit{g}}_{jk}^{\pi }={\mathbf{B}}_i^T\left([{\varphi }_{ij},1]\otimes \nabla \left[\text{log}\left({\pi }_k\left({\mathit{\xi }}_{ij}\right)\right)\right]\right),$$ (A.3)

with

$${\mathbf{B}}_i^T=\frac{\partial {\mathbf{S}}_i}{\partial {\mathbf{a}}_i},$$ (A.4)

and

$${\mathbf{S}}_i=\left[\begin{array}{cc}{\mathbf{T}}_i& {\mathbf{t}}_i\\ 0& 1\end{array}\right]\text{.}$$ (A.5)

$$\begin{array}{l}\frac{{\partial }^2{\mathcal{D}}^{\left(i\right)}}{\partial {\mathbf{a}}_i^2}=\stackrel{N_i}{\sum _{j=1}}\left(\stackrel{K}{\sum _{k=1}}\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_{ij}\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_{ij}\right)}{\mathit{g}}_{jk}^{\pi }\right)\\ \times {\left(\stackrel{K}{\sum _{k=1}}\frac{w_{ik}\left({\pi }_k\left({\mathit{\xi }}_{ij}\right)\right)}{{\sum }_{c=1}^K{w}_{ic}\left({\pi }_k\left({\mathit{\xi }}_{ij}\right)\right)}{\mathit{g}}_{jk}^{\pi }\right)}^T\\ -\stackrel{N_i}{\sum _{j=1}}\stackrel{K}{\sum _{k=1}}\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_{ij}\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_{ij}\right)}{\mathit{g}}_{jk}^{\pi }{\left({\mathit{g}}_{jk}^{\pi }\right)}^T\text{.}\end{array}$$ (A.6)

Gradients and Hessians of the penalty term are instead given by

$$\frac{\partial {\mathrm{ℛ}}_{af}^{\left(i\right)}}{\partial {\mathbf{a}}_i}=-{\mathbf{\Sigma }}_{\mathit{a}}^{-1}{\mathbf{a}}_i,$$ (A.7)

$$\frac{{\partial }^2{\mathrm{ℛ}}_{af}^{\left(i\right)}}{\partial {\mathbf{a}}_i^2}=-{\mathbf{\Sigma }}_{\mathit{a}}^{-1}\text{.}$$ (A.8)

Appendix B. Derivatives of the lower bound with respect to the initial velocities

Optimisation of the initial velocities, for each image i, requires maximising the following objective function

$$\begin{array}{l}{\mathcal{E}}_{dif}^{\left(i\right)}={\mathcal{D}}^{\left(i\right)}+{\mathrm{ℛ}}_{dif}^{\left(i\right)}\\ =\stackrel{N_i}{\sum _{j=1}}\stackrel{K}{\sum _{k=1}}{\gamma }_{ijk}\text{log}\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_{ij}\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_{ij}\right)}-\frac{1}{2}\stackrel{M}{\sum _{i=1}}\left|\right|{\mathbf{L}}_{\mathbf{u}}{\mathbf{u}}_i|{|}_{L^2}^2,\end{array}$$ (B.1)

with respect to ${\mathbf{u}}_i$.

Here, we report the first and second derivatives of this objective function, which are useful to solve the registration problem using gradient-based techniques, such as the Gauss-Newton algorithm.

The gradient of the matching term ${\mathcal{D}}^{\left(i\right)}$ with respect to ${\mathbf{u}}_i$ is given by

$$\begin{array}{l}\frac{\partial {\mathcal{D}}^{\left(i\right)}}{\partial {\mathbf{u}}_i}=\stackrel{K}{\sum _{k=1}}{\gamma }_{ijk}\frac{\partial }{\partial {\mathbf{u}}_i}\left(\text{log}\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_i\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_i\right)}\right)\\ =\stackrel{K}{\sum _{k=1}}{\gamma }_{ijk}\left({\mathit{g}}_k^{\pi }-\stackrel{K}{\sum _{c=1}}\frac{w_{ic}{\pi }_c\left({\mathit{\xi }}_i\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_i\right)}{\mathit{g}}_c^{\pi }\right),\end{array}$$ (B.2)

which, making use of ${\sum }_{k=1}^K{\gamma }_{ijk}=1$, can be rewritten as

$$\frac{\partial {\mathcal{D}}^{\left(i\right)}}{\partial {\mathbf{u}}_i}=\stackrel{K}{\sum _{k=1}}\left({\mathit{\gamma }}_{ik}-\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_i\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_i\right)}\right){\mathit{g}}_k^{\pi },$$ (B.3)

where ${\mathit{g}}_k^{\pi }$ is computed, at each voxel j, by

$${\mathit{g}}_{jk}^{\pi }={\left({\mathbf{T}}_i,{\mathbf{J}}_{ij}^{\mathit{\xi }}\right)}^T\nabla \left[\text{log}\left({\pi }_k\left({\mathit{\xi }}_{ij}\right)\right)\right],$$ (B.4)

and ${\mathbf{J}}_i^{\mathit{\xi }}$ indicates the Jacobian matrix of ${\mathit{\xi }}_{ij}$.

An approximated positive semidefinite Hessian of $\mathcal{D}$ can instead be computed by discarding the second derivatives of the logarithm of tissue priors

$$\frac{{\partial }^2}{\partial {\mathit{y}}^2}\text{log}\left(\frac{w_{ik}\left({\pi }_k\left({\mathit{\xi }}_i\left(\mathit{y}\right)\right)\right)}{{\sum }_{c=1}^K{w}_{ic}\left({\pi }_c\left({\mathit{\xi }}_i\left(\mathit{y}\right)\right)\right)}\right)=0,\forall \mathit{y}\in {\text{Ω}}_i,$$ (B.5)

to give

$$\begin{array}{l}\frac{{\partial }^2{\mathcal{D}}^{\left(i\right)}}{\partial {\mathbf{u}}_i^2}=\left(\stackrel{K}{\sum _{k=1}}\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_i\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_i\right)}{\mathit{g}}_k^{\pi }\right)\\ \times {\left(\stackrel{K}{\sum _{k=1}}\frac{w_{ik}\left({\pi }_k\left({\mathit{\xi }}_i\right)\right)}{{\sum }_{c=1}^K{w}_{ic}\left({\pi }_k\left({\mathit{\xi }}_i\right)\right)}{\mathit{g}}_k^{\pi }\right)}^T\\ -\stackrel{K}{\sum _{k=1}}\frac{w_{ik}{\pi }_k\left({\mathit{\xi }}_i\right)}{{\sum }_{c=1}^K{w}_{ic}{\pi }_c\left({\mathit{\xi }}_i\right)}{\mathit{g}}_k^{\pi }{\left({\mathit{g}}_k^{\pi }\right)}^T\text{.}\end{array}$$ (B.6)

Finally, the first and second derivatives of the penalty term $\mathrm{ℛ}$, which are also required to optimise (B.1), can be computed by

$$\frac{\partial {\mathrm{ℛ}}_{dif}^{\left(i\right)}}{\partial {\mathbf{u}}_i}=-{\mathbf{L}}_{\mathbf{u}}^{†}{\mathbf{L}}_{\mathbf{u}}{\mathbf{u}}_i,$$ (B.7)

$$\frac{{\partial }^2{\mathrm{ℛ}}_{dif}^{\left(i\right)}}{\partial {\mathbf{u}}_i^2}=-{\mathbf{L}}_{\mathbf{u}}^{†}{\mathbf{L}}_{\mathbf{u}}\text{.}$$ (B.8)

Appendix C. Variational Gaussian mixtures: inference of missing data

The variational Bayes EM algorithm for fitting Gaussian mixture models, described in Blaiotta et al. (2016), can be generalised to handle the case where some components of the D-dimensional observation ${\mathbf{x}}_j$ are missing.

Having denoted

$${\mathbf{x}}_j=\left[\begin{array}{c}{\mathbf{o}}_j\\ {\mathbf{h}}_j\end{array}\right],$$ (C.1)

with ${\mathbf{o}}_j$ being the observed data and ${\mathbf{h}}_j$ the missing (hidden) data, the Gaussian likelihood $p\left({\mathbf{x}}_j\right|z_{jk}=1,{\mathit{\mu }}_k,{\mathbf{\Sigma }}_k)$ can be expressed as

$$p\left({\mathbf{x}}_j\right|z_{jk}=1,{\mathit{\mu }}_k,{\mathbf{\Lambda }}_k)=\mathcal{N}\left(\left[\begin{array}{c}{\mathbf{o}}_j\\ {\mathbf{h}}_j\end{array}\right]|\left[\begin{array}{c}{\mathit{\mu }}_k^o\\ {\mathit{\mu }}_k^h\end{array}\right],{\left[\begin{array}{cc}{\mathbf{\Lambda }}_k^{o,o}& {\mathbf{\Lambda }}_k^{o,h}\\ {\mathbf{\Lambda }}_k^{o,h}& {\mathbf{\Lambda }}_k^{h,h}\end{array}\right]}^{-1}\right)\text{,}$$ (C.2)

by making use of block matrix notation to partition the mean vector ${\mathit{\mu }}_k$ and the precision matrix ${\mathbf{\Lambda }}_k$.

In this case ${\mathbf{h}}_j$ is treated as an unobserved random variable. Thus, in a variational Bayes setting, an additional posterior factor can be introduced for each missing data point ${\mathbf{h}}_j$ to give

$$\begin{array}{l}q\left(\mathbf{H},\mathbf{Z},{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)=q\left(\mathbf{H}\right)q\left(\mathbf{Z}\right)q\left({\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)\\ =q\left(\mathbf{Z}\right)q\left({\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)\stackrel{N}{\prod _{j=1}}q\left({\mathbf{h}}_j\right)\text{.}\end{array}$$ (C.3)

Making use of the general result $q_{\hat{s}}\left({\text{Θ}}_{\hat{s}}\right)\propto \text{exp}\left({\mathbb{E}}_{s\ne \hat{s}}\left[\text{log}p\left(\mathbf{X},\text{Θ}\right)\right]\right)$ (Bishop, 2006), an approximated posterior on the missing data point ${\mathbf{h}}_j$ can be computed by

$$\begin{array}{l}\text{log}q\left({\mathbf{h}}_j\right)={\mathbb{E}}_{\mathbf{Z},{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[\text{log}p\left({\mathbf{x}}_j,{\mathbf{z}}_j,{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}|{\text{Θ}}_{\pi }\right)\right]+\text{const}\\ ={\mathbb{E}}_{\mathbf{Z},{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[\text{log}p\left({\mathbf{x}}_j|{\mathbf{z}}_j,{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)\right]+{\mathbb{E}}_{\mathbf{Z}}\left[\text{log}p\left({\mathbf{z}}_j|{\text{Θ}}_{\pi }\right)\right]\\ +{\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[\text{log}p\left({\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)\right]+\text{const},\end{array}$$ (C.4)

where ${\text{Θ}}_{\pi }$ denotes the mixing proportion parameter set, treated here via maximum likelihood, and $p\left({\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)$ is a conjugate Gaussian-Wishart prior on the means and covariances of the model.

Ignoring the terms independent from ${\mathbf{h}}_j$, equation (C.4) can be rewritten as

$$\begin{array}{l}\text{log}q\left({\mathbf{h}}_j\right)=\stackrel{K}{\sum _{k=1}}{\gamma }_{jk}{\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[\text{log}\mathcal{N}\left({\mathbf{x}}_j|{\mathit{\mu }}_k,{\mathbf{\Sigma }}_k\right)\right]+\text{const}\\ =\frac{1}{2}\stackrel{K}{\sum _{k=1}}{\gamma }_{jk}{\mathbf{h}}_j^T{\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathbf{\Lambda }}_k^{h,h}\right]{\mathbf{h}}_j\\ +\stackrel{K}{\sum _{k=1}}{\gamma }_{jk}{\mathbf{h}}_j^T{\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathbf{\Lambda }}_k^{o,h}\right]\left({\mathbf{o}}_j-{\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathit{\mu }}_k^o\right]\right)\\ -\stackrel{K}{\sum _{k=1}}{\gamma }_{jk}{\mathbf{h}}_j^T{\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathbf{\Lambda }}_k^{h,h}\right]{\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathit{\mu }}_k^h\right]+\text{const}\text{.}\end{array}$$ (C.5)

The previous equation indicates that the unobserved value ${\mathbf{h}}_j$ is drawn from a Gaussian mixture distribution with mixing proportions equal to the posterior (after having observed ${\mathbf{o}}_j$) membership probabilities ${\left\{{\gamma }_{jk}\right\}}_{k=1,\dots ,K}$, while the Gaussian means ${\left\{{\mathbf{n}}_{jk}\right\}}_{k=1,\dots ,K}$ and covariances ${\left\{{\mathbf{P}}_{jk}\right\}}_{k=1,\dots ,K}$ are given by

$$\begin{array}{l}{\mathbf{n}}_{jk}={\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathit{\mu }}_k^h\right]+{\left({\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathbf{\Lambda }}_k^{h,h}\right]\right)}^{-1}\\ \times {\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathbf{\Lambda }}_k^{o,h}\right]({\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathit{\mu }}_k^o\right]-{\mathbf{o}}_j),\end{array}$$ (C.6)

$${\mathbf{P}}_k={\mathbb{E}}_{{\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}}\left[{\mathbf{\Lambda }}_k^{h,h}\right]\text{.}$$ (C.7)

Given the posteriors $q\left(\mathbf{Z}\right)$ and $q\left(\mathbf{H}\right)$, the following sufficient statistics of $\mathbf{X}$ can be computed

$${\mathit{s}}_{1k}=\left[\begin{array}{c}\stackrel{N}{\sum _{j=1}}{\gamma }_{jk}{\mathbf{o}}_j\\ \stackrel{N}{\sum _{j=1}}{\gamma }_{jk}{\mathbf{n}}_{jk}\end{array}\right],$$ (C.8)

$${\mathit{S}}_{2k}=\left[\begin{array}{cc}\stackrel{N}{\sum _{j=1}}{\gamma }_{jk}{\mathbf{o}}_j{\mathbf{o}}_j^T& \stackrel{N}{\sum _{j=1}}{\gamma }_{jk}{\mathbf{o}}_j{\mathbf{n}}_{jk}^T\\ \stackrel{N}{\sum _{j=1}}{\gamma }_{jk}{\mathbf{n}}_{jk}{\mathbf{o}}_j^T& \stackrel{N}{\sum _{j=1}}{\gamma }_{jk}\left({\mathbf{n}}_k{\mathbf{n}}_{jk}^T+{\left({\mathbf{P}}_k\right)}^{-1}\right)\end{array}\right]\text{.}$$ (C.9)

Once such sufficient statistics have been evaluated, they can be used to update the Gaussian-Wishart posteriors $q\left({\text{Θ}}_{\mu },{\text{Θ}}_{\text{Σ}}\right)$ in the exact same way as in Blaiotta et al. (2016). Such posteriors are in turn used to compute the expectations that appear in equation (C.6), (C.7), in an iterative EM fashion.

An inner product naturally induces a norm by $\left|\right|\mathit{x}\left|\right|={\langle \mathit{x},\mathit{x}\rangle }^{1/2}$, therefore every inner product space is also a normed vector space (Dieudonné, 2013).