Motion-invariant variational autoencoding of brain structural connectomes

To comprehensively understand brain function and organization, enormous efforts have been dedicated to imaging and analyzing the brain structural connectome, defined as a collection of white matter fiber tracts connecting different brain regions. Recent advancements in noninvasive neuroimaging techniques have led to a rapid increase in brain imaging datasets, such as the Human Connectome Project (HCP) (Glasser et al., 2016) and the Adolescent Brain Cognitive Development (ABCD) study (Bjork et al., 2017). These developments have inspired a substantial body of literature (Zhang, Allen, et al., 2019; Zhang et al., 2018) that focuses on analyzing structural connectomes derived from diffusion-weighted magnetic resonance imaging (dMRI) data. However, these studies often encounter challenges due to subject head displacement during image acquisition, resulting in image misalignment and artifacts that hinder unbiased connectome reconstructions and analysis (Andersson et al., 2016). Addressing these biases is crucial, which has motivated us to develop novel statistical methods capable of modeling and analyzing brain structural connectomes in an invariant manner (Achille & Soatto, 2018; Zemel et al., 2013).

Our goal is to learn a low-dimensional representation of brain connectomes invariant to motion artifacts in the data, thereby facilitating prediction and inference on the relationship between brain structure and human traits such as cognitive abilities. A common motion correction technique in neuroimaging involves registering motion-corrupted images to a geometrically correct template (Bhushan et al., 2015; Kochunov et al., 2006). However, in dMRI, motion frequently causes geometric distortion and signal dropout, significantly complicating registration (Ben-Amitay et al., 2012) and impacting subsequent brain network analyses (Le Bihan et al., 2006). Current motion correction methods often utilize the FSL eddy tool (Andersson & Sotiropoulos, 2016) to estimate and correct for eddy current-induced distortions and subject movements. Recent advancements include techniques such as SHORELine (Cieslak et al., 2021) and SHARD (Christiaens et al., 2021) for correcting multishell dMRI sequences. Despite these improvements, studies have shown that even after motion correction in diffusion signals, in-scanner head motion can still have notable impacts on the recovered structural connectivity (Baum et al., 2018). Moreover, not all sources of image misalignment are due to motion (Yendiki et al., 2014). Various imperfections in the imaging system can lead to artifacts causing image misalignment. Given these challenges, our major motivation is to develop a method that addresses motion-induced artifacts and adjusts for residual misalignment during the brain network modeling process, rather than in the data preprocessing stage. Here, motion-induced artifacts refer to the in-scanner head displacement that can be approximately quantified using image registration algorithms (Andersson & Sotiropoulos, 2016). Residual misalignment encompasses some remaining artifacts measured after eddy motion correction. By incorporating these considerations into our modeling approach, we aim to improve the robustness and reliability of brain connectome analysis in the presence of motion-induced artifacts.

In biomedical imaging, removing unwanted effects in high-dimensional data has been an important problem for at least the last decade. Without proper adjustment, the association between the signal of interest and one or more nuisance variables can degrade the ability to infer the signal accurately. This issue is particularly evident in brain imaging studies, and various tools have been developed. A popular matrix factorization method (Alter et al., 2000) uses a singular value decomposition of the original matrix to filter out the singular vector associated with the nuisance variable, and reconstructs an adjusted matrix with the remaining singular vectors. Similarly, the Sparse Orthogonal to Group (SOG) (Aliverti et al., 2021) proposes an adjusted dataset based on a constrained form of matrix decomposition. In addition, model-based methods are popular alternatives for removing unwanted associations in the data. For example, distance-weighted discrimination (Benito et al., 2004) adapts Support Vector Machines (SVM) for batch effect removal; the ComBat method (Johnson et al., 2007) adjusts for batch effects in data using empirical Bayes. Our work falls under the model-based framework, and we are particularly interested in incorporating batch effect removal into a deep learning framework so that the nonlinear batch effect can be modeled and removed.

Removing unwanted information in the deep learning literature is formulated as an invariant representation learning problem. Existing deep learning methods learn invariant representations through either adversarial or nonadversarial methods. Adversarial learning methods, prevalent in domain adaptation, involve a generative adversarial network (GAN) (Creswell et al., 2018) where the generator extracts features to deceive the discriminator, while the discriminator aims to distinguish between different data domains. However, adversarial learning has limitations, including GANs’ mode collapse issues and sensitivity to hyperparameter choices (Papernot et al., 2016). Conversely, nonadversarial invariant methods, including recent group invariance techniques such as contrastive learning (Chen et al., 2020; He et al., 2020; Khosla et al., 2020) and conditionally invariant VAE (Aliee et al., 2023), aim to develop predictors invariant to a group’s action on the input space. Despite their strengths, these methods face a notable limitation: they are designed to handle categorical nuisance variables and may encounter challenges when confronted with continuous-valued nuisance variables. Some deep learning methods aim to mitigate the impact of continuous-valued nuisances, often incorporating mutual information loss. For instance, Greenfeld and Shalit (2020) utilize the Hilbert Schmidt Independence Criterion (HSIC), while Yu et al. (2021) replace HSIC with matrix-based mutual information, yielding performance improvements. Another variant, squared-loss MI (SMI) (Y. Liu et al., 2021), is also popular in this context. However, these mutual information-based methods focus solely on capturing dependencies in data in an unsupervised setting and do not ensure that the learned representations are task relevant. In our work, we advocate for learning motion-invariant representations guided by the information bottleneck principle (Alemi et al., 2016; Moyer et al., 2018). Our approach removes the mutual information between latent representations and undesirable artifacts, with latent representations learned via graph variational autoencoders (VAE) (M. Liu et al., 2021). Situated in the information bottleneck literature, our work addresses the challenge of finding representations that maximize the relationship with the target response while minimizing mutual information with nuisance variables such as subject motion and residual image misalignment. This framework not only addresses the noted limitations of existing methods but also correlates learned invariant representations with human behavior and cognition-related attributes.

