INDIVIDUALIZED TRAJECTORY PREDICTION OF EARLY DEVELOPING FUNCTIONAL CONNECTIVITY

Abstract

Predicting the development of functional connectivity (FC) derived from resting-state functional MRI is pivotal for elucidating the intrinsic brain functional organization and modeling its dynamic development during infancy. Existing deep learning methods typically predict FC at a target timepoint from each available FC independently, yielding inconsistent predictions and overlooking longitudinal dependencies, which introduce ambiguity in practical applications. Furthermore, the scarcity and irregular distribution of longitudinal rs-fMRI data pose significant challenges in accurately predicting and delineating the trajectories of early brain functional development. To address these issues, we propose a novel Triplet Cycle-Consistent Masked Autoencoder (TC-MAE) for the trajectory prediction of the development of infant FC. Our TC-MAE has the capability to traverse FC over an extended period, extract unique individual characteristics, and predict target FC at any given age in infancy with longitudinal consistency. Extensive experiments on 368 longitudinal infant rs-fMRI scans demonstrate the superior performance of the proposed method in longitudinal FC prediction compared with state-of-the-art approaches.

1. INTRODUCTION

Functional connectivity (FC) derived from resting-state functional magnetic resonance imaging (rs-fMRI) plays a vital role in unraveling the intricate brain functional organization [1]. Researches on FC and its applications span a variety of fields. These include mutual prediction between FC and structural connectivity [2, 3], prediction between awake-state FC and sleep-state FC [4], and early detection of brain diseases based on FC [5]. Of note, subtle abnormal changes in FC during infancy may undergo significant amplification as various developmental processes unfold, resulting in more profound adverse effects [6], compared to FCs of adults. Predicting FC can help address issues such as longitudinal data augmentation, missing data imputation, identification of crucial biomarkers in functional organization, advancing our understanding of developmental patterns and mechanisms of individualized functional organization. However, current approaches primarily treat FC prediction at different timepoints as isolated tasks, neglecting the cross-correlation of FC across ages, resulting in poor temporal coherence and ambiguity in practical applications. Moreover, due to the strong motion effect and long acquisition time, current limited longitudinal infant FC data with highly irregular temporal distributions causes huge difficulties in individual feature learning and prediction. Despite the adoption of various advanced techniques [2, 7], the above issues remain inadequately addressed. As a result, predicting the individualized FC trajectory in infants becomes a task that is highly desired but challenging.

To address these challenges, we propose a novel Triplet Cycle-Consistent Masked Autoencoder (TC-MAE) for the longitudinal individualized prediction of early developing FC. The model can traverse FC over an extended period, extract unique individual characteristics, and predict the developmental trajectories of FC across ages during infancy. Specifically, 1) instead of predicting FC at different timepoints independently, we introduced the dynamic time warping loss [8], which is widely used to assess the similarity between time series, along with a cycle-consistent scheme to utilize longitudinal trajectory information across subjects’ ages, thus enabling more accurate predictions of the developmental trajectories of FC; 2) to address the issue of limited infant FC data, inspirated by Masked Autoencoder [9], we designed a Masked Autoencoder tailored specifically for FC prediction. To this end, we leveraged abundant adult FC data to pretrain a robust feature extraction encoder and extend its effectiveness to infant FC; 3) we introduced a dual-triplet network aiming at effectively disentangling age-related and identity(ID)-related information, thus accurately separating brain development-related features from ID-related ones. Then, ID-related features will be fused with developmental information by a dedicatedly designed identity conditional module to generate FCs for infants at all ages. Experiments on 368 longitudinal infant rs-fMRI scans demonstrate the superior performance of the proposed method in longitudinal FC prediction compared with state-of-the-art approaches.

2. METHODS

2.1. Architecture Overview

The architecture of our proposed Triplet Cycle-Consistent Masked Autoencoder (TC-MAE) is shown in Fig. 1. Specifically, we separately design a model backbone and three key stages of training: Pretraining, Triplet Reconstruction, and Longitudinal Prediction. Following these three stages, the model backbone can generate individualized FC development trajectory for any subject during the testing phase.

Fig. 1.

The framework of the proposed Triplet Cycle-Consistent Masked Autoencoder (TC-MAE) includes three stages: pretraining, triplet reconstruction, and longitudinal prediction. In the pretraining stage, only the encoder and decoder are trained, while all components are involved in the subsequent stages.

