Unsupervised Deep Learning for Bayesian Brain MRI Segmentation

Abstract

Probabilistic atlas priors have been commonly used to derive adaptive and robust brain MRI segmentation algorithms. Widely-used neuroimage analysis pipelines rely heavily on these techniques, which are often computationally expensive. In contrast, there has been a recent surge of approaches that leverage deep learning to implement segmentation tools that are computationally efficient at test time. However, most of these strategies rely on learning from manually annotated images. These supervised deep learning methods are therefore sensitive to the intensity profiles in the training dataset. To develop a deep learning-based segmentation model for a new image dataset (e.g., of different contrast), one usually needs to create a new labeled training dataset, which can be prohibitively expensive, or rely on suboptimal ad hoc adaptation or augmentation approaches. In this paper, we propose an alternative strategy that combines a conventional probabilistic atlas-based segmentation with deep learning, enabling one to train a segmentation model for new MRI scans without the need for any manually segmented images. Our experiments include thousands of brain MRI scans and demonstrate that the proposed method achieves good accuracy for a brain MRI segmentation task for different MRI contrasts, requiring only approximately 15 seconds at test time on a GPU.

1. Introduction

Bayesian segmentation of medical images, particularly in the context of brain MRI scans, is a well-studied problem. Most probabilistic models for image segmentation exploit atlas priors, and account for variations in contrast and imaging artifacts such as MR inhomogeneity [32,34]. Most of the popular neuroimage processing pipelines rely on segmentation algorithms based on these ideas, including FreeSurfer [13], SPM [3], and FSL [28]. While these tools achieve high robustness to changes in MRI contrast of the input scan, a significant drawback is that they are computationally demanding (e.g., 23 minutes using a multithreaded setup, in a recent study [29]), which limits their deployment at scale and in time-sensitive applications. Therefore, there is a need for computationally efficient methods that are contrast-adaptive, requiring no additional labeled training images to segment a new dataset.

Recently, there has been a surge in the application of deep learning (DL) techniques to medical image segmentation, often based on convolutional neural network (CNN) architectures that excel at learning contextually important multi-scale features. An advantage of these methods is their computational efficiency at test (segmentation) time, offering the potential to use automatic segmentation in new application areas, such as those involving very large test datasets [19,30]. Moreover, these algorithms can be combined with atlas priors for increased robustness [9,22,26] However, DL based techniques are notoriously sensitive to changes in the image intensity data distribution. For example, an upgrade to the MRI scanner or a change in the pulse sequence might alter contrast properties that can dramatically reduce the performance of a CNN-based segmentation model [17]. This issue can be alleviated via domain adaptation, which usually requires some amount of labeled training data for the new conditions, or data augmentation, which requires the user to simulate expected variations. However, even with additional data, these methods only partially close the gap with the fully supervised setting [27]. Furthermore, the dependency on manually annotated datasets means that existing DL approaches are only applicable if enough resources are available to compile the required training data. This is often infeasible, for example in the context of continuously upgrading imaging technologies.

In this paper, we consider the scenario in which we have a general probabilistic atlas prior and a collection of images with no manual delineations. The probabilistic atlas is a volume where each voxel has an associated vector with the prior probabilities of observing the different segmentation labels at that location. Our approach assumes the availability of such an atlas (in brain imaging, they are readily available), and is independent of how it was created. For example, it could have been obtained by averaging a collection of manually annotated volumes of a different imaging modality. Alternatively, it could have been derived from an anatomical template, after applying spatial blurring to account for variability in location.

Several recent methods tackle segmentation tasks in the presence of small training datasets. Most assume at least one manually segmented image from the same modality as the main task, and leverage data augmentations techniques and exploit priors to enable the use of supervised methods [5,36]. Other methods require no labelled examples from the target modality, but leverage a large collection of segmentation maps from other datasets [9,18].

