Empirical Bayesian Mixture Models for Medical Image Translation

TL;DR

This paper introduces an interpretable Bayesian mixture model for medical image translation that can predict missing imaging modalities from limited data, demonstrating effectiveness across multiple clinical scenarios.

Contribution

It presents a novel probabilistic generative model capable of handling missing data and training on small datasets for medical image translation tasks.

Findings

Effective prediction of missing MR contrasts and CT images.

Model performs well with limited training data.

Validated on three clinically relevant scenarios.

Abstract

Automatically generating one medical imaging modality from another is known as medical image translation, and has numerous interesting applications. This paper presents an interpretable generative modelling approach to medical image translation. By allowing a common model for group-wise normalisation and segmentation of brain scans to handle missing data, the model allows for predicting entirely missing modalities from one, or a few, MR contrasts. Furthermore, the model can be trained on a fairly small number of subjects. The proposed model is validated on three clinically relevant scenarios. Results appear promising and show that a principled, probabilistic model of the relationship between multi-channel signal intensities can be used to infer missing modalities -- both MR contrasts and CT images.

Equations

$$p\left(\vec{X}\mid \vec{Z},{\vec{\mu }}_{1\dots K},{\vec{\Sigma }}_{1\dots K}\right)=\prod _{d=1}^D\prod _{k=1}^K\mathcal{N}{\left({\vec{x}}_d\mid {\vec{\mu }}_k,{\vec{\Sigma }}_k\right)}^{z_{dk}}.$$

$$p\left(\vec{Z}\right)=\prod _{d=1}^D\mathrm{Cat}\left({\vec{z}}_d\mid {\vec{\pi }}_d\right),{\vec{\pi }}_d\in {\mathbb{R}}^K,{\pi }_{dk}=\frac{\exp ({\omega }_k+{\vec{a}}_{dk}(\vec{\varphi }))}{{\sum }_{j=1}^K\exp ({\omega }_j+{\vec{a}}_{dj}(\vec{\varphi }))},$$

$$p\left({\vec{\mu }}_k,{\vec{\Sigma }}_k^{-1}\right)=$$

$$=$$

$$q\left(\vec{Z},{\vec{\mu }}_{1\dots K},{\vec{\Sigma }}_{1\dots K}^{-1}\right)=\left[\prod _{d=1}^{D}q\left({\vec{z}}_d\right)\right]\left[\prod _{k=1}^{K}q\left({\vec{\mu }}_k,{\vec{\Sigma }}_k^{-1}\right)\right],$$

$$\mathcal{L}=$$

$$-\sum _{d=1}^D{D}_{\mathrm{KL}}\left(q_{{\vec{z}}_d}\Vert p_{{\vec{z}}_d}\right)-\sum _{k=1}^K{D}_{\mathrm{KL}}\left(q_{{\vec{\mu }}_k,{\vec{\Sigma }}_k}\Vert p_{{\vec{\mu }}_k,{\vec{\Sigma }}_k}\right).$$

$$p\left(\vec{g}\mid {\vec{\mu }}_k,{\vec{\Sigma }}_k,z_k=1\right)=\mathcal{N}\left(\vec{g}\mid {\vec{\mu }}_{k\vec{o}},{\vec{\Sigma }}_{k\vec{o}\vec{o}}\right),$$

$$p\left(\vec{h}\mid \vec{g},{\vec{\mu }}_k,{\vec{\Sigma }}_k,z_k=1\right)=$$

$$\mathcal{N}\left(\vec{h}\mid {\vec{\mu }}_{k\vec{m}}-{\left({\vec{\Lambda }}_{k\vec{m}\vec{m}}\right)}^{-1}{\vec{\Lambda }}_{k\vec{m}\vec{o}}\left(\vec{g}-{\vec{\mu }}_{k\vec{o}}\right),{\left({\vec{\Lambda }}_{k\vec{m}\vec{m}}\right)}^{-1}\right),$$