Specifically, to predict infant FC across ages and capture individualized, temporally coherent FC progression, our proposed model backbone, TC-MAE, is composed of three key components, i.e., an encoder and a decoder, a feature disentanglement module, and an identity condition module. The training stage of this model contains three distinct sub-stages. In the Pretraining stage, the encoder and decoder of a Masked Autoencoder are pretrained using abundant adult FC data to capture the underlying relationships between different patches in FC. In the Triplet Reconstruction stage, we establish a new triplet network configuration consisting of three parallel branches of the backbone that can share weights. All components are trained at this stage to capture more latent relationships in FC. Simultaneously, an Identity Conditional Module (ICM) is employed to incorporate age-related information, while the triplet network disentangles individual ID-related information and age-related information. The Longitudinal Prediction stage introduces a ring prediction mechanism across multiple timepoints and leverages specific loss to capture longitudinal information in infant FC development. This novel approach aims to preserve temporal consistency in prediction results, ensuring a coherent developmental trajectory of FC. Finally, in the testing stage (Fig. 1(e)), the trained model takes an available FC as input to jointly predict multi-time-point FC during infancy.

2.2. Pretraining Stage

To address the scarcity of infant FC data, we leverage abundant adult FC data for pretraining to learn and extract FC domain features. We specifically tailored the Masked Autoencoder [9] for FC, as depicted in Fig.1(a), which has outstanding ability to extract domain and holistic information from FC as a pretraining framework. In our model, unlike dividing images into square patches in traditional MAE, we dissect the FC matrix $\mathit{x}$ into $n$ patches by rows to keep the FC features of each row corresponding to a brain region. We use random mask module, which selectively obscures most of patches each time with a ratio $\phi$, denoted as $\mathit{\text{Masked}}(\mathit{x},\phi )$. The consecutive reconstruction task is designed to compel the encoder $\mathit{E}$ to focus on global information from the unmasked patches and restore the entire FC by the decoder $\mathit{G}$. Based on the mean square loss, the pretraining loss of the entire Pretraining Stage is defined as follows:

$${ℒ}_{\mathit{\text{PT}}}=\Vert \mathit{G}\left(\mathit{E}\right(\mathit{\text{Masked}}(\mathit{x},\phi )\left)\right)-\mathit{x}{\Vert }_2^2$$ (1)

2.3. Triplet Reconstruction Network

To effectively disentangle mixed age-related and ID-related information and extract features from the source FC, we introduce a triplet design [10] and structure the training samples in units $\left({\mathit{x}}_{\mathit{i}}^{t_1},{\mathit{x}}_{\mathit{i}}^{t_2},{\mathit{x}}_{\mathit{j}}^{t_1}\right)$, where the first two elements represent FCs from the same individual $\mathit{i}$ at different ages $t_1$ and $t_2$, and the third element is an FC from a different individual $\mathit{j}$ at age $t_1$. As depicted in Fig.1 (b), during the stage of triplet reconstruction, the encoder $\mathit{E}$ and decoder $\mathit{G}$ trained in the previous stage will be retained and fine-tuned using infant FC data. A shared ID extractor and an age extractor disentangle the the output of the encoder into two distinct components: $\mathit{I}$ and $\mathit{A}$, representing ID-related and age-related information, respectively. $\mathit{A}$ is then channeled into an age predictor $\mathit{P}$, a key downstream task that enhances the effectiveness of the disentanglement. The age prediction loss is defined as follows: ${ℒ}_{\mathit{\text{age}}}={\Sigma }_{\mathit{x}={\mathit{x}}_{\mathit{i}}^{t_1},{\mathit{x}}_{\mathit{i}}^{t_2},{\mathit{x}}_j^{t_1}}\left|y_{\mathit{x}}-\mathit{P}\left({\mathit{A}}_{\mathit{x}}\right)\right|$. Where $y_{\mathit{x}}$ is the real age of the FC $\mathit{x}$.

