A Longitudinal Method for Simultaneous Whole-Brain and Lesion Segmentation in Multiple Sclerosis

The segmentation prior is composed of two components $p(\mathbf{l}|\mathbf{x})$ and $p(\mathbf{z}|\mathbf{h},\mathbf{x})$ that encode spatial information regarding the neuroanatomical labels $\mathbf{l}$ and the lesion map $\mathbf{z}$ respectively: $p(\mathbf{l},\mathbf{z}|\mathbf{h},\mathbf{x})=p(\mathbf{l}|\mathbf{x})p(\mathbf{z}|\mathbf{h},\mathbf{x}).$ The first component is a deformable probabilistic atlas, encoded as a tetrahedral mesh [17] with node positions $\mathbf{x}$ and with a deformation prior distribution defined as:

$$p(\mathbf{x})\propto\exp\left[-K\sum_{d}U_{d}(\mathbf{x},\mathbf{x}_{ref})\right].$$

Here $K$ controls the stiffness of the mesh deformations, $d$ loops over the tetrahedra in the mesh, and $U_{d}(\mathbf{x},\mathbf{x}_{ref})$ is a cost [18] associated with deforming the $d^{th}$ tetrahedron from its shape in the atlas’s reference position $\mathbf{x}_{ref}$. Letting $p(l_{i}=k|\mathbf{x})$ denote the probability of observing label $k$ at voxel $i$ for a given deformation, assuming conditional independence of the labels between voxels yields $p(\mathbf{l}|\mathbf{x})=\prod_{i=1}^{I}p(l_{i}|\mathbf{x}).$

The second component of the segmentation prior is a model of the form: $p(\mathbf{z}|\mathbf{h},\mathbf{x})=\prod_{i=1}^{I}p(z_{i}|\mathbf{h},\mathbf{x}),\textbf{ }p(\mathbf{h})=\mathcal{N}(\mathbf{h}|\mathbf{0},\mathbf{I}),$ where $p(z_{i}=1||\mathbf{h},\mathbf{x})$ is the probability that voxel $i$ is part of a lesion. This model takes into account both a voxel’s spatial location within its neuroanatomical context (through $\mathbf{x}$), as well as lesion shape constraints through a variational autoencoder (VAE) [19] that “decodes” a low-dimensional latent code $\mathbf{h}$ using a convolutional neural network.

For the likelihood, which links segmentations $\{\mathbf{l},\mathbf{z}\}$ to intensities $\mathbf{D}$, we use a multivariate Gaussian intensity model for each structure, and model the MRI bias field artifact as a linear combination of spatially smooth basis functions that is added to the local voxel intensities [20, 21]. Letting $\boldsymbol{\theta}_{z}=\{\boldsymbol{\mu}_{z},\boldsymbol{\Sigma}_{z}\}$ denote the mean and variance of lesion intensities, and $\boldsymbol{\theta}$ the collection of bias field parameters and intensity means and variances $\{\boldsymbol{\mu}_{k},\boldsymbol{\Sigma}_{k}\}$ of all $K$ anatomical structures, the likelihood is defined as $p(\mathbf{D}|\mathbf{l},\mathbf{z},\boldsymbol{\theta},\boldsymbol{\theta}_{z})=\prod_{i=1}^{I}p(\mathbf{d}_{i}|l_{i},z_{i},\boldsymbol{\theta},\boldsymbol{\theta}_{z}),$ where

$\displaystyle p(\mathbf{d}_{i}|l_{i}=k,z_{i},\boldsymbol{\theta},\boldsymbol{\theta}_{z})=\begin{cases}\mathcal{N}(\mathbf{d}_{i}|\boldsymbol{\mu}_{z}+\mathbf{C}^{T}\boldsymbol{\phi}_{i},\boldsymbol{\Sigma}_{z})&\text{if }z_{i}=1,\\ \mathcal{N}(\mathbf{d}_{i}|\boldsymbol{\mu}_{k}+\mathbf{C}^{T}\boldsymbol{\phi}_{i},\boldsymbol{\Sigma}_{k})&\text{otherwise}.\end{cases}$

Here $\boldsymbol{\phi}_{i}$ evaluates the bias field basis functions at the $i^{th}$ voxel, and $\mathbf{C}=(\mathbf{c}_{1},\dots,\mathbf{c}_{N})$, where $\mathbf{c}_{n}$ denotes the parameters of the bias field model for the $n^{th}$ contrast. The model is completed by a flat prior on $\boldsymbol{\theta}$, and a weak conditional prior $p(\boldsymbol{\theta}_{z}|\boldsymbol{\theta})$ that ensures that the method can be robustly applied to scans with no or very small lesion loads [16].

Given an MRI scan $\mathbf{D}$, segmentation proceeds by approximating the segmentation posterior using point estimates of the parameters $\mathbf{x}$ and $\boldsymbol{\theta}$:

$$p(\mathbf{l},\mathbf{z}|\mathbf{D})\simeq p(\mathbf{l},\mathbf{z}|\mathbf{D},\hat{\boldsymbol{\theta}},\hat{\mathbf{x}}),$$

(1)

