Probabilistic medical image imputation via deep adversarial learning

Abstract

The ability to impute missing images from a sequence of medical images plays an important role in enabling the detection, diagnosis and treatment of disease. Motivated by this, in this manuscript we propose a novel, probabilistic deep-learning algorithm for imputing images. Within this approach, given a sequence of contrast enhanced CT images, we train a generative adversarial network (GAN) to learn the underlying probabilistic relation between these images. Thereafter, given all but one member from a sequence, we infer the probability distribution of the missing image using Bayesian inference. We make the inference problem computationally tractable by mapping it to the low-dimensional latent space of the GAN. Thereafter, we use Markov Chain Monte Carlo (MCMC) techniques to learn and sample the inferred distribution. Moreover, we propose a novel style loss unique to contrast-enhanced computed tomography (CECT) imaging to improve the texture of the generated images, and apply these techniques to infer missing CECT images of renal masses collected during an IRB-approved retrospective study. In doing so, we demonstrate how the ability to infer the probability distribution of the missing image, as opposed to a single image recovery, can be used by the end-user to quantify the reliability of the imputed results.

Appendices

Appendix A Convergence of MCMC algorithm

MCMC algorithms are known to suffer from challenges when used in high-dimensional spaces, where the mass of the target density is typically concentrated in narrow regions on a lower dimensional manifold [25]. To fully explore the regions of interest, the requirements on the length of the Markov chains [26] can make the algorithm computationally infeasible. Thus, posing the inference problem on the lower-dimensional latent space can alleviate this issue.

A number of diagnostic tools are available to analyse the convergence of the MCMC algorithm, and thus determine the termination length of the generated Markov chains. We direct interested readers to [27] for a summary of such techniques. In the present work, we use the Gelman-Rubin diagnostic [28] which estimates the convergence by considering multiple Markov chains and evaluating the between-chains and within-chains variances.

We consider M chains of length N, each of which is generated by the MCMC algorithm from different random initial points. Let $\mu _m$ and $\sigma _m^2$ denote the sample mean and variance of the mth chain, and $\mu$ denote the overall mean across all chains, i.e., $\mu = \sum _{m=1}^M \mu _m/M$. Then, we estimate the within-chain variance W, the between-chain variance B and the pooled variance $\widehat{V}$ as

$$\begin{aligned} \begin{aligned} W&= \frac{1}{M} \sum _{m=1}^M \sigma _m^2, \quad B = \frac{N}{M-1}\sum _{m=1}^M(\mu _m - \mu )^2,\\ \widehat{V}&= \frac{N-1}{N} W + \frac{B}{N}. \end{aligned} \end{aligned}$$

(A1)

Finally, we evaluate the potential scale reduction factor $\widehat{R} = \sqrt{\widehat{V}/W}$. Assuming that the initial points of the chains were sampled from an over-dispersed distribution compared to the target distribution, $\widehat{V}$ is expected to overestimate the variance of the target distribution, while W underestimates it. Thus, the closer that value of $\widehat{R}$ is to 1, the more assured we are about the convergence of the chains.

To demonstrate the utility of this tool, we consider the chains generated to impute the missing images at the 4 time-points for Subject 1. At each time-point, we use $M=4$ chains for each of the lengths N = 256 K, 512 K, 1024 K. Since the chains are generated for latent variable $z \in {\mathbb {R}}^{N_z}$, we obtain a vector $\widehat{R} \in {\mathbb {R}}^{N_z}$ for each configuration. To simplify the analysis, we condense this vector to a scalar by considering the dimensional mean ± standard deviation of $\widehat{R}$, which is listed in Table 1. Note that these scalar value moves closer to 1 as N increases, indicating convergence. In practice, a value of $\widehat{R} < 1.2$ is considered as a good termination threshold. To balance the convergence of the chains and the associated computational cost, we use $N=1024$ K for all results presented in this work.

Table 1 $\widehat{R}$ values of Markov chains in z used to impute the images at the four time-points for Subject 1

Appendix B Architecture and hyper-parameters

Table 2 Generator architecture

Table 3 Discriminator architecture

We use the axial slice of the 3000 CECT images for training. The WGAN-GP models, whose architectures are described in Tables 2 and 3, are trained with the gradient penalty parameter set to 10. We use the ADAM optimizer with learning rate of $2 \times 10^{-4}$ and momentum parameters $\beta _1 = 0.0, \beta _2 = 0.9$. We perform 5 gradient updates of critic per gradient update of generator and train both networks in TensorFlow with a batch size of 64. For posterior inference we sample Markov chain using Hamiltonian Monte Carlo (HMC) with No-U Turn Sampler (NUTS) [29] and implement it in TensorFlow Probability [30]. We use initial step size of 1.0 for HMC and adapt it following [31] based on the target acceptance probability. A burn-in period of 50% is used for all HMC simulations. These hyper-parameters are chosen to ensure the convergence of the chains.