Simultaneously, the $\mathit{I}$ is directed to the identity conditional module (ICM), designed to learn identity-restricted age-conditional progression patterns. The ICM, as shown in Fig. 1(d), containing a fully connected feed-forward network, can convert the input $\mathit{I}$ into a feature representation based on the age group and select the corresponding feature according to the input target age condition. This process is repeated six times to enhance extraction effect, and finally, the result is input into the decoder $\mathit{G}$ to complete the reconstruction. The reconstruction is evaluated using Pearson’s correlation coefficient (pcc) and mean absolute error (mae), which measure the similarity between the source FC $\mathit{x}$ and the recovered FC $\stackrel{\mathit{ˆ}}{\mathit{x}}=\mathit{G}\left(\mathit{\text{ICM}}\right(\mathit{I},t\left)\right)$. The pcc loss, mae loss and the joint reconstruction loss are defined as:

$${ℒ}_{\mathit{\text{recon}}}^{\mathit{\text{pcc}}}={\Sigma }_{\mathit{x}={\mathit{x}}_{\mathit{i}}^{t_1},{\mathit{x}}_{\mathit{i}}^{t_2},{\mathit{x}}_{\mathit{j}}^{t_1}}\mathbb{C}(\mathit{x},\mathit{G}(\mathit{\text{ICM}}(\mathit{I},t)\left)\right)$$ (2)

$${ℒ}_{\mathit{\text{recon}}}^{\mathit{\text{mae}}}={\Sigma }_{\mathit{x}={\mathit{x}}_{\mathit{i}}^{t_1},{\mathit{x}}_{\mathit{i}}^{t_2},{\mathit{x}}_{\mathit{j}}^{t_1}}\Vert \mathit{x}-\mathit{G}\left(\mathit{\text{ICM}}\right(\mathit{I},t\left)\right){\Vert }_1$$ (3)

$${ℒ}_{\mathit{\text{recon}}}={ℒ}_{\mathit{\text{recon}}}^{\mathit{\text{mae}}}-\beta {ℒ}_{\mathit{\text{recon}}}^{\mathit{\text{pcc}}}$$ (4)

where $t\in \left\{t_1,t_2\right\}$, corresponding to the subjects’ age, $\mathbb{C}$ represents the Pearson’s correlation and $\beta$ balances the respective loss terms.

To enhance the discrimination of $\mathit{I}$ and $\mathit{A}$ among different subjects, we incorporate a dual intensive triplet loss. Specifically, ${\mathit{I}}_{\mathit{i}}^{t_2}$ and ${\mathit{I}}_{\mathit{j}}^{t_1}$ are serving as the ID-positive and ID-negative samples for ${\mathit{I}}_{\mathit{i}}^{t_1}$, while ${\mathit{A}}_{\mathit{j}}^{t_1}$ and ${\mathit{A}}_{\mathit{i}}^{t_2}$ represent the age-positive and age-negative samples for ${\mathit{A}}_{\mathit{i}}^{t_1}$. This dual intensive triplet loss is defined to enforce the similarity of the features of similar-age pairs and same-ID pairs as follows:

$${ℒ}_{\mathit{\text{tri}}}^{\mathit{\text{ID}}}=\mathbb{C}\left({\mathit{I}}_{\mathit{i}}^{t_1},{\mathit{I}}_{\mathit{j}}^{t_1}\right)+\mathbb{C}\left({\mathit{I}}_{\mathit{i}}^{t_2},{\mathit{I}}_{\mathit{j}}^{t_1}\right)-2*\mathbb{C}\left({\mathit{I}}_{\mathit{i}}^{t_1},{\mathit{I}}_{\mathit{i}}^{t_2}\right)$$ (5)

$${ℒ}_{\mathit{\text{tri}}}^{\mathit{\text{Age}}}=\mathbb{C}\left({\mathit{A}}_{\mathit{i}}^{t_1},{\mathit{A}}_{\mathit{i}}^{t_2}\right)+\mathbb{C}\left({\mathit{A}}_{\mathit{j}}^{t_1},{\mathit{A}}_{\mathit{i}}^{t_2}\right)-2*\mathbb{C}\left({\mathit{A}}_{\mathit{i}}^{t_1},{\mathit{A}}_{\mathit{j}}^{t_1}\right)$$ (6)

$${ℒ}_{\mathit{\text{tri}}}=L_{\mathit{\text{tri}}}^{\mathit{\text{ID}}}+L_{\mathit{\text{tri}}}^{\mathit{\text{Age}}}$$ (7)