$$q\left(\mathcal{H},\vec{Z},{\vec{\mu }}_{1\dots K},{\vec{\Sigma }}_{1\dots K}^{-1}\right)=\left[\prod _{d=1}^{D}q\left({\vec{h}}_d\mid {\vec{z}}_d\right)q\left({\vec{z}}_d\right)\right]\left[\prod _{k=1}^{K}q\left({\vec{\mu }}_k,{\vec{\Sigma }}_k^{-1}\right)\right],$$

$$q({\vec{h}}_d)=\sum _{k=1}^K{\tilde{z}}_{dk}\mathcal{N}\left({\vec{h}}_d\mid {\widetilde{\stackrel{⃗}{h}}}_{dk},{\widetilde{\stackrel{⃗}{S}}}_{dk}\right).$$

$$\mathcal{L}=$$

$$-\sum _{d=1}^D{D}_{\mathrm{KL}}\left(q_{{\vec{z}}_d}\Vert p_{{\vec{z}}_d}\right)-\sum _{k=1}^K{D}_{\mathrm{KL}}\left(q_{{\vec{\mu }}_k,{\vec{\Sigma }}_k}\Vert p_{{\vec{\mu }}_k,{\vec{\Sigma }}_k}\right)$$

$$\mathcal{L}=$$

$$-\sum _{d=1}^D{D}_{\mathrm{KL}}\left(q_{{\vec{z}}_d}\Vert p_{{\vec{z}}_d}\right)-\sum _{k=1}^K{D}_{\mathrm{KL}}\left(q_{{\vec{\mu }}_k,{\vec{\Sigma }}_k}\Vert p_{{\vec{\mu }}_k,{\vec{\Sigma }}_k}\right).$$

$${\tilde{z}}_{dk}=$$

$${\tilde{b}}_k=$$

$${\widetilde{\stackrel{⃗}{\mu }}}_k=$$

$${\tilde{\nu }}_k=$$

$${\widetilde{\stackrel{⃗}{V}}}_k^{-1}=$$

$$\begin{array}{rl}\mathbb{E}{\left[z_{dk}{\vec{x}}_d\right]}_{\vec{o}}=& {\tilde{z}}_{dk}{\vec{g}}_d,\\ \mathbb{E}{\left[z_{dk}{\vec{x}}_d\right]}_{\vec{m}}=& {\tilde{z}}_{dk}{\widetilde{\stackrel{⃗}{h}}}_{dk},\\ \mathbb{E}{\left[z_{dk}{\vec{x}}_d{\vec{x}}_d^{\mathrm{T}}\right]}_{\vec{o}\vec{o}}=& {\tilde{z}}_{dk}{\vec{g}}_d{\vec{g}}_d^{\mathrm{T}},\end{array}\begin{array}{rl}\mathbb{E}{\left[z_{dk}{\vec{x}}_d{\vec{x}}_d^{\mathrm{T}}\right]}_{\vec{m}\vec{m}}=& {\tilde{z}}_{dk}\left({\widetilde{\stackrel{⃗}{h}}}_{dk}{\widetilde{\stackrel{⃗}{h}}}_{dk}^{\mathrm{T}}+{\widetilde{\stackrel{⃗}{S}}}_{dk}\right),\\ \mathbb{E}{\left[z_{dk}{\vec{x}}_d{\vec{x}}_d^{\mathrm{T}}\right]}_{\vec{o}\vec{m}}=& {\tilde{z}}_{dk}{\vec{g}}_d{\widetilde{\stackrel{⃗}{h}}}_{dk}^{\mathrm{T}},\\ \mathbb{E}{\left[z_{dk}{\vec{x}}_d{\vec{x}}_d^{\mathrm{T}}\right]}_{\vec{m}\vec{o}}=& {\tilde{z}}_{dk}{\widetilde{\stackrel{⃗}{h}}}_{dk}{\vec{g}}_d^{\mathrm{T}},\end{array}$$

