InVA: Integrative Variational Autoencoder
for Harmonization of Multi-modal Neuroimaging Data

Image-on-image regression involves predicting one imaging modality based on other imaging modalities. This approach is commonly used when the imaging modality to be predicted is either too costly to acquire or when a clear version of the image is not available (Jeong et al., 2021; Subramanian et al., 2023; Onishi et al., 2023). In the domain of image-on-image regression, the prevailing method involves conducting region-by-region regression analyses between images. For example, in a study related to Multiple Sclerosis, (Sweeney et al., 2013) applied region-wise logistic regression models, incorporating T1-weighted, T2-weighted, FLAIR, and PD volumes to predict lesion incidence. However, a notable limitation of these region-wise approaches is their inability to capture associations between different regions, resulting in reduced accuracy of predicting the output image from input images. To overcome this limitation, some methods employ adaptive smoothing techniques to integrate information from neighboring regions (Hazra et al., 2019). A more general and effective approach involves smoothing coefficients connecting outcome and input images using spatially varying coefficient models, well-suited for exploring regression relationships between spatially structured images (Niyogi et al., 2023; Mu et al., 2018; Zhu et al., 2014; Mu, 2019). Extending this direction of research, spatial latent factor models have been introduced to capture more complex non-linear spatial dependencies between outcome and input images (Guo et al., 2022).

The utilization of spatially varying coefficients in image-on-image regression is effective for leveraging information between regions. However, these methods tend to be computationally expensive, even when either the sample size or number of regions is moderately large. Another research direction treats both response and input images as tensors, leading to the emergence of tensor-on-tensor regression approaches (Lock, 2018; Gahrooei et al., 2021; Miranda et al., 2018; Guhaniyogi & Rodriguez, 2020; Guha & Guhaniyogi, 2021; Guhaniyogi & Spencer, 2021). While these methods implicitly consider smoothing among neighboring regions, they often necessitate scaling down images due to significant computational demands and low signal-to-noise ratios. Furthermore, these approaches have not yet addressed the modeling of non-linear associations between images. A third avenue of research focuses on developing multivariate support vector machines for predicting missing spatial information in EEG from fMRI (De Martino et al., 2011) or missing temporal information in fMRI from EEG data (Jansen et al., 2012).

Compared with the original high-dimensional images, the low-dimensional features of images can facilitate estimating relationships between an outcome image and input images. With this motivation, we are particularly intrigued by deep neural networks (DNNs) for non-linear dimension reduction. Generative algorithms, such as Variational Auto-Encoders (VAEs) (Kingma & Welling, 2013; Doersch, 2016; Girin et al., 2020; Zhao & Linderman, 2023), have proven successful in representing images via low-dimensional latent variables. VAEs model the population distribution of image data through a simple distribution, often Gaussian distribution, for the latent variables combined with a complex non-linear mapping function. A key to the success of such methods is the use of flexible probability fields to represent the important information of images, as well as rapid computation with high-dimensional images and large sample. Despite the success of VAEs in imaging analysis, they do not fully explore shared information between input images to enhance performance in predicting the outcome image. More precisely, the existing approaches on VAEs synthesize multi-modal imaging data at two levels, input-level and decision-level. Input-level fusion usually involves merging multiple inputs together before modeling, which can result in a large feature vector. Deciding how to merge the images is not an easy task, and naive merging may lead to poor performance. In contrast, in decision-level fusion, each type of image is used to train a model, and the output of the model is fused to make a final decision. This type of fusion separates each type of image information during training and ignores the contribution of complementary information between different types of images to the training. Ignoring interconnection between input images not only limits the biological plausibility and interpretation of predictive inference for the outcome image, but broadly has been shown to reduce statistical efficiency (Dai & Li, 2021), and increase sensitivity to noise in the images (Calhoun & Sui, 2016).

Hierarchical Bayesian methods allow structured information to be borrowed explicitly among image inputs via joint prior structure on coefficients at different layers of hierarchy and offer inference via the joint posterior distribution (Jin et al., 2020; Lee et al., 2020; Lei et al., 2021; Su et al., 2022; Kaplan et al., 2023). However, this perspective has been under-utilized due to computational bottlenecks and lack of appropriate modeling architecture.

Motivated by the hierarchical Bayesian principle of leveraging shared information, this paper introduces an integrative variational autoencoder (InVA) designed for the harmonization of multi-modal neuroimaging data in a computationally efficient manner. Our primary contributions are outlined as follows. First, we construct DNN-based encoder $e_{k}$ and decoder $d_{k}$ corresponding to the $k$th input image, $k=1,...,K$. These encoders and decoders are utilized to map the $k$th input image to low-dimensional multivariate Gaussian distributions, independently for each $k=1,...,K$, all of the same dimension. The construction preserves sufficient flexibility of modeling each input image while providing a shallow-level representation of the images in the context of hierarchical Bayesian modeling.