and Markov chain Monte Carlo sampling to marginalize over the remaining, lesion-specific parameters $\boldsymbol{\theta}_{z}$ and $\mathbf{h}$. For the purpose of finding the point estimates $\mathbf{\hat{x}}$ and $\boldsymbol{\hat{\theta}}$, a simplified model is fitted to the data:

$$\boldsymbol{\hat{\Omega}}=\operatornamewithlimits{argmax}_{\boldsymbol{\Omega}}p(\boldsymbol{\Omega}|\mathbf{D})\quad\mathrm{with}\quad\boldsymbol{\Omega}=\{\mathbf{x},\boldsymbol{\theta},\boldsymbol{\theta}_{z}\},$$

(2)

where the lesion-shape encoding VAE and its parameters $\mathbf{h}$ are temporarily removed to simplify the optimization process. More details can be found in [16].

In order to obtain temporal consistency in the segmentation prior, we use the concept of a “subject-specific atlas” [22]: a deformation of the cross-sectional atlas to represent the average subject-specific anatomy across all time points. In particular,

$$p(\{\mathbf{x}_{t}\}_{t=1}^{T}|\mathbf{x}_{0})=\prod_{t=1}^{T}p(\mathbf{x}_{t}|\mathbf{x}_{0}),\quad p(\mathbf{x}_{t}|\mathbf{x}_{0})\propto\exp\left[-K_{1}\sum_{d}U_{d}(\mathbf{x}_{t},\mathbf{x}_{0})\right],$$

where $\mathbf{x}_{0}$ are latent atlas node positions encoding subject-specific brain shape, with prior $p(\mathbf{x}_{0})\propto\exp\left[-K_{0}\sum_{d}U_{d}(\mathbf{x}_{0},\mathbf{x}_{ref})\right].$ Here the mesh stiffnesses $K_{0}$ and $K_{1}$ are hyperparameters of the model; by choosing $K_{0}=\infty$ and $K_{1}=K$ the model devolves into the cross-sectional segmentation prior for each time point separately.

In a similar vein, we also introduce subject-specific latent variables to encourage temporal consistency in the Gaussian intensity models. For each anatomical structure, we condition the Gaussian parameters $\{\boldsymbol{\mu}_{t,k},\boldsymbol{\Sigma}_{t,k}\}$ on latent variables $\{\boldsymbol{\mu}_{0,k},\boldsymbol{\Sigma}_{0,k}\}$ using a normal-inverse-Wishart (NIW) distribution: $p(\{\boldsymbol{\theta}_{t}\}_{t=1}^{T}|\boldsymbol{\theta}_{0})=\prod_{t=1}^{T}p(\boldsymbol{\theta}_{t}|\boldsymbol{\theta}_{0})$ with

$\displaystyle p(\boldsymbol{\theta}_{t}|\boldsymbol{\theta}_{0})\propto\prod_{k=1}^{K}\mathcal{N}(\boldsymbol{\mu}_{t,k}|\boldsymbol{\mu}_{0,k},P_{0,k}\boldsymbol{\Sigma}_{0,k})\mathrm{IW}(\boldsymbol{\Sigma}_{t,k}|P_{0,k}\boldsymbol{\Sigma}_{0,k},P_{0,k}-N-2).$

Here $\boldsymbol{\theta}_{0}=\{\boldsymbol{\mu}_{0,k},\boldsymbol{\Sigma}_{0,k}\}_{k=1}^{K}$ with prior $p(\boldsymbol{\theta}_{0})\propto 1$, and $P_{0,k}\geq 0$ is a hyperparameter that governs the strength of the regularization across time for label $k$. Note that choosing $P_{0,k}=0,\forall k$ yields the cross-sectional likelihood for each time point independently.

We follow the same overall segmentation strategy as in the cross-sectional setting: we first compute point estimates $\{\boldsymbol{\hat{\theta}}_{t},\mathbf{\hat{x}}_{t}\}_{t=1}^{T}$ using a simplified model in which the lesion shape codes $\{\mathbf{h}_{t}\}_{t=1}^{T}$ are removed, and subsequently obtain segmentations as described in the cross-sectional setting, i.e., by using (1) for each time point separately. As in the cross-sectional case, we obtain the required point estimates by fitting the longitudinal model to the data:

$$\{\boldsymbol{\hat{\Omega}}_{1},\ldots,\boldsymbol{\hat{\Omega}}_{T},\boldsymbol{\hat{\theta}}_{0},\mathbf{\hat{x}}_{0}\}=\operatornamewithlimits{argmax}_{\{\boldsymbol{\Omega}_{1},\ldots,\boldsymbol{\Omega}_{T},\boldsymbol{\theta}_{0},\mathbf{x}_{0}\}}p(\boldsymbol{\Omega}_{1},\ldots,\boldsymbol{\Omega}_{T},\boldsymbol{\theta}_{0},\mathbf{x}_{0}|\mathbf{D}_{1},\ldots,\mathbf{D}_{T}),$$

(3)

