A Matrix Autoencoder Framework to Align the Functional and Structural Connectivity Manifolds as Guided by Behavioral Phenotypes

Functional Connectivity Reconstruction: By construction, the correlation matrices $\mathbf{\Gamma}_{n}$ belong to the manifold of symmetric positive semi-definite matrices $\mathcal{P}^{+}_{P}$. Our matrix autoencoder estimates a latent functional embedding $\mathbf{F}_{n}\in\mathcal{P}_{K}^{+}$ using a 2D fully connected (2D FC-NN) layer [9, 17]. Formally, this mapping $\mathbf{\Phi}_{\text{ec}}(\cdot):\mathcal{P}^{+}_{P}\rightarrow\mathcal{P}^{+}_{K}$ is parametrized by weights $\mathbf{W}\in\mathcal{R}^{P\times K}$ and is computed as a cascade of two linear layers with tied weights: $\mathbf{F}_{n}=\mathbf{\Phi}_{\text{ec}}(\mathbf{\Gamma}_{n})=\mathbf{W}^{T}\mathbf{\Gamma}_{n}\mathbf{W}$. Our decoder is another 2D FC-NN that estimates $\tilde{\mathbf{\Gamma}}_{n}$ from $\mathbf{F}_{n}$ via a similar transformation $\mathbf{\Phi}_{\text{dc}}(\cdot):\mathcal{P}^{+}_{K}\rightarrow\mathcal{P}^{+}_{P}$ that shares weights with the encoder. Mathematically, our FC reconstruction loss is represented as follows:

The second term of Eq. (1) encourages the columns of the brain basis $\mathbf{W}$ to be orthonormal. Conceptually, this specialized loss helps us learn uncorrelated patterns that explain the rs-fMRI data well while acting as an implicit regularizer.

Structural Connectivity Estimation: The structural connectivity matrices $\mathbf{A}_{n}$ are derived from DTI tractography and belong to the manifold of symmetric (non PSD) matrices $\mathcal{S}_{P}$. Our alignment decoder first generates an SC embedding $\mathbf{S}_{n}\in\mathcal{R}^{K\times K}$ from $\mathbf{F}_{n}$ via a 2D FC-NN layer $\mathbf{\Phi}_{\text{align}}(\cdot):\mathcal{P}^{+}_{K}\rightarrow\mathcal{P}^{+}_{K}$, followed by a second 2D FC-NN layer $\mathbf{\Phi}_{\text{est}}(\cdot):\mathcal{P}^{+}_{K}\rightarrow\mathcal{P}^{+}_{P}$ which maps to the structural connectivity matrices. For stability our SC matrices do not have self-connections and are normalized to $\lVert\mathbf{A}_{n}\rVert_{1}=1$. Accordingly, at the output layer, we apply a 2D softmax $\mathcal{SF}(\cdot)$, and then suppress the diagonal elements to generate the final output $\tilde{\mathbf{A}}_{n}\in\mathcal{R}^{P\times P}$. Our SC estimation objective is represented as follows:

where $\circ$ is the element-wise Hadamard product. $\mathbf{1}\in\mathcal{R}^{P\times 1}$ is the vector of all ones, and $\mathcal{I}_{P}$ is the identity matrix of dimension $P$. Conceptually, the loss in Eq. (2) is akin to manifold alignment [37] between the functional and structural embeddings based on a two sided Procrustes-like objective.

Phenotypic Prediction: We map the intermediate representation $\mathbf{X}_{n}=\mathbf{\Gamma}_{n}\mathbf{W}$ $\in\mathcal{R}^{P\times K}$ learned by the FC encoder to the phenotypes $\mathbf{y}_{n}$ via a cascade of a 1D convolutional layer and an ANN. The convolutional layer $\mathcal{F}_{\text{conv}}(\cdot)$ collapses $\mathbf{X}_{n}$ along its rows via a weighted sum to generate a $K$ dimensional feature vector. This feature vector is input to a simple two layered ANN $\mathcal{G}(\cdot)$ to jointly estimate the elements in $\hat{\mathbf{y}}_{n}$. We use a Mean Squared Error (MSE) loss function:

This prediction task is a secondary regularizer that encourages our matrix autoencoder to learn representations predictive of inter-subject variability.

Implementation Details: We train our framework on a joint objective that combines Eqs. (1), (2) and (3) as follows:

where $\gamma_{1}$ and $\gamma_{2}$ balance the tradeoff for the SC estimation and phenotypic prediction relative to the FC reconstruction objective. We employ an ADAM optimizer [22] with learning rate $0.005$ and weight decay regularization [29] ($\delta=0.0005$) run for a maximum of $400$ epochs. Optimization parameters were fixed based on a validation set consisting of $30$ additional patients from the HCP database. We use this strategy to set the the dimensionality of our autoencoder embedding at $K=15$ and loss penalities to $\{\gamma_{1},\gamma_{2}\}=10^{3},3$. Finally, we utilize a spectral initialization for the encoder-decoder weights $\mathbf{W}$ in Eq. (1) based on the top $K$ eigenvectors of the average patient correlation matrix $\bar{\mathbf{\Gamma}}_{n}$ for the training set. We use a similar initialization based on $\bar{\mathbf{A}}_{n}$ for $\mathbf{\Phi}_{\text{est}}(\cdot)$, and default initialization [24] for the remaining layers. Our model has a runtime of 10-12 minutes on an $8$ core machine with $32$GB RAM implemented in PyTorch (v1.5.1).