Second, we sample $h_{k}$ from each learned multivariate Gaussian distribution, representing the features of the $k$th input image. Subsequently, we construct a DNN-based encoder $\tilde{e}$ and decoder $\tilde{d}$ shared over input images to map these features $h_{k}$ to another multivariate Gaussian distribution $N(m_{k},b_{k})$. The shared DNN architecture $\tilde{e}$ and $\tilde{d}$ across input images, combined with the independent distributions $N(m_{k},b_{k})$ for each $k$, strike a favorable balance between individual and shared information. This configuration represents a deeper-level of hierarchical Bayesian modeling.

Third, the parameters from the encoders and decoders specific to each input image, along with those from the shared encoder and decoder, are collectively fed into a Deep Neural Network (DNN)-based predictor designed for prediction of the output image. The joint learning of input parameters from all imaging inputs facilitates the sharing of information across different inputs, leading to more accurate predictions of the output image. The performance of InVA surpasses that of ordinary Variational Autoencoders (VAEs) constructed independently using different input images, as well as other popular image-on-image regression competitors, in predicting the output image across various simulation studies. The exceptional performance of InVA in predicting a costly image from inexpensive imaging inputs in multi-modal neuroimaging data underscores the importance of borrowing information from input images.

Our InVA approach has incorporated novel modeling architecture over the existing literature on hierarchical VAEs. In the hierarchical VAE literature, DRAW (Gregor et al., 2015) introduces a deep and recurrent approach that combines a sequence of VAEs to obtain more realistic image generation. Ladder Variational Autoencoders (Sønderby et al., 2016) also designs a new version of layered VAE that recursively corrects the generative distribution and produces highly expressive models. This is further generalized to other hierarchical variational models to get expressive variational distribution as well as efficient computation (Ranganath et al., 2016). Hierarchical priors (Klushyn et al., 2019) are then proposed in VAE to avoid the overregularization that results from the standard normal prior distribution and to encourage the properties desirable for model learning, such as smoothness and simple explanatory factors. More recently, NVAE (Vahdat & Kautz, 2020) has also designed a hierarchical VAE, which utilizes a deep hierarchical structure to achieve more stable and accurate image generation.

Despite their good performance, these hierarchical VAEs are designed in pursuit of better performance when modeling a single source imaging data and are not designed to adequately extract shared information from multiple imaging inputs in the prediction of an output image. In contrast to these hierarchical VAEs, our InVA has a more flexible structure to handle both single-source and multi-source imaging data. On the one hand, InVA has a hierarchical structure to capture complex patterns. On the other hand, InVA introduces both input-specific and shared model components. This new architecture allows better integration and information borrowing when facing multiple imaging inputs, leading to more accurate output image prediction. At the same time, even when dealing with a single imaging input, this new architecture enables better generalization of hierarchical VAE literature by allowing the same images from multiple views (Yu et al., 2023; Yan et al., 2021). Therefore, our InVA fills an important gap in bridging the literature between hierarchical VAEs and image-on-image regression, broadening applications of VAEs in the analysis of multi-modal neuroimaging data.

Autoencoder (AE) is a widely-used unsupervised learning method that utilizes an encoder to compress data and reconstruct the data from the encoded features through a decoder (Geng et al., 2015; Tschannen et al., 2018; Chorowski et al., 2019; Nazari et al., 2023; Hao & Shafto, 2023). To cope with different scenarios, variants of autoencoders have also been inspired (Ng & Autoencoder, 2011; Rifai et al., 2011a, b; Chen et al., 2012, 2014; Ranjan et al., 2017; Kingma & Welling, 2013; Tolstikhin et al., 2017; Pei et al., 2018; Vahdat & Kautz, 2020). Based on AE, variational autoencoders (VAEs) are designed to model the data distribution (Doersch, 2016; Girin et al., 2020; Zhao & Linderman, 2023), which maps the input data into latent Gaussian distribution through the encoder (Kviman et al., 2023; Hao & Shafto, 2023; Janjos et al., 2023).

In VAE, the learning of encoder and decoder weights is usually based on variational inference (Blei et al., 2017; Nakamura et al., 2023), where the loss is defined as the negative variational lower bound on the marginal likelihood (Kingma & Welling, 2013). To be more specific, let the input image $X_{k,i}$ be mapped to the latent variables $z_{k,i}\in\mathbb{R}^{p}$ which follows a prior distribution $p(z_{k,i})=N(0,I_{p})$. Assume that $q(z_{k,i}|X_{k,i})$ denotes the variational distribution for $z_{k,i}$, $p_{\phi}(X_{k,i}|z_{k,i})$ denotes the likelihood of $X_{k,i}$, and KL stands for the Kullback–Leibler divergence. Assuming $q(z_{k,i}|X_{k,i})=\prod_{j=1}^{p}N(\mu_{i,j},\sigma_{i,j})$, the training objective is to minimize the negative of the evidence lower bound (ELBO), given by,