We aim to model brain connectomes and learn their representations invariant to nuisance factors, particularly motion-related artifacts. Our analysis of ABCD and HCP data, preprocessed with FSL eddy software (Andersson & Sotiropoulos, 2016), reveals that structural connectome data are significantly affected by motion. This impact is observed not only in the motion quantified by FSL eddy but also in residual misalignment in the diffusion MRI data after FSL eddy correction. Notably, individuals exhibiting substantial head displacement during data acquisition tend to have fewer reconstructed fiber connections across nearly all brain regions. To address this issue, we develop a generative model that learns the distribution of brain connectomes while adjusting for undesirable artifacts induced by motion. This model also learns low-dimensional representations of structural connectomes that are invariant to motion. We achieve this by introducing a penalized VAE objective to eliminate mutual information between connectome representations and motion measures. The learned invariant representation enhances our understanding of motion’s impact on structural connectivity reconstructions. It can also be used to produce motion-adjusted connectomes for downstream analysis tasks, such as predicting behavior- and cognition-related traits and inferring structural connections between brain regions. Our analysis of ABCD and HCP data demonstrates that these motion-invariant representations improve the prediction of many cognitive traits by 5% to 25% compared with connectome representations without motion adjustment.

The rest of the paper is organized as follows: Section 2 introduces the dMRI datasets used for deriving structural brain connectomes and outlines the motion measures utilized in this study. In Section 3, we present our method for adjusting for motion-induced artifacts during the modeling of structural brain networks. Section 4 details a simulation study verifying the effectiveness of our method in removing unwanted information. Section 5 applies our proposed method to two large neuroimaging studies, demonstrating its efficacy in correcting for motion-induced artifacts and establishing a more accurate analysis between structural connectomes and cognitive traits. Section 6 includes a discussion. In the Appendix, we present additional results by considering the residual misalignment after FSL eddy processing as an alternative source of artifacts. Furthermore, we expand the simulation study to assess the robustness of our method when handling extreme motion-related artifacts in the Appendix.

The ABCD study tracks the brain development of over 10,000 adolescents in the United States. We download and process the diffusion MRI (dMRI) data for 8,646 subjects from the NIH Data Archive. Details about data acquisition and preprocessing can be found in Casey et al. (2018). The Human Connectome Project (HCP) characterizes brain connectivity in more than 1,000 adults; see Van Essen et al. (2012) for details about data acquisition. The latest dMRI data from the HCP can be accessed through ConnectomeDB. We download and process 1,065 subjects from the HCP.

To extract structural connectomes from the raw dMRI data, a reproducible tractography algorithm (Maier-Hein et al., 2017) is used to generate the whole-brain tractography for each subject. We use the Desikan–Killiany parcellation (Desikan et al., 2006) to define brain regions of interest (ROIs) corresponding to different nodes in the brain network. The Desikan–Killiany atlas contains 68 ROIs with 34 ROIs belonging to each hemisphere. We use Freesurfer (Iannopollo et al., 2019) to perform brain registration and parcellation. We extract streamlines connecting ROI pairs from the tractography data and the brain parcellation. Fiber count is often used to measure the coupling strength between ROI pairs in the current literature (Fornito et al., 2013; Zhang, Allen, et al., 2019), and, therefore, we summarize each brain connection with its fiber count.

ABCD and HCP use reliable measures to assess a wide range of human behaviors and functions (Gershon et al., 2013). We focus on inferring the relationship between structural connectomes and cognition-related traits, for example, scores from the picture vocabulary and oral reading recognition tests.

Our study considers two sources of motion artifacts: head motion and residual misalignment. Head motion refers to the displacement of the head in the scanner estimated by the eddy tool (Andersson & Sotiropoulos, 2016). Residual misalignment, on the other hand, refers to any remaining misalignment after applying the eddy correction tool. While we analyze both types of artifacts, our primary focus is on head motion because we are interested in the impact of head motion on the diffusion-weighted images used for tractography and the subsequent construction of structural connectivity. We provide additional analysis of the residual misalignment in Appendix G. This section explores how head motion estimates are obtained and its impact on brain structural connectomes.

We primarily rely on the FSL eddy tool to estimate head motion. The eddy software corrects head motion in diffusion MRI data through a multistep process that involves modeling and correcting for both eddy currents and subject movement. It outputs a file named “my_eddy_output.eddy_movement_rms,” which provides summaries of the total movement in each volume. This is calculated by determining the displacement of each voxel, averaging the squares of those displacements across all intracerebral voxels, and finally taking the square root of that average. The file contains two columns: the first column lists the RMS (root mean square) movement relative to the first volume, and the second column lists the RMS movement relative to the previous volume. We use the second measure as our head motion metric in this paper (as it is the only head motion measure accessible in the ABCD data).