Finally, the loss of the triplet reconstruction network, ${ℒ}_{\mathit{\text{TR}}}$, integrating the age prediction loss, intensive triplet loss, and joint reconstruction loss, is designed as: ${ℒ}_{\mathit{\text{TR}}}={ℒ}_{\mathit{\text{recon}}}+\alpha {ℒ}_{\mathit{\text{age}}}+\delta {ℒ}_{\mathit{\text{tri}}}$, where $\alpha$ and $\delta$ govern the impact of age prediction loss and dual intensive triplet loss, respectively.

2.4. Longitudinal Prediction

To reinforce the temporal consistency between longitudinal FCs, we introduce a Longitudinal Prediction Network (Fig. 1(c)), to predict the trajectory ${𝒯}_i=\left\{{\stackrel{\mathit{ˆ}}{\mathit{x}}}_{\mathit{i}}^1,\dots ,{\widehat{\mathit{x}}}_{\mathit{i}}^m\right\}$, which consists of a set of predicted FCs across the first two years of age. A dynamic time warping (DTW) loss [8] is adopted to ensure the predicted trajectories based on different inputs for the same subject are similar. The DTW loss for $N$ subject is calculated as:

$${ℒ}_{\mathit{\text{DTW}}}={\Sigma }_{i=1}^N\mathit{\text{softDT}}{W}_{\gamma }\left({𝒯}_i^{t_1},{𝒯}_i^{t_2}\right)$$ (8)

where $softDT{W}_{\gamma }$ is the dynamic time warping distance, $\gamma$ is set to 0.1 and ${𝒯}_i^{t_1}$ and ${𝒯}_i^{t_2}$ are the predicted trajectories based on $t_1$ and $t_2$, respectively. Meanwhile, inspired by [11] in image translation, we incorporate a decoder with additional constraints on backward mapping, ensuring the consistency of longitudinal predictions. After generating the FC trajectory from $t_1$, we select one FC, ${\stackrel{\mathit{ˆ}}{\mathit{x}}}_i^{t_2}$, which has a ground truth at the corresponding timepoint $t_2$, and perform a reverse mapping to predict the original input FC: ${\widehat{\mathit{x}}}_i^{t_2}=\mathit{G}\left(\mathit{\text{ICM}}\left({\mathit{I}}_i^{t_1},t_2\right)\right),{\widehat{\widehat{\mathit{x}}}}_i^{t_1}=\mathit{G}\left(\mathit{\text{ICM}}\left({\widehat{\mathit{I}}}_i^{t_2},t_1\right)\right)$. We then impose a cycle-consistency loss to ensure both the generated FC ${\stackrel{\mathit{ˆ}}{\mathit{x}}}_i^{t_2}$ and the reconstructed FC ${\widehat{\widehat{\mathit{x}}}}_i^{t_1}$ are meaningful representations of the target FC ${\mathit{x}}_i^{t_2}$ and the original FC ${\mathit{x}}_i^{t_1}$. The cycle-consistent loss ${ℒ}_{\mathit{\text{cycle}}}$ and the overall prediction loss ${ℒ}_{\mathit{\text{Pre}}}$ are defined as follows:

$${ℒ}_{\mathit{\text{cycle}}}=\epsilon {‖{\mathit{x}}_i^{t_2}-{\widehat{\mathit{x}}}_i^{t_2}‖}_2^2+{‖{\mathit{x}}_i^{t_1}-{\widehat{\widehat{\mathit{x}}}}_i^{t_1}‖}_2^2$$ (9)

$${ℒ}_{\mathit{\text{Pre}}}=L_{\mathit{\text{DTW}}}+L_{\mathit{\text{cycle}}}$$ (10)

3. EXPERIMENT

3.1. Dataset

During the pretraining stage, we employ 2,171 resting-state fMRI scans of 1,086 adult subjects from the Human Connectome Project (HCP) dataset [12]. In the subsequent stages, we utilize 368 longitudinal resting-state fMRI scans of infants aged younger than 600 days from the Baby Connectome Project (BCP) dataset [13]. After a detailed investigation of the dataset, we divide the BCP dataset into six temporal partitions, which are [1, 60], [61, 120], [121, 210], [211, 300], [301, 450] and [451, 600] days, corresponding to six early developing periods. This non-uniform division captures infants’ rapid developmental changes in the first 6 months with finer temporal granularity. For infant structural and functional MRI processing, we follow the methodologies detailed in [10, 14, 15, 16] to extract the fMRI time-series for each vertex on middle cortical surfaces, all-time series within each cortical region are then further averaged based on the parcellation scheme in [17]. We compute functional connectivity using Pearson’s correlation and Fisher’s r-to-z transformation, focusing on the absolute value of FCs [18]. All subjects are stratified and split into training and testing set with a portion of 4:1, and a 5-fold cross-validation in training set is used for tuning the parameters.