where $\hat{X}_{k}$ is the reconstruction of $X_{k}$ through the decoder. Importantly, the standard VAE architecture does not allow borrowing of information between multiple imaging inputs.

We adopt an architecture inspired by hierarchical Bayesian modeling to enhance the learning of the latent variable distribution from multiple imaging inputs. Our integrative variational autoencoder incorporates image-specific encoders and decoders for each input image at a shallow level to capture image-specific features. Additionally, it includes encoders and decoders shared by all input images at a deeper level to facilitate information borrowing and capture shared features of the input images.

Image-specific encoder: For every input image $X_{k}$, we employ a DNN-based image-specific encoder $e_{k}$ to project it onto a hidden $p$-variate Gaussian distribution $N(\mu_{k},\sigma_{k})$. Subsequently, we sample $h_{k}\in\mathbb{R}^{p}$ from this distribution to depict the shallow-level imaging features. This process effectively maps various input images to a common latent feature space, facilitating finer feature extraction.

Shared encoder: To leverage on the shared information provided by the input images and enhance feature extraction, we construct a shared DNN-based encoder $\tilde{e}$. This encoder maps the hidden features $h_{k}$ from each image to a deeper hidden $p$-variate Gaussian distribution $N(m_{k},b_{k})$. Subsequently, we sample $z_{k}\in\mathbb{R}^{p}$ from this distribution to represent deep-level features corresponding to each input image. These features are analogous to hyperparameters in Bayesian hierarchical models, facilitating information borrowing.

Shared decoder: Next, we incorporate a shared DNN-based decoder $\tilde{d}$ to predict image-specific features $h_{k}$ from finer features $z_{k}$. Through minimizing the difference between $h_{k}$ and the fitted $\hat{h}_{k}$, we enable information borrowing and optimize the weights of the shared encoder $\tilde{e}$ and decoder $\tilde{d}$ to achieve a well-fitted deep hidden distribution $N(m_{k},b_{k})$.

Image-specific decoder: After the shared decoder $\tilde{d}$, we also use a image-specific decoder $d_{k}$ for each imaging input, which maps the predicted $\hat{h}_{k}$ to the input image space and predicts $\hat{X}_{k}$. By minimizing the difference between the input image $X_{k}$ and the predicted $\hat{X}_{k}$, the modality-specific information can be used to learn the weights in $e_{k}$ and $d_{k}$ and thus better estimation of the shallow-level hidden distribution.

To further predict the output image $Y$ based on the extracted hidden data distributions, we utilize a multi-level conditional structure in our InVA. Specifically, to predict the shared response $Y$ with both modality-specific information and shared knowledge, we concatenate the extracted distribution parameters at different levels, i.e., $\{\mu_{k},\sigma_{k},m_{k},b_{k}\}$ where $k\in\{1,\cdots,K\}$. The obtained vector is then sent to a DNN-based predictor, which maps it to the response space and predicts $\hat{Y}$. In order to learn the data distribution and predict the response $Y$ at the same time, we add the difference between the response $Y$ and the prediction $\hat{Y}$ to the loss function and then minimize the loss. Without loss of generality, we take $K=2$ as an example and illustrate the conditional structure in Figure 2.

The learning of the weights of our InVA is based on variational inference. The marginal likelihood of imaging inputs consists of the sum of the marginal likelihoods for individual data from different data sources and can be written as Eq. (3), where $H_{(i)}=\{h_{1,i},\cdots,h_{K,i}\}$ and $Z_{(i)}=\{z_{1,i},\cdots,z_{K,i}\}$ are hidden features of $\{X_{1,i},...,X_{K,i}\}$ at shallow and deeper levels, respectively, for $i=1,...,n$.

The first term in Eq. (3) is KL divergence between the posterior distribution of $Z_{(i)}$ and its variational approximation which is non-negative. Thus, to maximize the first term in Eq. (3), we minimize the second term $\mathcal{L}(\theta,\phi,X_{(i)})$ which is called the variational lower bound on the marginal likelihood of i-th subject and can be written as Eq. (4):

where $p_{\theta}(X_{k,i}|z_{k,i})$ is the data likelihood, $q_{\phi}(h_{k,i}|X_{k,i})$ and $q_{\phi}(z_{k,i}|h_{k,i})$ represent the variational distribution of $h_{k,i}$ and $z_{k,i}$, respectively. We assume these variational distributions assume the form of $p$-variate Gaussian distributions, i.e., $q_{\phi}(h_{k,i}|X_{k,i})=\prod_{j=1}^{p}N(\mu_{k,i,j},\sigma_{k,i,j})$ and $q_{\phi}(z_{k,i}|h_{k,i})=\prod_{j=1}^{p}N(m_{k,i,j},b_{k,i,j})$. The priors for $z$ and $h$ are set as standard normal distribution, i.e., $p_{\theta}(h_{k,i})=N(0,I_{p})$ and $p_{\theta}(z_{k,i})=N(0,I_{p})$. During the minimization of Eq. (4), the first term represents the difference between input data $X$ and the reconstructed $\hat{X}$, which can be rewritten as