Figure 1(A) shows distributions of head motion for the ABCD and HCP studies. We notice that the motion distribution of ABCD subjects is more right skewed with more subjects experiencing bigger motion during image acquisition. This disparity in motion measure distributions between the ABCD and HCP studies could be attributed to the demographic differences in their respective cohorts. Our hypothesis suggests that the HCP study focuses on adults, who tend to follow instructions better and exhibit less movement during image acquisition compared with adolescents in the ABCD study.

To explore the impact of head motion on structural connectomes, we assign individuals to a large and a small motion group corresponding to each movement type. Figure 1 shows the cutoff lines for both ABCD and HCP data. Specifically, in the ABCD data, individuals whose motion estimates are ranked among the top (bottom) 5 percent are assigned to the large (small) motion group. In the HCP data, we use the top and the bottom 10 percent as the threshold. We then quantify brain connection differences between groups. We subtract edge-specific means of the small motion group from those of the large motion group and show the results as adjacency matrices in Figure 1(B). In addition, we visualize between-group differences using a circular layout in Figure 1(C) to examine how connection patterns and strength differ across ROIs. Note that the dMRI data used here are preprocessed with the FSL eddy tool and its motion correction procedures. Figure 1(B and C) reveals that brain networks of large-motion subjects show systematic differences in fiber counts across multiple ROIs compared with those of small-motion subjects, implying that motion artifacts still affect structural connection reconstruction even after motion correction in the preprocessing stage (Andersson & Sotiropoulos, 2016; Andersson et al., 2016). Motivated by these findings, we propose the motion-invariant graph variational autoencoder in Section 3 to remove such systematic artifacts from the connectome data and downstream analyses.

The brain connectome for an individual is represented as an undirected graph $G=(V,ℰ)$ with nodes $V$ and edges $ℰ$⁠. The collection of nodes $V$ refers to the partitioned brain regions, and the set of edges $ℰ$ represents connections between pairs of ROIs. The ROIs are prealigned so that each node in $V$ corresponds to the same brain region across different subjects. We use the symmetric adjacency matrix $A$ to represent $G$ with $Auv$ denoting the number of streamlines connecting ROIs $u$ and $v$⁠; $Auu=0$ since we do not consider self-connections. For $N$ subjects, we denote the network data as $Ai,i=1,2,…,N$⁠. Also, denote the cognitive trait score for an individual as $yi$⁠, and the amount of motion as $ci$⁠. Note that $ci$ can be a scalar quantifying overall head displacement or a vector (⁠$ci∈ℝC$⁠) with each dimension measuring different aspects of motion such as the amplitudes of translation or rotation. Moreover, $ci$ can represent other undesirable artifacts, such as the residual misalignment, depending on the research goal.

We aim to infer the relationship between individuals’ structural connectomes, recovered from neuroimaging data, and their measured cognitive traits, while adjusting for motion-induced artifacts during data acquisition. The key to achieving this goal is to learn a low-dimensional feature representation of $Ai$ as $zi$ that is invariant to motion $ci$⁠, and then relate this feature $zi$ to $yi$⁠. Inspired by the variational autoencoder (VAE) framework in D. P. Kingma (2013), we achieve our goal in three steps: (1) an encoder module (inference model) that learns the mapping from the observation $Ai$ to the motion-invariant latent variable $zi$⁠; (2) a decoder module (generative model) that specifies how to construct $Ai$ from $zi$ and motion $ci$⁠; and (3) a regression module to link the motion-invariant $zi$ to the cognitive trait $yi$⁠; see Figure 2 for the model architecture. Intuitively, in the training process, our approach seeks to maximize the joint likelihood of $(Ai,yi)$ conditional on $ci$ while minimizing the dependence between $zi$ and $ci$⁠. This is achieved by characterizing the joint likelihood through the brain network generative model in Section 3.1 and addressing the dependence through mutual information between the low-dimensional representations of brain networks and motion, as detailed in Section 3.2.

In the following sections, we start with introducing our graph generative model and then discuss the encoder and regression modules.

For each individual $i$⁠, we assume the edges $Ai[uv]$ are conditionally independent given $zi$ and $ci$⁠. The likelihood of the set of edges in $Ai$ is

where $pθ(Ai|zi,ci)$ is the generative model for $Ai$ given $zi$ and $ci$⁠, parameterized by $θ$ learned via neural networks. We assume $Ai$ to be generated from the following process:

where $K$ is the dimension of the latent feature space. The fiber count between ROIs $u$ and $v$ is modeled using a Poisson distribution with a parameter $λi[uv](zi,ci)$⁠, which relates to both the encoded feature $zi$ and motion $ci$⁠.

To capture the varying effect of motion across individuals and brain connections, we further model $λi[uv](zi,ci)$ as

where $⊕$ denotes the concatenation operator and $z˜i∈ℝK+C$ is a concatenated vector representing motion-affected latent features. $ξuv$ is a baseline parameter controlling the population connection strength between regions $u$ and $v$⁠, representing shared structural similarity across individuals, and $ψuv(z˜i)$ characterizes individual connection strength after accounting for motion.

