A cross-center smoothness prior for variational Bayesian brain tissue segmentation

TL;DR

This paper introduces a smoothness prior for brain tissue segmentation in MRI that leverages cross-center data to improve unsupervised Bayesian segmentation, addressing generalization issues caused by acquisition differences.

Contribution

It proposes a novel cross-center smoothness prior for unsupervised Bayesian segmentation, enhancing generalization across different medical centers without requiring labeled data.

Findings

The prior improves segmentation smoothness consistency across centers.

The semi-supervised model eliminates the need for manual cluster interpretation.

The approach enhances cross-center tissue segmentation accuracy.

Abstract

Suppose one is faced with the challenge of tissue segmentation in MR images, without annotators at their center to provide labeled training data. One option is to go to another medical center for a trained classifier. Sadly, tissue classifiers do not generalize well across centers due to voxel intensity shifts caused by center-specific acquisition protocols. However, certain aspects of segmentations, such as spatial smoothness, remain relatively consistent and can be learned separately. Here we present a smoothness prior that is fit to segmentations produced at another medical center. This informative prior is presented to an unsupervised Bayesian model. The model clusters the voxel intensities, such that it produces segmentations that are similarly smooth to those of the other medical center. In addition, the unsupervised Bayesian model is extended to a semi-supervised variant, which needs no visual interpretation of clusters into tissues.

Tables

Table 1. Table 1: Mean classification error and standard errors of the means (in brackets) of each of the segmentation models on all pairwise combinations of one data set as the source (rows) and another as the target (columns).

	Methods	Brainweb1.5T	MRBrainS13	IBSR
Brainweb1.5T	U-net	-	0.448 (.008)	0.384 (.019)
	1NN	0.117 (.027)	0.288 (.019)	0.518 (.075)
	UGM	0.142 (.044)	0.339 (.023)	0.525 (.088)
	SGM	0.116 (.031)	0.268 (.017)	0.527 (.089)
	UHP	0.147 (.053)	0.337 (.023)	0.511 (.087)
	SHP	0.117 (.032)	0.253 (.017)	0.519 (.078)
MRBrainS13	U-net	0.257 (.003)	-	0.589 (.022)
	1NN	0.103 (.010)	0.282 (.018)	0.513 (.075)
	UGM	0.112 (.024)	0.345 (.022)	0.521 (.090)
	SGM	0.102 (.011)	0.282 (.021)	0.502 (.093)
	UHP	0.119 (.030)	0.344 (.018)	0.507 (.091)
	SHP	0.102 (.008)	0.277 (.020)	0.503 (.076)
IBSR	U-net	0.334 (.007)	0.425 (.015)	-
	1NN	0.102 (.005)	0.282 (.064)	0.492 (.034)
	UGM	0.125 (.037)	0.369 (.071)	0.508 (.039)
	SGM	0.103 (.015)	0.260 (.043)	0.502 (.039)
	UHP	0.123 (.023)	0.350 (.068)	0.509 (.041)
	SHP	0.103 (.008)	0.255 (.049)	0.496 (.041)

Equations

$$\displaystyle p(X\>|\>Y;\ \pi,\mu,\Lambda)=\ \prod_{i=1}^{N}\prod_{k=1}^{K}\big{[}\pi_{k}\ \mathcal{N}(x_{i}\>|\>\mu_{k},\Lambda_{k}^{-1})\big{]}^{y_{ik}}\,.$$

$${\pi }_k\sim \mathcal{D}({\alpha }_{0k}),{\mu }_k\sim \mathcal{N}({\upsilon }_{0k},({\gamma }_{0k}{\Lambda }_k{)}^{-1}),{\Lambda }_k\sim \mathcal{W}({\nu }_{0k},{\Delta }_{0k}).$$

$$p(Y|\beta )=\prod _{i=1}^{N}p(y_i|y_{{\delta }_i},\beta ).$$

$$\log p(y_i|y_{{\delta }_i},\beta )=$$

$$=$$