The main contribution of this paper is the integration of mathematical ideas from the Bayesian segmentation literature with an unsupervised deep learning framework. Specifically, we assume a probabilistic model, which requires estimation of scan-specific parameters comprising an atlas deformation and image intensity statistics. The estimation of the atlas warp has traditionally relied on classic deformable registration algorithms [31], which are based on iterative, numerical optimization, and are therefore computationally expensive. Instead, we leverage recent advances in learning-based registration [4,8,21,33] to efficiently estimate the warp jointly with the intensity parameters. We use a novel loss function, which is derived from the probabilistic model with Bayesian inference, and is thus principled and interpretable. Integrating DL with Bayesian segmentation, we attain two highly desirable features. First, given a probabilistic atlas, the method is unsupervised, and hence contrast adaptive: given a new dataset with a previously unobserved MRI contrast (e.g., a change of pulse sequence in MRI acquisition), we train our network without the need to label any MRI scans. Second, the segmentation is efficient and runs in approximately 15 seconds on a GPU.

2. Method

2.1. Segmentation as Bayesian inference

Let I represent the intensities of a 3D brain MRI scan, defined over a discrete domain $\Omega \subset {ℝ}^3$. Let S be a corresponding discrete segmentation into L neuroanatomical labels. Bayesian segmentation relies on Bayes’ rule to derive the posterior probability distribution of the segmentation S given the input image I. Then, the segmentation $\widehat{\mathit{S}}$ is estimated as the mode of this posterior:

$$\widehat{\mathit{S}}=\underset{\mathit{S}}{\arg \max }p(\mathit{S}|\mathit{I})=\underset{\mathit{S}}{\arg \max }p(\mathit{I}|\mathit{S})p(\mathit{S}).$$ (1)

The posterior distribution p(S|I) depends on two terms: a prior p(S) and a likelihood p(I|S). This is in contrast to discriminative segmentation approaches, which model p(S|I) directly. The prior represents knowledge about the spatial distribution of labels in the segmentation, and often has the form of a probabilistic atlas endowed with a deformation model. The likelihood models the relationship between the segmentation (i.e., underlying anatomy) and image intensities, including image artifacts such as noise and bias field. Both the prior and likelihood may have a set of associated parameters, which we define as θ_S and θ_I, respectively. The former describes attributes such as label probabilities and atlas deformation, while the latter typically includes image intensity statistics as a function of label and possibly location.

The likelihood parameters may be global for a training dataset, or estimated specifically for each test scan. Here we are interested in a subset of Bayesian segmentation models that follow the latter approach [3,29,32,34,35], which enables these models to adapt to the intensity characteristics of the input scans, making them robust to changes in MRI contrast. Expanding Eq. (1) to include model parameters, which we treat as random variables, yields:

$$\widehat{\mathit{S}}=\underset{\mathit{S}}{\arg \max }{\int }_{{\mathit{\theta }}_S}{\int }_{{\mathit{\theta }}_I}p\left(\mathit{S}|{\mathit{\theta }}_S,{\mathit{\theta }}_I,\mathit{I}\right)p\left({\mathit{\theta }}_S,{\mathit{\theta }}_I|\mathit{I}\right)d{\mathit{\theta }}_{\mathit{S}}d{\mathit{\theta }}_{\mathit{I}},$$ (2)

which is intractable. A standard approximation is to use point estimates for the parameters. First, one estimates the mode of the posterior distribution for the parameters:

$$\begin{gathered} \left\{{\widehat{\mathit{\theta }}}_S,{\widehat{\mathit{\theta }}}_I\right\}=\underset{\left\{{\mathit{\theta }}_S,{\mathit{\theta }}_I\right\}}{\arg \max }p\left({\mathit{\theta }}_S,{\mathit{\theta }}_I|\mathit{I}\right) \\ =\underset{\left\{{\mathit{\theta }}_S,{\mathit{\theta }}_I\right\}}{\arg \max }p\left({\mathit{\theta }}_S\right)p\left({\mathit{\theta }}_I\right)\sum _{\mathit{S}}p\left(\mathit{I}|\mathit{S},{\mathit{\theta }}_I\right)p\left(\mathit{S}|{\mathit{\theta }}_S\right), \end{gathered}$$ (3)

where we have assumed independence between the parameters of the prior and likelihood. The computation often requires estimating an atlas deformation in θ_S and intensity parameters in θ_I, and is typically achieved with a combination of numerical optimization and the Expectation Maximization (EM) algorithm [10]. Given point estimates, the final segmentation is computed efficiently as:

$$\widehat{\mathit{S}}=\underset{\mathit{S}}{\arg \max }p\left(\mathit{S}|{\widehat{\mathit{\theta }}}_S,{\widehat{\mathit{\theta }}}_I,\mathit{I}\right)=\underset{\mathit{S}}{\arg \max }p\left(\mathit{I}|\mathit{S},{\widehat{\mathit{\theta }}}_I\right)p\left(\mathit{S}|{\widehat{\mathit{\theta }}}_S\right),$$ (4)

and is often produced directly by the same EM algorithm.

2.2. Proposed Model

Our model instantiation builds on existing work [3,29,32]. The prior is defined by a given probabilistic atlas A, such that A(l, x) provides the probability of observing each neuroanatomical label l = 1, …, L at each location x ∈ Ω. The atlas is deformed by a diffeomorphic transform ϕ, parameterized by a stationary velocity field υ, (i.e., ϕ_υ = exp[υ], see [2]) that parametrizes the prior such that θ_S = υ. Assuming independence over voxels:

$$p\left(\mathit{S}|{\mathit{\theta }}_S;\mathit{A}\right)=p(\mathit{S}|\mathit{v};\mathit{A})=\prod _{j\in \Omega }A\left(S_j,{\varphi }_v\left({\mathit{x}}_j\right)\right),$$ (5)

where S_j is the segmentation at voxel j, and x_j is its spatial location.

We discourage strongly varying deformations by penalizing the spatial gradient ∇u_υ of displacement u_υ, where ϕ_υ = Id + u_υ:

$$p\left({\mathit{\theta }}_S;\lambda \right)=p(\mathit{v};\lambda )\propto \exp \left[-\lambda {\Vert \nabla {\mathit{u}}_v\Vert }^2\right].$$ (6)

The hyperparameter λ controls the weight for the atlas deformation penalty.

Conditioned on a segmentation, we assume that the observed intensities at different voxel locations are independent samples of Gaussian distributions:

$$p\left(\mathit{I}|\mathit{S},{\theta }_I\right)=p\left(\mathit{I}|\mathit{S},\mathit{\mu },{\sigma }^2\right)=\prod _{j\in \Omega }\mathcal{N}(I_j;{\mu }_{S_j},{\sigma }_{S_j}^2),$$ (7)

where $\mathcal{N}\left(\cdot ;\mu ,{\sigma }^2\right)$ is the Gaussian distribution, I_j is the image intensity at voxel j, and the likelihood parameters θ_I = {μ, σ²} are L means μ_l and variances ${\sigma }_l^2$, each associated with a different label l. We complete the model with a flat prior for these parameters: p(θ_I) ∝ 1. The model can be easily extended to the multispectral case (i.e., inputs with multiple MRI contrasts) by replacing means and variances by mean vectors and covariance matrices, respectively.

2.3. Learning

To avoid computationally expensive optimization typically required for maximum a posteriori (MAP) estimation in Eq. (3), we propose to train a CNN to estimate these parameters directly from an input scan. Specifically, we design a CNN $g_{{\mathit{\theta }}_C}(\mathit{I},\mathit{A})=\left({\mathit{\theta }}_S,{\mathit{\theta }}_I\right)=\left(\mathit{v},\mathit{\mu },{\mathit{\sigma }}^2\right)$ with convolutional parameters θ_C that takes as input a scan I and the probabilistic atlas A, and outputs the model parameters υ, μ, σ² for that scan. To learn the neural network parameters θ_c, we use a pool of N unlabeled scans ${\left\{I^n\right\}}_{n=1}^N$ to minimize the negative log posterior distribution of the image-specific parameters given the training images:

$$\begin{gathered} -\sum _{n=1}^N\log p\left({\mathit{v}}^n,{\mathit{\mu }}^n,{\left[{\mathit{\sigma }}^2\right]}^n|I^n;\mathit{A},\lambda \right) \\ =-\sum _{n=1}^N\sum _{j\in \Omega }\log \left[\sum _{l=1}^L\mathcal{N}\left(I_j^n;{\mu }_l^n,{\left[{\sigma }_l^2\right]}^n\right)A\left(l,{\varphi }_{v_m}\left({\mathit{x}}_j\right)\right)\right]+\lambda {\Vert \nabla {\mathit{u}}_v^n\Vert }^2 \\ -K(\lambda )+\text{const}, \end{gathered}$$ (8)