We model $ψuv(z˜i)$ based on the latent space model in Hoff et al. (2002). The key idea is that the connection strength (fiber count) between regions $u$ and $v$ depends on their distance in some network latent space in $ℝR$⁠. In this network latent space, each brain region is represented by a vector in the $ℝR$ space, with the distance between two points reflecting the degree of association between the corresponding brain regions. The closer two points are in this space, the stronger their connection is expected to be. The network latent space is different from the latent space of $z˜i$⁠, which represents a low-dimensional space where the original high-dimensional brain structural connectomes are encoded into a lower dimensional representation.

We construct a mapping $X$ from $ℝK+C$ to $ℝV×R$ that maps the motion-affected feature $z˜i$ to the network latent space that captures the latent positions of all brain regions, that is, $X(z˜i)⊤=(X1(z˜i),…,XV(z˜i))$ with $Xu(z˜i)=(Xu1(z˜i),...,XuR(z˜i))⊤∈ℝR$ for brain region $u$⁠. Specifically, we assume each node $u∈V$ for individual $i$ lies in a latent space $ℝR$ and represent $ψuv(z˜i)$ as

where $αr>0$ is a weight controlling the importance of the $r$-th dimension. A large positive inner product between $Xu(z˜i)$ and $Xv(z˜i)$ implies a large fiber count between this ROI pair. Therefore, $Xu(z˜i)$ incorporates the geometric collaborations among nodes, and broadcasts the impact of motion to all brain connections. M. Liu et al. (2021) proposed a graph convolutional network (GCN) to learn a nonlinear mapping from the encoded feature $zi$ to $(X1(zi),…,XV(zi)),$ leveraging on unique geometric features of the brain networks. The same procedure is adopted in this paper with the additional adjustment for motion $ci$⁠, and we defer the detailed GCN framework to Appendix A.

We are interested in studying the relationship between brain structural networks and human traits, correcting for motion-related artifacts. The key is to find an optimal encoder that produces $zi$ from $Ai$ invariant to $ci$⁠, and then formulate a regression of the trait $yi$ with respect to the latent feature vector $zi$⁠. We represent the inference model (encoder) as $qϕ(zi|Ai)$⁠, the generative model (decoder) as $pθ(Ai|zi,ci)$⁠, and regression as $pθ(yi|zi)$⁠. Our invariant encoding and trait prediction task is to find $qϕ(zi|Ai)$ and $pθ(Ai|zi,ci)$ that maximize $EAi,yi,ci[logpθ(Ai,yi|ci)]$⁠, subject to $zi$ being independent of $ci$⁠. Here $pθ(Ai,yi|ci)$ is the joint likelihood of $Ai,yi$ conditional on $ci$⁠. Since correlation is a limited notion of statistical dependence, which does not capture nonlinear dependencies, we instead use mutual information between $zi$ and $ci$ to quantify their dependency similar to Moyer et al. (2018), and formulate the objective function as follows:

where $I(zi,ci)$ is the mutual information between $zi$ and $ci$⁠, and $λ$ is a trade-off parameter between the joint likelihood and the mutual information. $ϕ$ are parameters of the encoder that are learned via neural networks.

To simplify our model, we assume (1) the human trait $yi$ and the brain connectivity $Ai$ are conditionally independent given the latent representation $zi$ for individual $i$ and (2) only $Ai$ is affected by $ci$ so $yi$ is not affected. Assumption (2) indicates that motion happens randomly given the population considered and does not relate to the human traits (e.g., cognitive ability) considered. In certain scenarios, this assumption might not hold. For example, in the older population, people with mild cognitive impairment (MCI) tend to move more than normal controls (Haller et al., 2014), and if $yi$ is an indicator of MCI, Assumption (2) does not hold here. Assumption (1) indicates that the encoded feature $zi$ contains all the information from $Ai$ needed to predict $yi$⁠, and the residuals are independent of each other. Intuitively, we assume the motion $ci$ explains part of the variability in brain networks $Ai$⁠, but this part of variability in $Ai$ does not relate to $yi$⁠. Of course, we can modify our model to add $ci$ into the prediction of $yi$ in the case that $ci$ is not independent of $yi$⁠. But it will make the final results difficult to interpret since we are mostly concerned with the relationship between $Ai$ and $yi$⁠.

Under Assumptions (1) and (2), we can express the conditional likelihood of $Ai,yi$ given $ci$ as

The derivation is based on Jensen’s inequality; see Appendix B. Intuitively, the lower bound of $logpθ(Ai,yi|ci)$ in (8) consists of three parts. The first term, $log pθ(Ai|zi,ci)$⁠, is a reconstruction error that utilizes log-likelihood to measure how well the generative model in (3) reconstructs $Ai$⁠. The second term, $logpθ(yi|zi)$⁠, measures the trait prediction accuracy. The third KL divergence term is a regularizer that pushes $qϕ(zi|Ai)$ to be close to its prior $N(0,IK)$ so that we can sample $zi$ easily.

For the second term in the objective function (7), we can upper bound $I(zi,ci)$ as

which consists of two parts: the KL divergence between $qϕ(zi|Ai)$ and its marginal $qϕ(zi)$ to ensure less variability across the input data $Ai$⁠, and the reconstruction error. See Appendix C for detailed derivations of (9).

Combining the log-likelihood and the mutual information terms, our training objective in (7) can be expressed as