$$p(X,Y,$$

$$=p(X|Y,\pi ,\mu ,\Sigma )p(Y|\beta )p(\pi |\alpha )p(\mu |\upsilon ,(\gamma \Lambda {)}^{-1})p(\Lambda |\nu ,\Delta ).$$

$$\log p(X|\beta )=$$

$$=$$

$$\ge$$

$$\mathcal{L}(q)\propto$$

$$=$$

$$\displaystyle\log\$$

$$\propto$$

$$\displaystyle\mathbb{E}_{\pi_{k}}\big{[}\log\pi_{k}\big{]}=\psi(\alpha_{k})-\psi\big{(}\sum_{k=1}^{K}\alpha_{k}\big{)}\,,$$

$$\displaystyle\mathbb{E}_{\mu_{k},\Lambda_{k}}\big{[}\log{\cal N}(x_{i}\>|\>\mu_{k},\Lambda_{k}^{-1})\big{]}=-\frac{D}{2}\log 2\mathrm{\pi}+\frac{1}{2}\mathbb{E}_{\Lambda_{k}}\big{[}\log|\Lambda_{k}|\ \big{]}-\frac{1}{2}\mathbb{E}_{\mu_{k},\Lambda_{k}}\big{[}\tilde{x}_{ik}\big{]}\,,$$

$$\displaystyle\mathbb{E}_{\Lambda_{k}}\big{[}\log|\Lambda_{k}|\ \big{]}=\ \sum_{d=1}^{D}\psi\big{[}(\nu_{k}+1-d)/2\big{]}+D\log 2+\log|\Delta_{k}|\,,$$

$$\displaystyle\mathbb{E}_{\mu_{k},\Lambda_{k}}\big{[}\tilde{x}_{ik}\big{]}=\ \frac{D}{\gamma_{k}}+\nu_{k}(x_{i}-\upsilon_{k})\Delta_{k}(x_{i}-\upsilon_{k})^{\top}\,,$$

$$\mathcal{L}(q)\propto$$

$$=$$

$${\alpha }_k=$$

$${\upsilon }_k=$$

$${\Delta }_k^{-1}=$$

$$S_k^0=\sum _{i=1}^N{\rho }_{ik},S_k^1=\sum _{i=1}^N{\rho }_{ik}{x}_i,S_k^2=\sum _{i=1}^N{\rho }_{ik}(x_i-S_k^1)(x_i-S_k^1{)}^{\top }.$$

$$\displaystyle p(X,\tilde{Y}\>|\>Y;\ \theta)=\prod_{j\in{O}}\prod_{i\centernot\in{O}}\prod_{k=1}^{K}\big{[}\pi_{k}\mathcal{N}(x_{j}\>|\>\mu_{k},\Lambda_{k}^{-1})\big{]}^{\tilde{y}_{jk}}\big{[}\pi_{k}$$

$$\hat{\beta }=\underset{\beta \in {\mathbb{R}}^{+}}{\mathrm{argmax}}\sum _{i=1}^N\log p(y_i|y_{{\delta }_i},\beta ).$$

$$\frac{\partial }{\partial \beta }\log p(y_i|y_{{\delta }_i},\beta )=$$

Full text

1 Introduction

Many modern automatic brain tissue segmentation methods are based on machine learning models. One of the limitations of these models is that they generalize poorly beyond the domain of the data they are trained on. In medical imaging, an example of a domain is the medical center itself. Data collected at different centers varies due to experimental, acquisition and annotation protocols. Most notably, the voxel intensity distributions of MR images are different, which means the mapping from scans to segmentations differs between centers. As a result, tissue classification models trained on examples from one center tend to perform poorly on data from another center [21].

But not all factors of variation are inconsistent across centers. Although the segmentations are different for each patient, certain aspects remain consistent. One such aspect is the spatial smoothness of each tissue. Radiologists and MR imaging experts know how smooth a segmentation is supposed to look like, and use this knowledge when segmenting a new scan. Essentially, we would like to give the tissue classifier that information as well. Our goal is to learn from segmentations at other medical centers and incorporate that knowledge into a Bayesian model for tissue segmentation.

1.1 Related Work

In transfer learning and domain adaptation, a model learns from a source domain and aims to generalize to a differently distributed target domain [16, 10]. In Bayesian transfer learning, the source domain can be interpreted as prior knowledge for the target task [17, 8]. For instance, in natural language processing, a document classification task can be performed using a Bayesian linear classifier trained on a bag-of-word encoding of the document [17]. Instead of imposing a weakly informative prior on how important each word of the dictionary is for the document classification task, one could fit the prior on data from Wikipedia. That produces a stronger, more informative prior over how important each word is. To our knowledge, no Bayesian transfer learning models have been proposed for medical imaging tasks. Our interest is to study what forms of prior knowledge can be obtained from large open access labeled data sets, and how that knowledge can be exploited for a specific task.

Hidden Markov Random Field (MRF) models are a form of Bayesian models for image segmentation. They pose a hidden state for each voxel that accounts for some intrinsic latent structure of the image [22]. For tissue segmentation, the latent state is assumed to be the tissue of the voxel, while the observed voxel intensity value is a sample from a probabilistic observation model. The observation model specifies the causal relations between the latent image and the observed image [24, 1]. Such assumptions are not unreasonable for the case of MR imaging, where T1 relaxation times depend on the tissue of the voxel.

Inference in hidden MRF’s is often done through Monte Carlo sampling [23]. However, sampling remains a computationally expensive procedure. An alternative is to use variational inference, where the joint distribution of an intractable Bayesian model is approximated [4, 7]. Variational inference is often much faster than sampling, depending on the form of the approximating distribution. We will employ variational Bayes to infer the underlying tissues of an observed MRI scan.

1.2 Outline

In Section 2, we will discuss a Bayesian model for tissue segmentation along with a hidden Markov Random Field prior. The variational approximation and the general inference procedure is presented in Section 2.1. Section 3 covers how to fit the MRF prior to segmentations produced at other medical centers. We perform a series of experiments in Section 4 where we pair up data sets from different medical centers. Model extensions and limitations are discussed in Section 5 and we draw conclusions in Section 6.

2 Method

Let $X\in[0,1]^{H\times W\times D}$ be an MR image and $Y\in\{0,1\}^{H\times W\times K}$ be its segmentation. $H$ and $W$ are the height of the width of the image, respectively, with $N=H\cdot W$ as the total number of voxels. $D$ refers to the number of channels of the image, which could be stacked filter response maps or additional modalities. In this paper, we consider only the MR image (i.e. D=1), but the update equations in Section 2.3 are general. $K$ corresponds to the number of tissues in the segmentation, also referred to as classes. Observed voxels are marked as $x$. Voxel labels are marked as $y$ and consist of $\{0,1\}$-valued vectors with $1$ on the $k$-th index if that voxel belongs to class $k$ (a.k.a. one-hot vectors).

2.1 Bayesian model

We assume a causal model $Y\rightarrow X$, such that the tissue causes the voxel intensity value. The measurement instrument, i.e. the MRI scanner, maps tissues to voxel intensities $f$, but imposes noise on the observation: $x=f(y)+\epsilon$. The mapping $f$ between $Y$ and $X$ is assumed to vary across experimental and acquisition protocols. We model the likelihood function of observing $X$ from $Y$ with a Gaussian mixture model, with one component for each tissue:

[TABLE]

The parameter $\pi_{k}$ is the proportion coefficient, $\mu_{k}$ is the mean intensity and $\Lambda_{k}$ is the precision of the $k$-th tissue. Note that this likelihood assumes that voxels are independent of each other, which is not valid in MR images. We model spatial relationships in Section 2.2 which introduces dependencies between pixels.

We select a Dirichlet distribution as the prior for the tissue proportions and a Normal-Wishart as the prior for the mean and precision parameters:

[TABLE]

The $\alpha_{0}$ are called the Dirichlet distribution’s concentration parameters, $\upsilon_{0}$ the hypermeans, $\gamma_{0}$ are precision-scaling hyperparameters, $\nu_{0}$ are the degrees of freedom of the Wishart distribution and $\Delta_{0}$ are the hyperprecisions. These priors are conjugate to the Gaussian likelihood.

2.2 Hidden Potts - Markov Random Field

Spatial properties of images can be described using Markov Random Fields. In general, MRF’s describe interactions between nodes in a graph by defining a probability distribution – to be precise, a Gibbs distribution – over configurations of states at the nodes [22]. The Markov property allows us to model this distribution in terms of local interactions. The Ising model is a classical MRF model, which describes the pairwise interactions between a binary-valued image pixel and its direct neighbours (i.e. up, down, left, right). The Potts model is its multivariate extension, using $K$ states.

We use the Potts model to capture how often a voxel’s label is equal to the labels of its neighbours. In other words, how smooth the segmentation is. The model incorporates a set of parameters, $\beta=(\beta_{1},\dots,\beta_{K})$, that explicitly describes each tissue’s smoothness. By fitting a hidden Potts model to a series of segmentations, it can act as an informative prior in the Bayesian model – a point we discuss in more detail in Section 3.

Hidden Potts models are usually defined for whole images. However, that induces a partition function with a discrete sum over $K^{N}$ states, which is computationally intractable. Instead, we consider a local variant, where voxels depend only on their direct neighbours $\delta_{i}$ [13, 12]:

[TABLE]

Voxels in the center of the image have four neighbours (i.e. up, down, left, right), while edge and corner voxels have three and two neighbours, respectively. Taking its logarithm, the Potts model has the following form:

[TABLE]

The sum with the subscript $\{y^{\prime}\}$ in (2) denotes summing over all possible states of $y$ (i.e. $[1,0,\dots 0],\ [0,1,\dots,0],\ \dots,\ [0,0,\dots,1]$). Since $y$ is a one-hot vector, it multiplies the terms in the sum that involve the $k$-th tissue with $1$ and multiplies the other terms with [math]. All but one term drop out, which means the sum over $\{y^{\prime}\}$ can be simplified to a sum over classes, as in (3).

2.3 Variational approximation

The hidden Potts-MRF describes spatial relationships in the segmentation and acts as a prior on the Gaussian mixture model. Including the hidden Potts model, the joint distribution of the full model becomes:

[TABLE]

In the following, the likelihood parameters are summarized as $\theta=(\pi,\mu,\Lambda)$. With the inclusion of the hidden Potts-MRF, the posteriors cannot be derived analytically. We perform a variational approximation of the joint distribution using a distribution over the segmentation and the likelihood parameters, $q(Y,\theta\>|\>\beta)$ [4]. This approximation relates to the marginal log-likelihood as follows:

[TABLE]

${\cal L}(q)$ is a function of the approximating distribution $q$ and is called the evidence lower bound. Here, the dependence on the hyperparameters is left out for notational convenience. We only maintain the dependence on $\beta$, as it is of importance in Section 3. In this framework, the objective is to find a parametric form for the variational approximation distribution $q(Y,\theta\>|\>\beta)$ such that it matches the true distribution as well as possible [4].

For computational reasons, we make the mean-field assumption that the segmentation and the likelihood parameters are independent of each other: $q(Y,\theta\>|\>\beta)=q(Y\>|\>\beta)\ q(\theta)$ [4]. The optimal form of each factor can be found by dropping terms in the lower bound that do not depend on the factor in question (as they are constants in the optimization procedure), and deriving the analytical solutions to the remaining expectations. For latent factor $q(Y\>|\>\beta)$, terms in the numerator and denominator of (5) not involving $Y$ and $\beta$ are ignored, producing:

[TABLE]

Note that in (6), the expectation with respect to the other factor, $q(\theta)$, is moved to the numerator. It can now be seen that the latent factor is optimal when: \log q^{*}(Y\>|\>\beta)\propto\mathbb{E}_{\theta}\big{[}\log p(X\>|\>Y,\theta)\big{]}+\log p(Y\>|\>\beta). Using the full Bayesian model specified in (4) and the hidden Potts from (3), the $i$-th voxel of the segmentation factor can be written as [4]:

[TABLE]

where

[TABLE]

and $\tilde{x}_{ik}=(x_{i}-\mu_{k})\Lambda_{k}(x_{i}-\mu_{k})^{\top}$. $\psi$ refers to the digamma function.

In Equation 7, we recognize a multinomial distribution: $\log q^{*}(y_{i}\>|\>\beta)\propto\ \sum_{k}^{K}y_{ik}\log r_{ik}$. The proportionality is due to the ignored terms. As these terms only serve to normalize the probabilities to the $[0,1]$ interval, we can replace their computation by the following normalization: $\rho_{i}=r_{ik}/\sum_{k=1}^{K}r_{ik}$ [4]. The $\rho_{ik}$ are called the responsibilities, referring to the probability for the $i$-th voxel to belong to the $k$-th class. Note that $\beta$ has not been integrated out. It will be estimated in a cross-medical center fashion (see Section 3).

Similar steps are taken to compute an optimal form for $q(\theta)$. This time, we ignore all terms in the ratio in (5) that do not depend on $\theta$:

[TABLE]

The factor $q(\theta)$ is optimal when: \log q^{*}(\theta)\propto\mathbb{E}_{Y}\big{[}\log p(X\>|\>Y,\theta)\ \big{]}+\log p(\theta). This is a well-known result (the choice of a Gaussian likelihood with conjugate priors is made often) and extensive derivations are widely available [4]. It produces the following update equations:

[TABLE]

where parameters with the subscript [math] belong to the priors and

[TABLE]

Note that the smoothness parameters $\beta$ affect these hyperparameter estimates through the estimates of the responsibilities $\rho_{ik}$.

Inference consists of iteratively computing the responsibilities based on the current posterior hyperparameters followed by updating the posterior hyperparameters given the new responsibilities. This procedure, known as variational Bayes, is halted when the change in values between iterations is smaller than a set threshold [4, 13, 5].

2.4 Semi-supervised model

Unsupervised models are limited by the fact that cluster assignments are not tied to tissue labels. Manually labeling one voxel per tissue overcomes this limitation. In order to incorporate the given voxel labels, a split in the likelihood function between labeled samples and unlabeled samples needs to be introduced [11]:

[TABLE]

where $\tilde{Y}$ are the observed labels, $Y$ are the unobserved labels and $O\subset[1,\dots N]$ is the subset of indices that are observed.

To derive a semi-supervised hidden Potts Gaussian mixture requires substituting the unsupervised likelihood from (1) in the Bayesian model in (4) with the above semi-supervised likelihood. Using the same derivations as throughout Section 2, this results in equivalent update equations for both the $q(Y\>|\>\beta)$ and $q(\theta)$ with the following exception: the responsibilities of the observed voxels are fixed to $\rho_{ik}=1.0$ if class $k$ was observed, and to $\rho_{im}=0.0$ for $m\neq k$. These responsibilities remain fixed throughout the variational optimization procedure.

2.5 Initialization of posterior hyperparameters

Variational inference is a form of non-convex optimization, which means that different initializations lead to different local optima. Several initializations for variational mixture models have been proposed, most notably $k$-means for unsupervised Gaussian mixtures [15]. In that case, the responsibilities $\rho$ of a point are set by the negative exponential of the distance to each cluster center.

For the semi-supervised model, the responsibilities can be initialized based on the distance to the given labeled voxels. This corresponds to $k$-nearest-neighbour classifier. As long as the number of labeled pixels is small, this remains computationally efficient.

3 Cross-center empirical Bayes

The Potts model can be fit to other segmentations using a maximum likelihood approach. First, it is treated as a likelihood function in its own right, with $Y$ as the observed variable dependent on the smoothing parameters $\beta$. Using the log-likelihood, the estimator becomes:

[TABLE]

This log-likelihood function is convex in $\beta$, which means the optimal smoothing parameters can be obtained using gradient descent. Its partial derivative with respect to $\beta$ is:

[TABLE]

where $\bar{y}_{ik}=\sum_{j\in\delta_{i}}y_{jk}$. Using a constrained optimization procedure, where all $\beta$’s are lower bounded by [math], we can obtain a point estimate for each $\beta_{k}$.

4 Experiments

We perform a series of cross-center brain tissue segmentation experiments. The goal is to assign each pixel in the MR image the label ”background”, ”cerebro-spinal fluid”, ”gray matter”, or ”white matter”. One data set will act as the source and another as the target. For the sake of comparison, we include single-center experiments, where $\beta$ cannot be learned and is subsequently set to $0.1$. All classification errors are computed using the brain mask. In other words, we ignore all mistakes in the skull and outlying regions. With the unsupervised models, each cluster is interpreted as one of the tissues, so that classification errors can be computed. To test the performance gain that can be achieved with a small amount of supervision, we provide the semi-supervised models with 1 voxel label per tissue from the target image, sampled at random. The experiments are repeated 10 times.

4.1 Data sets

We will make use of 3 publicly available data sets: Brainweb1.5T, MRBrainS13, and IBSR. Each data set originates from one medical center. Brainweb1.5T is based on 20 realistic phantoms from Brainweb [2] and an MRI simulator (SIMRI; [3]). The simulator was set to use TE, TR and flip angle parameters of the 1.5T scanner in the Rotterdam Scan Study [9]. MRBrainS13 is a grand challenge for medical image tissue segmentation methods containing 5 scans for training [14]. The scans are 3T and have been fully manually annotated. IBSR is a classical data set of 18 patients and is automatically segmented but manually corrected [18]. Skulls are stripped off in all scans. Figure 1 visualizes examples from these sets.

4.2 Segmentation methods

We will compare the following methods: firstly, a U-net consisting of a mirrored VGG16 architecture pre-trained on ImageNet and fine-tuned on labeled data from the source medical center [19, 20]. This method represents the performance of a state-of-the-art tissue segmentation model without taking center-based variation into account. Secondly, we take both an unsupervised and a semi-supervised variational Gaussian mixture model (UGM, SGM), initialized using $k$-means and $1$-nearest-neighbours respectively. Thirdly, an unsupervised and a semi-supervised hidden Potts Gaussian mixture (UHP, SHP) are taken, also initialized using $k$-means and $1$-nearest-neighbours. Comparing these with the previous two models shows the influence of smoothing the segmentations. Lastly, we train a $1$-nearest-neighbours (1NN) based on the labeled voxels (1 per tissue) in the target image, as a baseline supervised tissue classifier. The maximum number of training iterations is set to 30 for all methods.

4.3 Results

We present mean classification errors (with standard errors of the means over 10 repetitions) of each method in Table 1. Firstly, comparing the performances of UHP and SHP in the experiments off the diagonal (multi-center) with their performances on the diagonal (single-center) shows that the learned smoothness parameters are more effective than the chosen ones. Secondly, the errors of the hidden Potts models versus the standard Gaussian mixtures tend to be lower or similar (UHP $<=$ UGM and SHP $<=$ SGM). Thirdly, the semi-supervised models tend to outperform the unsupervised ones (SGM $<$ UGM and SHP $<$ UHP). Taken the performance of 1NN into account, it shows that even 1 label per tissue is very informative. U-net performs poorly as it is not aware of the intensity and contrast shifts between data sets.

Figure 2 shows examples of each segmentation method on the MRBrainS13 data set, with Brainweb1.5T as the source center. For the unsupervised models we only show boundaries between clusters, to indicate that interpretation remains a necessary step. A couple of observations can be made: firstly, the hidden Potts models produce smoother segmentations. Secondly, the U-net over-predicts white matter in the whole image. Thirdly, the 1-nearest-neighbours classifier over-predicts background voxels in fluid regions. Lastly, all methods favour white matter over gray matter in ambiguous regions.

5 Discussion

Although segmentations remain relatively consistent across medical centers compared to the scans, annotator variation can be quite large. This is especially true if medical centers teach different annotation protocols. To account for this type of variation, it would be more appropriate to capture the uncertainty in smoothness and infer the posterior over $\beta$ [13].

In our formulation, the hidden Potts-MRF acts as a spatial regularizer on the responsibilities estimated by the variational mixture model. Spatial regularizers are not uncommon, but are often employed on the observed data: most models incorporate information on the smoothness in $X$ to estimate $Y$. Here, we explicitly look at smoothness in $Y$.

A limitation of the current model is that it is not appropriate for abnormality or pathology detection. That would require a different number of components for images with pathologies versus images without pathologies. However, it should be possible to extend variational Gaussian mixture models to incorporate a variable amount of components. In that case, the component weights are not modeled using a Dirichlet distribution, but a Dirichlet process [6].

6 Conclusion

We proposed to tackle center-specific variation in medical imaging data sets with Bayesian transfer learning. We fitted a spatial smoothness prior on the segmentations produced in one medical center and used this informative prior to perform brain tissue segmentation at the target center. Our results show improvements over non-spatially smoothed segmentations, and improvements with learned smoothness parameters over chosen ones.