HCP Dataset and Preprocessing: We download rs-fMRI and DTI scans for of $275$ healthy individuals from the Human Connectome Project (HCP) S1200 database [12]. Rs-fMRI data was pre-processed according to the standard HCP pipeline [35], which accounts for motion and physiological confounds. We selected a $15$-minute interval from the rs-fMRI scans to remain commensurate with standard rs-fMRI protocols. We used the Neurodata MR Graphs package [21] to pre-process the DTI scans and estimate fiber bundles via streamline tractography. Our phenotypic measure was the Cognitive Fluid Intelligence Score (CFIS) [10, 5] adjusted for age, which is obtained via a battery of tests measuring cognitive reasoning (dynamic range: $70-150$). We used the Automatic Anatomical Labelling (AAL) atlas with $116$ ROIs as our parcellation. From here, we compute the the Pearson correlation between the the mean rs-fMRI time course in the two regions; we subtract the contribution of the top (roughly constant) eigenvector to obtain the input $\mathbf{\Gamma}_{n}$. The input $\mathbf{A}_{n}$ are obtained by dividing the number of tracts between ROIs by the total fiber tracts across all ROI pairs.

Predicting Behavioral Phenotypes: Table 1 (and Fig. $1$ in Supplementary) compares the model against the baselines when predicting CFIS in a five-fold cross validated setting. Lower Median Absolute Error (MAE) and higher Normalized Mutual Information (NMI) and correlation coefficient (R) signify improved performance. Our framework outperforms the baselines during testing, though the model of [11] comes in a close second. This suggests that the Matrix Autoencoder faithfully models subject-specific variation even in unseen patients.

Functional to Structural Association: We evaluate three aspects of our functional to structural manifold alignment. First is our ability to recover structural connectivity matrices during testing. Here, we compare two distance metrics: (1) $F_{\text{self}}$ is the Frobenius norm between a test example $\mathbf{A}_{n}$ and the model prediction for the same example $\hat{\mathbf{A}}_{n}$, and (2) $F_{\text{other}}$ is $\hat{\mathbf{A}}_{n}$ and other SC matrices $\mathbf{A}_{m},(m\neq n)$. As shown in the left of Fig. 2(a), $F_{\text{self}}$ is consistently smaller than $F_{\text{other}}$, with statistical significance determined using the Wilcoxon rank sum test. This indicates that individual differences in SC are preserved by our framework. In the same plot, we also benchmark the recovery performance of our framework against a baseline Matrix encoder-decoder (gray box in Fig. 1) with the SC matrices as input and output. We also compare against a linear regression between the vectorized upper diagonal FC features (input) and SC features (output) to help evaluate the benefit of our matrix decomposition. As seen, our function $\rightarrow$ structure decoding achieves similar performance as directly encoding/decoding the structural connectivity. At the same time, the linear regression baseline performs worse than both of these techniques. This suggests that the ability to directly leverage the low rank matrix structure is key to preserving individual differences during reconstruction.

Examining the FC Biomarkers: We explore the functional connectivity patterns learned by our framework by first matching the brain bases (i.e., columns of $\mathbf{W}$) across the cross validation folds based on the absolute correlation coefficient. We run this experiment five times with different initializations for the ANN branch to check for consistency in the learned representation. Fig 3 displays the four most consistent bases, as projected onto the brain using the region definitions of the AAL atlas. We notice that while there is spatial overlap between the bases, the standard deviations are small, which indicates that our framework is learning stable patters in the data. Subnetwork $1$ highlights regions from the default mode network, which is widely inferred within the resting state literature, and known to play a critical role in consolidating working memory [34]. Subnetworks $1$, $3$ and $4$ highlight regions from the somatomotor network and visual cortex, together believed to be important functional biomarkers of cognitive intelligence [7]. Finally, Subnetwork $2$ and $4$ displays contributions from the frontoparietal network and the medial prefrontal network. These areas are believed to play a role in working memory, attention, and decision making, all of which are associated with cognitive intelligence [30].

Application to a Secondary Dataset: We evaluate our framework on a second clinical cohort of $57$ children with high-functioning Autism Spectrum Disorder (ASD). Rs-fMRI and DTI scans were acquired on a Philips $3T$ Achieva scanner (rs-fMRI: EPI, TR/TE = $2500/30$ms, res = $3.05\times 3.15\times 3$mm, duration = $128$ or $156$ time samples; DTI: EPI, TR/TE = $6356/75$ms, res = $0.8\times 0.8\times 2.2$mm, b-value = $700$s/$mm^{2}$, $32$ gradient directions). Rs-fMRI preprocessing includes motion correction, normalization to MNI, spatial and temporal filtering, and nuisance regression with CompCorr [3]. DTI data was preprocessed using the FDT pipeline in FSL [18], with tractography performed using the BEDPOSTx and PROBTRACKx functions in FSL [2]. We define $\mathbf{y}_{n}$ using three phenotypes that characterize various impairments associated with ASD: (1) the Autism Diagnostic Observation Schedule (ADOS) [33], (2) the Social Responsiveness Scale (SRS) [33], and (3) Praxis [32]. We carry forward the same parcellation scheme and model parameters as used for the HCP dataset. Table 2 compares the multi-score prediction testing performance of ADOS, SRS, and Praxis in a five fold cross validation setting. We observe that only our framework and the model of [11] can simultaneously predict all three measures. In contrast, the other baselines achieve good testing performance on one or two of the measures (for example, No DTI decoder baseline for ADOS and SRS) but cannot generalize all three. Our supplementary document includes additional results for SC recovery and FC biomarker extraction, and scatter plots for phenotypic prediction on the ASD datatset. Overall, our experiments on both healthy (HCP) and clinical (ASD) populations suggest that our model is robust across cohorts and generalizes effectively even with modest dataset sizes.