Our goal is to maximize $ℒ(Ai,yi,ci;θ,ϕ)$ through its lower bound (10). In practice, the expectation in (10) is intractable and we employ Monte Carlo approximation for computation; see Appendix D. We define the approximated objective as $ℒ˜(Ai,yi,ci;θ,ϕ)$⁠, and implement a stochastic variational Bayesian Algorithm 1 to update the parameters using minibatch training. For the regression module, we consider $pθ(yi|zi)$ as a univariate Gaussian, that is, $yi∼N(zi⊤β+b,σ2)$⁠, where $β,b,σ2∈θ$ are parameters to be learned. We list the model parameters and architecture in the Appendix E.

We first conduct a simulation study to verify the efficacy of our method in removing unwanted information from the data. We simulate random graphs with the community structure (Nowicki & Snijders, 2001) using the Python package NetworkX. Random graphs of two communities are considered: Nodes in the same group are connected with a probability of 0.25, and nodes of different groups are connected with a probability of 0.01. We simulate 1,000 such networks ${Ai}i=1n$⁠, where $Ai∈ℝV×V$ with $V=68$ nodes and $n=1,000$⁠. Figure 3(A) displays the average network across 1,000 simulated networks. Next, we simulate a nuisance variable, ${si}i=1n$⁠, by sampling $20%$ of its elements from $N(0.6, 0.05)$ and the rest from $N(1, 0.01)$⁠. We generate the nuisance-affected networks $A˜i$ by propagating $si$ across all edges in $Ai$ as follows:

Intuitively, $si∼N(0.6, 0.05)$ introduces large artifacts, since its multiplication with $Ai$ in (11) results in a $A˜i$ with reduced edge values; $si∼N(1, 0.01)$ generates small artifacts as the multiplication operation in (11) changes the edges of $Ai$ to a lesser extent. The average network across 1,000 such nuisance-affected networks is displayed in Figure 3(A). To visualize the impact of the simulated nuisance variable, the edge-specific differences of the large and small nuisance groups are computed for both simulated networks ${Ai}i=1n$ and nuisance-affected networks ${A˜i}i=1n$ (Fig. 3(B)). We observe apparent differences in edge values between the large and small nuisance groups of ${A˜i}i=1n$⁠. We further apply principal component analysis (PCA) to both ${Ai}i=1n$ and ${A˜i}i=1n$⁠, and visualize their projections onto the first two principal components (PCs) in Figure 3(C). The first two PCs of ${Ai}i=1n$ are indistinguishable. Note that these simulated networks are nuisance free, even though we color their projections with the simulated nuisance for visualization purposes. On the contrary, we observe an apparent separation between the large and small nuisance groups of ${A˜i}i=1n$ after introducing unwanted associations between the nuisance variable and the simulated networks.

To demonstrate the effectiveness of our method in removing artifacts introduced by the nuisance variable, we further compare inv-VAE with the graph autoencoder (GATE) developed in M. Liu et al. (2021). GATE has the same generative model as inv-VAE (see Section 3) except without adjusting for $ci$⁠. Specifically, GATE assumes the following generative process for the brain network $Ai$⁠:

Moreover, GATE learns the latent vector $zi$ using the following objective

instead of the motion-invariant objective in (10). We train GATE using ${A˜i}i=1n$ to learn latent representations of nuisance-affected networks. The learned latents are used by the generative model in (12) to reconstruct networks in Figure 3(A). Figure 3(B) and (C) shows the reconstructed edge-specific differences between the large and small nuisance group, and the first two PCs of the learned latent variables, respectively.

Next, we train inv-VAE with ${A˜i}i=1n$ to learn parameters $ϕ$ and $θ$ for the inference and generative model. The inference model then learns motion-invariant latent representations, from which the generative model produces nuisance-corrected networks (Fig. 3(A)) by setting $si=1$⁠. As previously defined, $si≈1$ corresponds to small nuisance artifacts, whereas a small $si$ represents large nuisance artifacts. This is because multiplication of $Ai$ with a small $si$ in (11) shrinks its edge values to a large degree, whereas multiplication with $si≈1$ keeps the edge values roughly unchanged. Between-group edge-specific differences of the nuisance-corrected networks, and the first two PCs of invariant latents are shown in Figure 3(B, C). Figure 3(B, C) together suggests that GATE reconstructs the artifacts introduced by the nuisance variable, and its learned latent embeddings are affected by nuisance artifacts. On the contrary, such edge-specific differences between the large and small nuisance groups do not exist in the nuisance-corrected networks from inv-VAE, and projections of invariant embeddings of the large and small nuisance groups are indistinguishable. This observation implies that inv-VAE has the ability to remove unwanted information from the data.

A major motivation of our work is to understand how motion affects learning a low-dimensional representation of the structural connectome and to remove motion-induced artifacts from the connectome data. For this purpose, we apply our inv-VAE to the ABCD (8,646 subjects) and HCP data (1,065 subjects), and compare inv-VAE with its competitors on two tasks: (1) learning a latent connectome representation that is uninformative of motion and (2) generating a motion-adjusted connectome. This section focuses on the FSL eddy head motion. We provide an additional analysis in Appendix G on addressing residual misalignment after eddy motion correction.