$$\begin{array}{rl}{\widetilde{\stackrel{⃗}{h}}}_{dk}=& {\widetilde{\stackrel{⃗}{\mu }}}_{k\vec{m}}-{\widetilde{\stackrel{⃗}{\Lambda }}}_{k\vec{m}\vec{m}}^{-1}{\widetilde{\stackrel{⃗}{\Lambda }}}_{k\vec{m}\vec{o}}\left({\vec{g}}_d-{\widetilde{\stackrel{⃗}{\mu }}}_{k\vec{o}}\right),\end{array}\begin{array}{rl}{\widetilde{\stackrel{⃗}{S}}}_{dk}=& {\widetilde{\stackrel{⃗}{\Lambda }}}_{k\vec{m}\vec{m}}^{-1},\end{array}$$

$${\vec{\mu }}_{0k}=$$

$$b_{0k}^{-1}=$$

$${\vec{V}}_{0k}=$$

$$\frac{\partial \mathcal{L}}{\partial {\nu }_{0k}}=$$

$$\frac{{\partial }^2\mathcal{L}}{\partial {\nu }_{0k}^2}=$$

$$\text{PSNR}=10{\log }_{10}\frac{maxva{l}^2}{MSE},$$

Full text

1 Introduction

This paper concerns a relatively simple method of synthesising data of one medical image modality, from data of other modalities. This is known as ‘image translation’. Applications of medical image translation are numerous, and include e.g. harmonising data across scanners; synthesising computed tomography (CT) images from magnetic resonance (MR) images for positron emission tomography (PET) attenuation correction [1], or decrease the need for radiating a patient; simplifying the problem of multi-modal image registration [2]; or generalising machine learning techniques by transferring out-of-distribution input data to the domain of the model’s training data [3].

Mapping from the signal intensities of one modality to those of another can be loosely categorised as either optimisation- or learning-based. Optimisation-based methods rely only on the data at hand to optimise a mapping between modalities, and do not use training data. Examples include using non-parametric joint histograms [4], estimating an intensity transformation during image registration [5], and biophysical models [6]. Learning-based methods use training data to learn the mapping, and can be applied to translating an unseen image from one domain into another. Some examples in this category use clustering [7], random forests [8], patch-matching [9] and dictionaries [10]. Learning-based methods based on various convolutional neural network architectures are currently the most popular approach for this. Trained end-to-end, on either paired or unpaired training data [11, 12, 13], they show promising results at this task, although they can run the risk of hallucinating unwanted features [14].

This paper presents a more interpretable generative modelling approach to image translation. It could be classed as an optimisation-based approach, although it does use training data to learn priors that inform the optimisation of mappings. More specifically, we show how a generative model for group-wise normalisation and segmentation of neuroimaging data can be extended to handle missing data. The generative model has a Gaussian mixture model component, which can naturally handle missing data [15]. In this paper, we extend this missing data model to a variational Gaussian mixture. Fitting this model to various populations of medical images allows us to predict, from a few MR contrasts, entirely missing modalities (e.g., non-acquired MR contrasts or CT images).

2 Methods

The prediction of one modality from another is here cast as a joint intensity modelling problem. The workhorse of the proposed method is the unified segmentation model [16], which uses mixtures of Gaussians with non-stationary tissue priors derived from a deformable template. When a large dataset is available, the optimisation of the template can be interleaved with the mixture model fit to each individual subject [17]. Furthermore, priors over the intensity parameters of the Gaussian mixture – its means and covariances – can be defined and optimised as well. This type of learning, where subject-specific parameters are marginalised while population parameters are optimised, is known as parametric empirical Bayesian methods [18]. Here, exact marginalisation is intractable, so we resort to a variational approximation.

2.0.1 Fully observed model:

Let $\vec{X}\in\mathbb{R}^{D\times M}$ be a multimodal dataset from one subject, where $M$ is the number of modalities and $D$ is the number of voxels in the images. Each voxel is assumed to belong to one of $K$ classes, where the classification is encoded by the label matrix $\vec{Z}\in\left[0,1\right]^{D\times K}$, with $z_{dk}=1$ iff. voxel $d$ belongs to class $k$. Each tissue class is associated with a multivariate Gaussian distribution of dimension $M$, which encodes the intensities’ mean ($\bm{\mu}_{k}\in\mathbb{R}^{M}$) and covariance ($\bm{\Sigma}_{k}\in\mathbb{R}^{M\times M}$) over the modalities. The Gaussian mixture model can then be written as a conditional probability that factorises across voxels:

[TABLE]

Subject-specific parameters (the label matrix and Gaussian parameters) are assumed to be drawn from prior distributions that describe their variability at the population level. Labels are drawn from a categorical distribution whose probabilities are encoded by a deformable template $\vec{a}\in\mathbb{R}^{D_{a}\times K}$. This template is mapped to the subject’s brain using a non-linear deformation field $\vec{\phi}$. This assumption can be written as the conditional likelihood:

[TABLE]

where $\vec{\omega}\in\mathbb{R}^{K}$ is a vector of global class proportions, which can be optimised to account for variable amounts of different classes (an example when modelling brain images could be atrophy due to ageing). The Gaussian parameters are drawn from their conjugate Gauss-Wishart distribution:

[TABLE]

Assuming that all population parameters are fixed, a mean-field approximation is made so that the posterior distribution over all latent, subject-specific parameters factorises as:

[TABLE]

with $q\left(\vec{z}_{d}\right)=\operatorname{Cat}\left(\vec{z}_{d}\mid\tilde{\vec{z}}_{d}\right)$ and $q\left(\vec{\mu}_{k},\vec{\Sigma}^{-1}_{k}\right)=\mathcal{NW}\left(\vec{\mu}_{k},\vec{\Sigma}^{-1}_{k}\mid\tilde{\vec{\mu}}_{k},\tilde{b}_{k},\tilde{\vec{V}}_{k},\tilde{\nu}_{k}\right)$. The posterior parameters (denoted by a tilde) can be optimised in turn by maximising the evidence lower bound (ELBO):

[TABLE]

When multiple subjects $\left\{\vec{X}_{n}\right\}_{n=1}^{N}$ are processed, the posterior distribution factorises across subjects and a combined ELBO can be written by summing the individual ELBOs ($\mathcal{L}=\sum_{n=1}^{N}\mathcal{L}_{n}$). In this case, empirical population priors can be obtained by optimising the combined ELBO with respect to the template ($\vec{a}$) and Gauss-Wishart prior hyper-parameters ($\vec{\mu}_{0k},b_{0k},\vec{V}_{0k},\nu_{0k}$). The means and scale matrices have closed form solutions, while the template and degrees of freedom must be optimised using an iterative scheme. Population prior parameters and subject posterior parameters can be optimised in turn, resulting in a variational Expectation-Maximisation (VEM) algorithm [19].

2.0.2 Missing modalities:

Let us assume that some modalities are missing in a voxel111For example, a multi-channel MRI might have three contrasts: T1w, T2w and PDw. In one voxel, only the T1w intensity is observed. The T2w and PDw intensities are then assumed missing in that voxel. Note that different voxels can have different combinations of contrasts/modalities missing.. We write as $\vec{o}$ the vector indexing observed modalities and as $\vec{m}$ the vector indexing missing modalities. Therefore, the observed channels can be written as $\vec{g}=\vec{x}_{\vec{o}}$ and the missing channels as $\vec{h}=\vec{x}_{\vec{m}}$, where the voxel index $d$ has been temporarily dropped for clarity. For a voxel in class $k$, the marginal distribution of the observed channels can then be written as [20]:

[TABLE]

and the conditional distribution of the missing channels as:

[TABLE]

where the precision matrix $\vec{\Lambda}=\vec{\Sigma}^{-1}$ is the inverse of the covariance matrix.

The set of all missing values in an image is written as $\mathcal{H}=\left\{\vec{h}_{d}\right\}_{d=1}^{D}$. The mean field approximation becomes:

[TABLE]

where $q\left(\vec{h}_{d}\mid\vec{z}_{d}\right)=\prod_{k=1}^{K}\mathcal{N}\left(\vec{h}_{d}\mid\tilde{\vec{h}}_{dk},\tilde{\vec{S}}_{dk}\right)^{z_{dk}}$. The marginal posterior over missing values is a mixture of Gaussians that can be obtained by marginalising the labels:

[TABLE]