For the second term, it is a penalty for the data-specific feature extraction and takes the form of a KL divergence between the variational distribution $q_{\phi}(h_{k,i}|X_{k,i})$ and the prior distribution on $h_{k,i}$, which assumes a closed form,

In addition, the third term is for the extracted deep features. This term also assumes a closed form given by,

Data Simulation: For the $i$-th subject, where $i=1,...,n$, we generate two input images, $X_{1,i}$ and $X_{2,i}$, with each being a 3-way tensor having dimensions $d\times d\times d$, comprising $d^{3}$ cells. Each cell entry of $X_{1,i}$ and $X_{2,i}$ is independently and identically simulated from the normal distribution N(0,1). The $j=(j_{1},j_{2},j_{3})$-th cell of the outcome image is generated from a polynomial of order $O$ with varying coefficients as follows: $y_{i}(j)=\sum_{o=1}^{O}\sum_{k=1}^{2}\beta_{o,k}(j)x_{k,i}(j)^{o}+\epsilon_{i}(j)$, where $\epsilon_{i}(j)\stackrel{{\scriptstyle i.i.d.}}{{\sim}}N(0,\sigma^{2})$ and $x_{k,i}(j)$ represents the $j$th cell of the $k$th input image for the $i$th sample. The coefficient $\beta_{o,k}(j)$ is generated from a Gaussian process with an exponential correlation kernel, allowing for a nonlinear relationship between the outcome and input images, as well as imposing correlation between cells of the outcome image. Further, the scale parameter for the exponential correlation kernel is set at $0.25$ for all these Gaussian processes to borrow information across input image coefficients. We explore a wide range of simulation scenarios by varying the order of the polynomial $O=1,2,3$ and the signal-to-noise ratio, represented by varying $\sigma=0.1,0.3,0.5$. We consider two different combinations of $(n,d)=(100,2),(800,3)$. The setting $n=800,d=3$ closely matches the scenario in multi-modal neuroimaging data in Section 4. For the test data, we further simulate the equivalent of 20% of the training data using the same simulation settings.

In the scenario where $n=100$ and $d=2$, the performance of all competitors deteriorates under increased noise variance or when the relationship between the outcome and input images becomes more complex, as evidenced by a higher order of the polynomial governing the outcome image. As indicated in Table 1, InVA consistently achieves a smaller MSE across various orders and noise levels $\sigma$ compared to TensorReg. This can be attributed to its effective capture of the non-linear association between the outcome and input images. Although VAE ($X_{1}$), VAE ($X_{2}$) and BART also capture the non-linear association between the outcome and input images, InVA significantly outperforms all of them. Notably, the smaller MSE of InVA compared to VAE ($X_{1}$) and VAE ($X_{2}$) highlights the advantage of borrowing information from multiple imaging inputs. Var-Coef performs exceptionally well when the order of the polynomial used to simulate the outcome image is $O=1$. In this case, the fitted Var-Coef model is identical to the data-generating model. However, InVA outperforms Var-Coef when the order $O$ is set to 2 or 3. This suggests that our approach is advantageous when the true relationship between images is complex and unknown.

In the case of $n=800$ and $d=3$, InVA continues to outperform the baseline competitors (refer to Table 2). Var-Coef is not included as a baseline due to computational challenges with $n=800$. Similar to Table 1, Table 2 demonstrates a decline in performance with increasing noise variance and the order of the true data-generating polynomial. Importantly, both tables establish significantly superior performance when information is suitably borrowed from the two input images in predicting the output image.

We do ablation studies to demonstrate the importance of each component in our InVA, where we train our InVA without shared components (InVA w/o Shd) and without input image-specific components (InVA w/o IS), respectively. In InVA w/o Shd, we train input image-specific encoders and decoders for each modality and average their predictions to obtain a final prediction without using a shared encoder and decoder. In InVA w/o IS, we combine different modalities and train a shared encoder and decoder without adding input image-specific encoders and decoders.

We continue using mean squared error as the comparison metric and summarize the results in Table 3. As we can see, across different signal-to-noise ratios and polynomial orders, our InVA always has a smaller mean squared error compared to InVA w/o Shd and InVA w/o IS, which indicates that both the input image-specific and shared components in our InVA are crucial for the harmonization of multi-modal neuroimaging data analysis.