We first average the motion-affected brain network$Ai$ over all individuals in the study to obtain the motion-affected edge-specific means (see Section 2.2) of the ABCD and HCP datasets, respectively. We call the resulting matrices as “motion-affected,” meaning no motion adjustment has been done. The first columns of Figure 4(A, D) show the motion-affected edge-specific means in both datasets. We assign subjects into a small and a large motion group for each type of movement; See Section 2.2 for the cutoff thresholds for the large and small motion groups in both datasets. The first columns of Figure 4(B, E) show the motion-affected edge-specific differences, obtained by subtracting the motion-affected edge-specific means of the small motion group from those of the large motion group. For both datasets, we notice apparent fiber count differences between the two motion groups in almost all pairs of brain regions. Such brain connection differences due to motion artifacts are more prominent when we visualize the first two PCs of brain networks between the two groups. In the first columns of Figure 4(C, F), each point represents an individual colored with the corresponding amount of motion. In the first column of Figure 4(C), the PCA projections of the large and small motion subjects are closer compared with those in the first column of Figure 4(F). This may be explained by the amount of motion-induced artifacts in the data: participants of the HCP study are young adults who may move less during data acquisition relative to the young adolescent ABCD participants.

Next, we apply GATE (M. Liu et al., 2021) to the ABCD and HCP data. For each dataset, GATE is trained using all brain networks $Ai$’s to learn the model parameters $θ$ and $ϕ$⁠. Equipped with well-trained estimators, GATE obtains the learned latent variables from the inference model and reconstructs brain networks from the generative model in (12). We denote the reconstructed brain network as $A^i$⁠. The second columns of Figure 4(A, D) show the reconstructed edge-specific means averaged across all $A^i$’s in both datasets. In the second columns of Figure 4(B, E), reconstructed edge-specific differences between the large and small motion group are shown. Although GATE can generate brain networks that resemble the observed data, motion artifacts are also undesirably reconstructed, as shown in the second columns of Figure 4(B, E). The second columns of Figure 4(C, F) show the first two PCs of the learned latents by GATE from the large and small motion groups. The gap between the two groups of latents suggests that representations learned by GATE are also corrupted by motion artifacts, which will further affect prediction and inferences.

For our inv-VAE, we use $ci∈ℝ$ to represent the $i$-th subject’s head motion, with 0 representing the smallest amount of motion. To adjust for motion artifacts, we first train the inv-VAE model using the ABCD and HCP data separately. The inference model of inv-VAE learns a set of motion-invariant representations $zi$⁠. Following the generative model in Section 3.1, $zi$ are subsequently used to generate motion-adjusted brain networks, denoted as $Ai*$⁠, by setting $ci$ to 0 to remove motion artifacts. For each dataset, we average across all $Ai*$ to obtain the motion-adjusted edge-specific means shown in the third columns of Figure 4(A, D), which resemble the observed edge-specific means. The third columns of Figure 4(B, E) show the motion-adjusted edge-specific differences by taking the difference between the motion-adjusted edge-specific means of the large and small motion groups. Compared with the color scale of the observed between-group edge-specific differences, the colors of the motion-adjusted between-group edge-specific differences are noticeably lighter, suggesting that brain connection differences between the large and small motion groups become smaller after motion adjustment. The third columns of Figure 4(C, F) show the first two PCs of invariant representations $zi$ for the ABCD and HCP data. The gap between large and small motion subjects is smaller after motion adjustment compared with the projections of the observed brain networks. Although the HCP subjects have smaller head movement and their connectomes are less affected by motion-induced artifacts, we still observe that motion adjustment minimizes the reconstructed structural connection difference between the large and small motion groups in the third column of Figure 4(F).

Motion can significantly impact the acquired diffusion data, introducing signal loss, geometric distortion, misalignment volumes, and compromised diffusion measures. There have been many studies to see how motion impacts the diffusion MRI signal and constructed structural connectomes. For example, Andersson and Sotiropoulos (2016) highlight the susceptibility of diffusion MRI to artifacts caused by subject movements, resulting in both signal loss and geometric distortion. Additionally, Baum et al. (2018) demonstrate that motion has a noticeable impact on a considerable percentage of edges in structural brain networks. This effect is not limited to network edges but extends to significantly affected node strength and total network strength. In this work, we aim to understand how motion impacts our understanding of structural brain connectivity. For both ABCD and HCP studies, we first obtain the motion-affected edge-specific means via averaging $Ai$’s across all individuals in the dataset. After training inv-VAE using all $Ai$’s, we generate motion-adjusted networks $Ai*$’s from the invariant embeddings. We then average across all $Ai*$’s to construct the motion-adjusted edge-specific means. The differences between the motion-adjusted and motion-affected edge-specific means are shown in Figure 5(A, B). In both datasets, we observe more connections between the cingulate and the other brain regions in both hemispheres after motion adjustment, implying that brain connections related to the cingulate are more severely impacted by motion artifacts. There are also notable impacts in other regions, including the frontal, temporal, occipital, and parietal lobes, where we observe significant changes in connection strength after motion adjustment.

Our method can effectively adjust for motion-induced artifacts in the connectome data. A natural question is whether motion-invariant representations further our understanding of the relationship between brain connectomes and cognition-related traits. From the ABCD dataset, we extract the following cognitive traits as $yi$⁠: (1) picture vocabulary test score for language comprehension; (2) oral reading recognition test score for language decoding; (3) fluid composite score for new learning; (4) crystallized composite score for past learning; and (5) cognition total composite score for overall cognition capacity. From the HCP dataset, we consider similar traits: (1) picture vocabulary test score; (2) pattern completion test score for processing speed; (3) picture sequence test score for episodic memory; (4) list sorting test score for working memory; and (5) fluid intelligence score; see Gershon et al. (2013) for details about these traits. All traits are age corrected.