where $\boldsymbol{\Omega}_{t}=\{\mathbf{x}_{t},\boldsymbol{\theta}_{t},\boldsymbol{\theta}_{t,z}\}$. For optimizing (3) we use coordinate ascent, updating one variable at a time in an iterative fashion. Because $p(\mathbf{x}_{t}|\mathbf{x}_{0})$ is of the same form as the cross-sectional deformation prior, and the NIW distribution used in $p(\boldsymbol{\theta}_{t}|\boldsymbol{\theta}_{0})$ is the conjugate prior for the mean and variance of a Gaussian distribution, estimating $\boldsymbol{\Omega}_{t}$ from $\mathbf{D}_{t}$ for given values of $\boldsymbol{\theta}_{0}$ and $\mathbf{x}_{0}$ simply involves performing an optimization of the form of (2) for each time point $t$ separately. Conversely, for given values $\{\boldsymbol{\Omega}_{t}\}_{t=1}^{T}$ the update for $\boldsymbol{\theta}_{0}$ is given in closed form:

$\displaystyle\boldsymbol{\mu}_{0,k}\leftarrow\left(\sum_{t=1}^{T}\boldsymbol{\Sigma}_{t,k}^{-1}\right)^{-1}\sum_{t=1}^{T}\boldsymbol{\Sigma}_{t,k}^{-1}\boldsymbol{\mu}_{t,k},\quad\boldsymbol{\Sigma}^{-1}_{0,k}\leftarrow\left(\frac{1}{T}\sum_{t=1}^{T}\boldsymbol{\Sigma}_{t,k}^{-1}\right)\frac{P_{0,k}}{P_{0,k}-N-2},$

whereas updating $\mathbf{x}_{0}$ involves the optimization (cf. [22])

$$\operatornamewithlimits{argmin}_{\mathbf{x}_{0}}\sum_{d}\left[K_{0}U_{d}(\mathbf{x}_{0},\mathbf{x}_{ref})+K_{1}\sum_{t=1}^{T}U_{d}(\mathbf{x}_{0},\mathbf{x}_{t})\right],$$

which we solve numerically using a limited-memory BFGS algorithm.

In order to avoid longitudinal processing biases resulting from not treated all time points in exactly the same way, we first compute an unbiased within-subject template using an inverse consistent registration method [23]. This template is a robust representation of the average subject anatomy over time, and we use it as an unbiased reference to register all time points to in a preprocessing step. We also use it to start the proposed iterative algorithm optimizing (3): we apply the cross-sectional method to the template, and use the estimated model parameters $\boldsymbol{\Omega}$ to initialize $\boldsymbol{\Omega}_{t},t=1,\ldots,T$. The proposed algorithm, which interleaves updating the latent variables $\boldsymbol{\theta}_{0}$ and $\mathbf{x}_{0}$ with updating the parameters $\{\boldsymbol{\Omega}_{t}\}_{t=1}^{T}$, is then run for five iterations, which we have found to be sufficient to reach convergence.

Based on initial pilot experiments on scans from the ADNI project¹¹1http://adni.loni.usc.edu/ (distinct from the ones used in the experiments below), we set the method’s hyperparameter values to $K_{1}=14K$ and $K_{0}=K$, where $K$ is the mesh stiffness in the existing cross-sectional method, and $P_{0,k}$ to the number of voxels assigned to class $k$ in the segmentation of the within-subject template.

Our implementation builds upon the C++ and Python code of [24, 16], and is publicly available from FreeSurfer²²2http://freesurfer.net/. Segmenting one subject takes approximately 15 minutes per time point on an Intel 12-core i7-8700K processor with a GeForce GTX 1060 graphics card.

We wished to assess whether the proposed method is able to reduce non-biological variations in longitudinal volume measurements, both within the short ($<3$ weeks) and longer ($<6$ years) time intervals of the test-retest and the Achieva datasets, respectively. For the test-retest dataset one can expect true biological changes to be minimal, and we therefore computed the coefficient of variation (ratio of the standard deviation to the mean) for each brain structure. The results, shown in Table 3, indicate that the longitudinal method indeed performs better in this respect than the cross-sectional one for almost all the structures.

For the Achieva dataset, one may assume that the true change in volume of a structure over the span of a few years is approximately linear, except for lesions whose temporal trajectory is affected by disease effects, with growing and shrinking lesions occurring at the same time. We therefore fitted, for each structure and for each subject, a linear regression model to the longitudinal volumes estimated by each method, and computed the ratio of the standard deviation of the residuals to the intersect (time of the first scan is taken as zero). The results are summarized in Table 3 and indicate that the proposed model indeed yields generally better results in this respect.

In order to ensure that the proposed method is not simply over-regularizing, we also assessed whether it can capture known group differences in the temporal evolution of certain brain structures better than the cross-sectional method. Towards this end, we compared the annualized percentage change (the slope of a linear regression model divided by its intersect) in the volume of the hippocampus between the CN, MCI and AD groups in the ADNI dataset. The results, shown in Fig. 3, indicate that the longitudinal method can indeed detect group differences better this way.

This project has received funding from the European Union’s Horizon 2020 research and innovation program under the Marie Sklodowska-Curie project TRABIT agreement No 765148, the German Research Foundation No 428223038, and the NIH NINDS No R01NS112161.