where K(λ) is a log-partition function that depends on the hyperparameter λ, and which does not affect the optimization. We emphasize that the network outputs different parameters μ, σ, and υ for each test image I.

We design the neural network $g_{{\mathit{\theta }}_C}(\cdot ,\cdot )$ based on a 3D UNet-style architecture [30] and the public VoxelMorph implementation [4]. The network consists of downsampling convolutional layers with 32 filters with 3×3 kernels, stride of 2, and LeakyReLu activations, followed by mirror upsampling layers and skip-connections. From this point, an additional convolutional layer is used to output υ, a dense 3D velocity field defined over Ω; and an additional pair of convolutional layers followed by a global max pooling operation to output the Gaussian parameters μ, σ². We compute ϕ = exp(υ) using a network integration layer that implements scaling and squaring [2,8,21], enabling the computation of the loss regularization term. We warp the probabilistic atlas A with a spatial transform layer. Combining the Gaussian parameters with the input image yields likelihood maps, which together with the warped atlas enable computation of the first term of the loss function (Fig. 1).

Fig. 1. Method overview.

The network block g_ψ(·,·) outputs a stationary velocity field υ, enabling alignment of the probabilistic atlas to the input volume, and likelihood Gaussian parameters μ, σ², which yield likelihood maps for each label.

2.4. Efficient segmentation

Given a trained network and a new test subject, the network efficiently provides the image-specific parameter point estimates $\widehat{\mathit{v}}$, and ${\widehat{\mathit{\theta }}}_I$ via a single forward pass. The optimal segmentation can be efficiently computed for each voxel:

$${\widehat{S}}_j=\underset{l}{\arg \max }\mathcal{N}\left(I_j;{\widehat{\mu }}_l,{\widehat{\sigma }}_l^2\right)A\left(l,{\varphi }_{\widehat{v}}\left({\mathit{x}}_j\right)\right).$$ (9)

Both terms in Eq. (9) are computed inside our GPU implementation (Fig. 1).

3. Experiments and results

3.1. Data

We evaluate our approach on three different image sets. The first dataset (“multi-site”) includes 8,332 T1-weighted scans from several public datasets: OASIS [23], ABIDE [11], ADHD200 [25], MCIC [15], PPMI [24], HABS [7], and Harvard GSP [16]. We randomly selected 7,332 scans to train and validate, and the remaining 1,000 were held out for testing. Manual delineations are not available for these scans, but we used automated segmentations produced by FreeSurfer [13] as a silver standard, for evaluation only. The second dataset (“T1”) consist of 38 T1-weighted scans, used only for testing, each with 36 manually delineated brain structures [13]. The third dataset (“PD”) consists of eight proton density-weighted (PD) scans, manually segmented with the same protocol [14]. All scans were preprocessed with FreeSurfer, including skull stripping, bias field correction, intensity normalization, affine registration to Talairach space, and resampling to 1 mm³ isotropic resolution [12].

3.2. Experimental setup

We perform three experiments, one for each dataset. In the first experiment, we fit our network to the 7,332 T1-weighted training scans of the multi-site dataset, and use the resulting model to segment the 1,000 test scans. Despite the lack of a manual gold standard, this experiment enables assessment of performance on a large, heterogeneous dataset. In a second experiment, we use the model already trained in the first experiment (i.e., on the 7,332 T1 scans) to segment scans from the separate T1 dataset. This experiment enables evaluation with manual ground truth on scans from a scanner and pulse sequence that were not observed by the neural network during training. In the third experiment, we train a network on the PD dataset, and then use it to segment those 8 PD scans. This is a different scenario than the first two experiments, since we learn to segment the test dataset directly. This experiment enables us to assess the ability of our algorithm to segment a substantially different MRI contrast, and to fit datasets of reduced size. In all experiments, we use our method with the publicly available atlas from [29]. We emphasize that all networks are trained in an unsupervised fashion, and segmentation maps are only used for evaluation.