According to Section 3.2, we assume that motion $ci$ is independent of the trait of interest $yi$⁠. However, there might be significant correlations between clinically and scientifically interesting traits and subject motion. To mitigate the concern of erroneously attributing part of the connectome differences to motion rather than underlying traits, we preprocess trait $yi$ by regressing out $ci$⁠. For our analysis, we predict the residuals of $yi$ after regressing out $ci$ and examine the relationship between motion-adjusted connectomes and cognition-related traits.

To investigate whether our method outperforms other approaches in relating structural connectomes to traits, we compare inv-VAE with the following competitors: (1) LR-PCA (motion adjusted) that uses PCA to obtain a lower dimensional projection from each individual’s brain matrix, and then links $yi$ to the projection via a linear regression (LR). Translation and rotation are covariates that LR adjusts for; (2) Sparse Orthogonal to Group (SOG) (Aliverti et al., 2021) that utilizes matrix decomposition to produce an adjusted dataset that is statistically independent of the motion variable. We then regress $yi$ onto the obtained low-dimensional factorizations via LR; (3) ComBat (Johnson et al., 2007) that uses empirical Bayes for batch-effect removal, and returns a motion-adjusted dataset. We then apply PCA to the adjusted data to obtain low-dimensional projections as input to LR; (4) GATE (M. Liu et al., 2021) that has the same architecture as inv-VAE, but does not adjust for motion; see Section 4 for details about the model specification.

We use the above competing methods to produce 68-dimensional connectome representations for each individual, consistent with the recommended latent dimension $K$ in inv-VAE; see Section 3. To assess trait prediction quality, we use the Pearson correlation coefficient to measure the correlation between the observed and predicted traits, and perform fivefold cross-validation (CV). A large correlation coefficient means the predicted trait scores more closely match the observed ones, indicating a stronger association between the low-dimensional representations of structural connectomes and the cognition-related traits.

A comparison of inv-VAE with the other approaches with respect to the trait-prediction quality is displayed in Figure 6, which shows that inv-VAE outperforms LR-PCA, ComBat, and SOG over all studied traits in both datasets, and either outperforms or is comparable with GATE over most selected traits. The prediction quality of SOG is comparable with GATE and inv-VAE for some cognitive traits: oral reading recognition and fluid intelligence score.

Another way to assess the predictive performance is to visualize the associations between low-dimensional representations of structural connectomes and cognition-related traits. Both inv-VAE and GATE produce latent features $zi$ for each individual, which can be used for visualizations. Particularly, for the ABCD study, we consider picture vocabulary test, oral reading recognition test, and crystalized composite score; for the HCP study, we consider picture vocabulary test, oral reading recognition test, and dimensional change card sort test. Both inv-VAE and GATE are trained with brain networks of all individuals in each dataset to obtain the $zi$’s for the $yi$’s. For each trait, latent features from 200 subjects are selected, with the first group of 100 having the lowest trait scores and the second group of 100 having the highest scores. Next, we visualize the first three PCs of the posterior means of $zi$ colored with their corresponding trait scores in Figure 7. For both inv-VAE and GATE, Figure 7 shows that representations from the two trait groups are separated, suggesting that structural connectomes are different among individuals with different cognitive abilities. We highlight that motion-invariant $zi$’s of the two groups are more separated than motion-affected $zi$’s from GATE, particularly for the oral reading recognition test. It implies that motion-invariant structural connectomes are more strongly associated with cognition-related traits than other approaches without motion adjustment.

We develop a motion-invariant variational autoencoder (inv-VAE) for learning low-dimensional, motion-adjusted representations of structural connectomes. We apply inv-VAE to the Adolescent Brain Cognitive Development (ABCD) and Human Connectome Project (HCP) data and discover noticeable motion artifacts in both, despite the incorporation of motion correction procedures in preprocessing. This observation reinforces the need for effective motion mitigation strategies in connectome analysis. Being invariant to motion artifacts during the connectome modeling phase, inv-VAE shows improved performance in understanding the correlation between structural connectomes and cognition-related traits. While our inv-VAE is motivated to handle the motion confounder, it can be easily extended to handle other confounders, such as site batch effect and variation due to machine settings. Compared with popular batch effect removing models such as ComBat, the advantages of our inv-VAE are that (1) it can handle nonlinear objects (brain network data) and (2) it can deal with nonlinear batch effects. Therefore, we believe that our inv-VAE framework can be extended to broader brain imaging analysis tasks, highlighting the importance of confounder or batch effect adjustment methodologies.

Note that inv-VAE is developed under a generative model framework, which allows us to simulate brain networks under various conditions. The simulation of brain networks is crucial since brain network data are generally not widely available, and the extraction of brain networks from raw imaging data is far from straightforward. Our inv-VAE enables us to simulate $(Ai,ci,yi)$—that is, the brain network, undesirable artifacts, and cognition ability measures—simultaneously. Such simulated data can be freely shared with biostatisticians and data scientists who wish to develop statistical models for brain networks, but prefer not to get involved in the intricate process of brain imaging preprocessing.

Our method can model the spatially heterogeneous impact of motion: edges that share the same node can have correlations under our model since they are functions of both $zi$ and $ci$⁠. In this work, we simplify $ci$ to encode only a summary of subject-level motion artifacts. One possible future direction is to explore the representations of motion that can improve our generative model.