3.2. Implementation Details and Ablation Study

The encoder and decoder adopt the ViT-base model [19]. The ID extractor, age extractor and age predictor consist of fully connected layers with dimensions (768, 668), (768, 100) and (100, 1), respectively. An ablation study is conducted to evaluate the contributions of all components. The experimental results (Table 1) indicate that all components of TC-MAE contribute to the result. Our Pretraining stage successfully captured domain features, and the Triplet Reconstruction Network along with Longitudinal Prediction Network further enhanced the model’s robustness and accuracy.

Table 1.

Ablation study of TC-MAE (Mean ± Std).

Method	mae	pcc
TC-MAE	0.104 ± 0.02*	0.790 ± 0.06*
w/o Pretraining	0.136 ± 0.03	0.726 ± 0.09
w/o Triplet Reconstruction	0.120 ± 0.03	0.780 ± 0.09
w/o Longitudinal Prediction	0.120 ± 0.02	0.752 ± 0.07

^*indicates statistically significantly better results than other methods with p-value<0.05.

3.3. Evaluation and Comparison

We conduct comprehensive comparisons with the following approaches: the MLP-type network [7], which incorporates source FC with the target age information by one-hot encoding; the multi-marginal Wasserstein GAN (MWGAN) [21], which defines a generator for the target FC prediction at each age; the multi-GCN based generative adversarial network (MGCN-GAN) [22], which generates FCs from different ages by adjusting the weight of each branch; and the Conditional Intensive Triplet Network (CITN) [20], which uses a similar triplet scheme for prediction. The comparison results for a representative individual are displayed in Fig. 2(a). Notably, in comparison to TC-MAE, the results from other methods lack details and fail to accurately preserve FC features. Table 2 further reinforces these observations by presenting average prediction results across subjects using various methods, highlighting the superior effectiveness of our proposed TC-MAE in predicting of FC development during infancy.

Fig. 2.

(a) FC predictions from TC-MAE and four competing methods are shown. The first column shows the input and ground truth, and the last five columns display the predictions and difference maps, with exemplary parts highlighted in red and yellow in the third row. (b) Comparison of predicted trajectories from various timepoints using different methods, with the source input indicated by the star.

Table 2.

Comparison between TC-MAE and other competing methods (Mean ± Std).

Method	mae	pcc
TC-MAE (proposed)	0.104 ± 0.02*	0.790 ± 0.06*
CITN [20]	0.116 ± 0.03	0.766 ± 0.06
MWGAN [21]	0.111 ± 0.01	0.755 ± 0.01
MGCN-GAN [22]	0.132 ± 0.06	0.710 ± 0.05
MLP [7]	0.148 ± 0.02	0.500 ± 0.01

^*indicates statistically significantly better results than other methods with p-value<0.05.

Furthermore, to compare the methods in preserving longitudinal consistency, we illustrated individual-level functional network trajectories predicted from different timepoints in Fig. 2(b). Each brain region in the Desikan-Killiany atlas was assigned to its corresponding network [23], and the inter- and intra-network FCs were calculated by averaging the FCs related to the corresponding regions. We selected the visual network as the representative to demonstrate the individual trajectories. As shown, compared with other approaches, TC-MAE better captures the dynamic features of longitudinal development and aligns more closely with the actual developmental trajectory.

4. CONCLUSION

This study introduces the Triplet Cycle-Consistent Masked Autoencoder (TC-MAE) for individualized infant FC trajectory prediction. We use Longitudinal Prediction network to reinforce the temporal coherence of different timepoints. We adapted the Masked Autoencoder to learn domain knowledge and holistic information from FC by performing pretraining leveraging abundant adult FC data. Additionally, we employed the Triplet Reconstruction Network to effectively disentangle mixed age-related and ID-related information, enabling individualized prediction. The promising results on the BCP dataset underscore the superiority of our model in predicting FC trajectories during infancy.