3.3. Baseline

We compare our method to a reimplementation of [32], which relies on an affine version of the aforementioned atlas and Gaussian likelihood functions. Specifically, the baseline method solves Eq. (8), but with no deformation (i.e., υ = u = 0, and ϕ_υ = Id), and the model parameters are estimated with the EM algorithm. Since the model does not include deformation, using the nonrigid version of the atlas would yield very low performance.

3.4. Evaluation

We used Dice scores for a subset of structures of interest (Fig. 2). We quantify the results on these structures, and also focus on deep structures such as the hippocampus, which is the target of many neuroimaging studies due to its significance in dementia.

Fig. 2. Segmentation Statistics.

Dice scores for: cerebral cortex (CT) and white matter (WM); lateral ventricle (LV); cerebellar cortex (CC) and white matter (CW); thalamus (TH); caudate (CA); putamen (P); pallidum (PA); brainstem (BS); hippocampus (HP); and amygdala (AM). Scores of contralateral structures are averaged. The number of outliers under the x axis is shown in red (baseline) and blue (ours).

3.5. Implementation Details

We group anatomical labels with similar intensity properties into eleven merged labels to force groups of original labels to share Gaussian parameters, increasing robustness [29]. Specifically, we group: contralateral structures (in general), gray matter structures (cerebral gray matter, hippocampus, amygdala, caudate, accumbens), and cerebrospinal fluid structures.

We implement our method using Keras [6] with a Tensorflow [1] backend and the ADAM optimizer [20]. We predict the velocity field υ and resulting deformation field ϕ at every second voxel in each dimension, due to memory constraints. We linearly interpolate to obtain a final dense deformation field. To set λ, the only free parameter of our framework, we visually evaluated segmentation results for several validation subjects (held out from the training dataset), and set λ = 10 in all experiments.

3.6. Results

Our method requires only 15 seconds per scan on an NVIDIA Titan Xp GPU. Fig. 2 reports segmentation statistics for all experiments. Our method achieves considerably higher Dice scores than the baseline on the multi-site dataset (average over all structures 83.5% vs. 79.0%), particularly in deep brain structures, such as the hippocampi (81.1% vs. 73.1%). Moreover, it largely reduces the number of outliers with very poor segmentation (e.g., there are over 100 cases with Dice lower than 50% in the caudate for the baseline approach, and none for our method). In the T1 dataset, the test intensity distribution is slightly different that of the training dataset. However, our approach successfully generalizes and outperforms the baseline (average 81.9% vs. 79.4%, hippocampi 79.9% vs. 73.5%). The results of the third experiment illustrate the ability of our method to adapt to contrasts other than T1, even when the data are limited, and out-perform the baseline (average 80.5% vs. 78.3%, hippocampi 76.6% vs. 69.8%).

Figure 3 shows two segmentations from the T1 and PD datasets. In the T1 scan, the atlas successfully deforms to match the large ventricles of the subject, producing more accurate segmentations than the baseline – not only for the ventricles (purple), but also for surrounding structures, e.g., the thalami (green). In the PD scan, despite the small dataset, our method manages to segment all structures including the amygdalae (light blue), which are missed by the baseline.

Fig. 3. Example Results.

Coronal slices of two scans (one from each of the T1 and PD datasets), along with the initial and deformed probabilistic atlas, and corresponding segmentations. In the atlas, the color of each pixel is a combination of the colors of different labels, weighted by their probabilities. In the segmentations, we show the contour of the labels in the corresponding colors. We use the FreeSurfer color map [12].

4. Conclusion

We propose a principled approach for unsupervised segmentation, which enables training a CNN for a dataset without the need for any manually annotated images. The likelihood model may be extended to incorporate more complex functions (such as mixtures of Gaussians) and artifacts such as partial voluming and bias field. In addition to segmentations, the method produces a dense nonlinear deformation field that is a useful output by itself, e.g., for tensor-based morphometry. Using a large dataset, we demonstrate that the proposed approach achieves state-of-the-art accuracy for unsupervised brain MRI segmentation in different MRI contrasts. Our method runs in under 15 seconds on a GPU, facilitating deployment on large studies and in time-sensitive applications.