Our proposed inv-VAE currently models the motion-affected structural connectome using a Poisson generative model in conjunction with a GCN. Future work could incorporate pre-existing neuroscience knowledge regarding the effects of motion-induced artifacts on specific brain regions or connections to refine our model and improve the reconstruction of motion-adjusted structural connectomes. This could be achieved by leveraging empirical evidence, neuroscientific findings, or theoretical frameworks that shed light on the disparate impacts of motion on individual brain regions or connections. Through this integration of knowledge, we aim to boost the precision and accuracy of our motion artifact mitigation process, enabling us to better capture the subtle influences of motion on brain connectivity.

Another future direction is to leverage recent developments of generative modeling frameworks, such as variational diffusion models (D. Kingma et al., 2021), score-based generative models (Song & Ermon, 2020), and Hierarchical Variational Autoencoders (HVAE) (Sønderby et al., 2016). These models can be understood under the general Evidence Lower Bound (ELBO) method used here in this paper but with different constraints in the latent space (Luo, 2022). Moreover, we use the mean motion estimate across frames, as computed by the FSL eddy tool, to represent motion during diffusion MRI acquisition and ignore the longitudinal or higher order information in the motion. This simplification may not fully represent the impacts of motion on diffusion MRI data, and thus the motion-adjusted connectomes from our inv-VAE may not be completely invariant to motion. Future developments may incorporate more complex summaries of dynamic motion.

The data used in this study contain publicly available datasets from the Human Connectome Project^* and the Adolescent Brain Cognitive Development Study^†. The model training and testing code is available at https://github.com/yzhang511/inv-vae^‡.

Conceptualization and Study Design: Yizi Zhang, Meimei Liu, Zhengwu Zhang. Data curation: Zhengwu Zhang. Methodology, Formal Analysis, Visualizations: Yizi Zhang, Meimei Liu. Manuscript Writing (original draft): Yizi Zhang. Manuscript Writing (review & editing): Yizi Zhang, Meimei Liu, Zhengwu Zhang, David Dunson. Supervision: David Dunson. All authors have substantively revised the work, reviewed the manuscript, approved the submitted version, and agreed to be personally accountable for their contributions.

The datasets were anonymized and not collected by the investigators, in which case the work is classified as nonhuman research.

The authors declare no competing interests.

This work is supported by grants from the National Science Foundation (NSF-DMS-2124535).

This section provides model details and additional numerical studies. The appendix is organized as follows:

• Section A: Specification of GCN details used in both inv-VAE and GATE.

• Section B: Derivation of the joint likelihood of brain networks $Ai$ and cognition traits $yi$ conditional on the nuisance factor $ci$ as defined in Equation (8).

• Section C: Derivation of the mutual information between latent variables $zi$ and $ci$ as defined in Equation (7).

• Section D and E: Implementation details, including Monte Carlo approximation of the training objective and model architecture details of our developed inv-VAE.

• Section F: Investigation of correlation between motion score and outlier measure to assess the relationship between head motion, residual misalignment, and outlier data.

• Section G: Analysis of the effect of residual misalignment adjustment on brain connectome reconstruction and trait prediction.

• Section H: Expanded simulation study to evaluate our method’s robustness in handling extreme motion-induced artifacts.

• Section I: Methods for extracting residual misregistration representations.

To exploit the local collaborative patterns among brain regions, M. Liu et al. (2021) propose to learn each region’s representation by propagating node-specific $k$-nearest neighbor information. We use the relative distance between a pair of brain regions to represent their locality. The relative distance is measured through the length of the white matter fiber tract connecting the pair, and stored in a matrix $D∈ℝV×V$⁠. $Duv$ denotes the fiber length between regions $u$ and $v$⁠. For $u$⁠, its $k$-nearest neighbors (kNNs) can be defined according to $D$⁠. Graph convolution is defined to learn the $r$-th latent coordinate $Xr(z˜i)$ for subject $i$⁠. In particular, an $m$-layer GCN is $Xr(m)(z˜i)=hm(Wr(m)Xr(m−1)(z˜i)+bm),$ for $2≤m≤M$⁠, with $Xr(1)(z˜i)=h1(Wr(1)z˜i+b1)$⁠. $Xr(m)(z˜i)$ denotes the output of the $m$-th GCN layer, $hm(⋅)$ is an activation function of the $m$-th layer, and $Wr(m)$ is a weight matrix characterizing the convolution operation at the $m$-th layer. The embedding feature of each region at the $m$-th layer is determined by the weighted sum of itself and its nearest neighbor regions at the $(m−1)$-th layer. When $m=1$⁠, $Wr(1)∈ℝV×(K+C)$ maps $z˜i∈ℝK+C$ to the latent space $Xr(1)(z˜i)∈ℝV$⁠. When $m≥2$⁠, $Wr(m)$ is a $V×V$ matrix with the $u$-th row $wu⋅(r,m)$ satisfying $wuv(r,m)>0$ if $v=u$ or $v∈krNN(u)$⁠, and $wuv(r,m)=0$ otherwise. Here, $krNN(u)$ denotes the $k$ nearest neighbors of region $u$ in the $r$-th dimension of the latent space.

The joint likelihood of $Ai$ and $yi$ given $ci$ is

By Jensen’s inequality,

Assuming that $zi$ is independent of $ci$⁠, we have $pθ(zi|ci)=p(zi)$⁠. We then have