Its expected value is $\mathbb{E}\left[\vec{h}_{d}\right]=\sum_{k}\tilde{z}_{dk}\tilde{\vec{h}}_{dk}$. This is the expression that we evaluate to predict missing voxels.

The set of all observed values is written as $\mathcal{G}=\left\{\vec{g}_{d}\right\}_{d=1}^{D}$. The ELBO can then be written in two equivalent forms:

[TABLE]

The first form is used to optimise the labels’ posterior parameters, while the second is used to optimise, in turn, the missing values and the Gaussian posterior parameters.

2.0.3 Model updates:

Optimising the ELBOs in (14) and (15) gives the subject-level update equations as:

[TABLE]

The update equations for the Gaussian parameters in the missing data case are very similar to the fully observed case, except that expectations are taken about the data. These expectations are evaluated as:

[TABLE]

where

[TABLE]

and $\tilde{\vec{\Lambda}}_{k}=\tilde{\nu}_{k}\tilde{\vec{V}}_{k}$ is the posterior expected precision matrix of a given class.

Finally, we provide the optimal updates of the Gaussian prior parameters, given a set of individual posterior parameters. All prior parameters have closed-form updates, except for the degrees of freedom of the Wishart distribution, which is updated using an iterative Gauss-Newton scheme. The update equations are:

[TABLE]

We do not provide update rules for the template ($\vec{a}$), as they can be found in [17].

3 Experiments and Results

In this section we aim to explore the translation (or inference) capability of the proposed model by conducting three experiments on publicly available data. We investigate: (1) inferring missing voxels of MRIs with differing field of views; (2) inferring entirely missing MRI contrasts; and (3), inferring CT scans from MRIs. The findings are quantified by computing the peak-signal-to-noise-ratio (PSNR) for an image channel $c$ as:

[TABLE]

where the mean-squared error is defined as $\text{MSE}=\frac{1}{D}\sum_{d=1}^{D}(\hat{x}_{cd}-(\mathbb{E}\left[\vec{h}_{d}\right])_{c})^{2}$, $maxval$ is the maximum channel intensity in the reference image $\hat{\vec{X}}$, and $\mathbb{E}\left[\vec{h}_{d}\right]$ from (12) is evaluated to predict missing voxels. The PSNR is a metric that is commonly used in the medical image synthesis literature [11, 13, 12]. Note that no voxels are excluded when computing the PSNR.

3.1 MRI Contrast Translation

This section evaluates translating between MR contrasts. The model is trained on 50 subjects from the publicly available IXI dataset222http://brain-development.org/ixi-dataset/, which was acquired on three different MR scanners333This scenario is more realistic in a clinical context. The results would improve if data from only one scanners was used.. Each IXI subject has three MR images: a T1-, T2- and PD-weighted scan (T1w, T2w and PDw). Furthermore, the images have approximately 1 mm isotropic voxels and all subjects are healthy. $K=12$ mixture components are used, resulting in the model shown in Fig. 1. Note that the template learned by the algorithm does not need to represent real tissues. Here, the model has been treated as a method of representing a probability density function, rather than as a way to do clustering. Any ‘meaningful’ clusters are incidental.

3.1.1 Inferring MRIs with Differing Fields of View:

Doctors often acquire routine clinical MR scans of multiple contrasts. Commonly, these contrasts have differing fields of view, meaning the brain coverage varies (cf. observed T1w and T2w images in Fig. 2). This can be problematic for image segmentation routines as voxels with non-observed contrasts need to be discarded. The model should prevent this issue by inferring the values of these missing voxels. To test this, T1w, T2w and PDw scans of 50 unseen IXI subjects are used444The model is trained on IXI subjects IXI[064-118], and tested on IXI[002-063].. All of the voxels are retained in the PDw image, while an increasing amount of voxels are removed from the T1w and T2w images (25%, 50%, 75% and 100%). The missing voxels are then inferred with the trained model. An example can be seen in Fig. 2. The mean PSNR computed between the known references and the inferred images are shown in Table 2. For routine clinical MRI, it is rare that more than 50% of the field of view is missing. The results therefore suggest that the model does a good job at filling in missing fields of view, which could be of value in segmenting hospital data.