GATE: Graph CCA for Temporal SElf-supervised Learning for Label-efficient fMRI Analysis

Medical imaging techniques such as Computed Tomography (CT), Magnetic Resonance Imaging (MRI), Positron Emission Tomography (PET), and X-ray have been employed for diagnosis and early detection of diseases. For example, Kam et al. utilized multiple convolutional neural networks to conduct early mild cognitive impairment diagnosis [26]. In neuroscience, MRI is widely used to study neurological disorders, which can be separated between structural MRI and functional MRI depending on the detected type of connectivity. In structural MRI, nodes are defined by anatomical connections between regions or brain tissues, while edges are constructed by the topology between them. For instance, Yao et al. presented a triplet GNN model for disease diagnosis on structural MRI data [27]. In functional MRI (fMRI), nodes represent functional regions of the brain, while edges are constructed by correlations between their activities. For example, Wang et al. proposed a similarity-driven multi-view model to conduct autism spectrum disease diagnosis on fMRI data [28]. In fact, it has been a prevalent way to exploit the functional organization of human brain by analyzing functional fMRI, as shifts in human attention (e.g., a stimulus or cognitive task) are associated with systematic changes in the activity of functional areas of brains [29, 30].

There is growing evidence that functional connectivity may show dynamic changes over a short period of time  [31, 32, 33]. Previous works have demonstrated that the using of sliding windows is beneficial to conduct dynamic connectivity analysis on fMRI data since the resting-state reflects different physiological states, such as wandering and muscle fatigue. For example, Wang et al. first used sliding windows to divide the rs-fMRI time series into multiple segments and then proposed a spatio-temporal convolutional-recursive neural network [34]. Mao et al. analyzed the dynamic functional differences between men and women based on sliding window measures and community detection [32]. Yao et al. proposed a temporal adaptive graph convolutional network to exploit the topological information and to model multi-level semantic information within the entire time series signals [35]. Yu et al. [36] presented a Transformer-based method to infer functional networks in both spatial and temporal space in a self-supervised manner.

Deep learning has revolutionized computer-aided diagnosis in the past decade [37, 38]. On this basis, deep GCNs takes the advantages of the relationship between nodes to find the common patterns/biomarkers, and there has been an increasing focus on GCNs in neuroscience [3]. Previous GCN-based methods on fMRI data can be categorized into two subgroups depending on the definition of nodes in the graph, i.e., population graph-based models and brain region graph based models [8]. Also, these two subgroups correspond to two separate tasks for fMRI analysis, i.e., population graph-based models for node classification task and brain region graph based models for graph classification tasks.

In population graph-based models, the nodes in the graph denote the subjects and edges represent the similarity between subjects. For example, Parisot  et al.  [5] exploits the GCN and involves representing populations as a sparse graph in which nodes are associated with imaging features, and edge weights are constructed from phenotype information (e.g., age, gene, and sex of the subjects). Kazi  et al.  [39] propose a GCN model with multiple filter kernel sizes on a population graph.

In brain region graph based models, the nodes denote anatomical brain regions and the edges represent functional or structural connectives among these brain regions. For example, Li  et al.  [3] conduct an interpretable GCN model on brain region graph to understand which brain regions are related to a specific neurological disorder. Xing  et al.   [40] consider a GCN model with long short term memory model on brain region graph for disease diagnosis.

Although previous GCN methods achieved promising results on both population graph and brain region graph, existing methods are still limited by large labeled data. Clearly, designing a GCN method for limited labeled fMRI data is an important yet unsolved challenge.

As an powerful unsupervised representation learning method, SSL has recently been developed in graph learning domain. Existing SSL methods can be broadly classified into three groups, i.e., contrastive-based SSL, reconstruction-based SSL and similarity-based SSL.

The contrastive-based SSL methods enhance the similarity between the representations from two views (e.g., global representation and local representation) by manually constructing positive and negative sample pairs. For example, Deep Graph Infomax (DGI) [41] employs InforNCE [14] contrastive loss function to contrast the local node representations and the global graph representations. MVGRL [42] proposes a contrastive multi-view representation learning method by contrasting local and global embeddings from two views. However, contrastive-based SSL methods rely on negative samples, which is not suitable with limited number of samples and small number of classes.

Recently, reconstruction-based SSL methods learn a representation by conducting reconstruction-based pretext tasks (e.g., image inpainting and recovering color channels). For example, He  et al.  [16] develop an SSL encoder-decoder model to reconstruct the original image from the latent representation and mask patches. Qiu  et al.   [17] propose a graph encoder-decoder model to reconstruct the relationships between nodes based on representation similarity. However, the reconstruction-based SSL needs to reconstruct input context (low dimensional features are transformed into high dimensional features), which is impractical due to fitting the low signal-to-noise ratio fMRI features may result in overfitting to spurious features.

More recently, similarity-based SSL takes the advantage of the coupling between multiple views of the same data (i.e., input context and its augmentation) to learn a good representation without labeled data. For example, Grill  et al.  [18] learn representations by encoding two augmented views through two distinct graph encoders. Although this technique avoids the limitations of selecting negative samples in the contrasitive-based SSL method and reconstruction loss in the reconstruction-based SSL method, designing a suitable strategy for generation multiple coupled views for fMRI data is highly needed.

The key idea of GATE is to ensure the consistency of two augmented views of an input content. Specifically, by adopting similarity-based SSL, spurious factors can be mitigated by augmenting its variations specifically. For example, if we train a model to identify “apple” from peach, these two views can be “red apple” and “yellow apple”. Both views contain the main characteristics (i.e., shape) of “apple” but vary on the spurious features (i.e., color).

Existing literature indicates that dynamic FC-based fMRI analysis is sensitive to the choice of the size of the sliding window $L$ [44, 45]. Therefore, we aim to impair the effect of perturbations on $L$, i.e., the spurious features in dynamic FC. Although the brain FC is not static, a short time sub-segment and its surrounding time sub-segments should contain the same characteristics (e.g., FC patterns) associated with neurological disease [46, 31]. Thus, we propose to augment fMRI data along its temporal domain with sliding windows to obtain two relevant views ($G^{a}=\{\mathbf{X}^{a},\mathbf{A}^{a}\}$ and $G^{b}=\{\mathbf{X}^{b},\mathbf{A}^{b}\}$) by step window augmentation (S-A) and multi-scale window augmentation (M-A).

Step window augmentation (S-A). As shown in Fig. 2 ($a$), S-A takes two neighboring sliding windows as two related views ($G^{a}$ and $G^{b}$). In this case, the raw BOLD signals $\mathbf{S}_{i}$ are divided into $M$ sub-segments $\{\mathbf{S}_{i}^{1},\dots,\mathbf{S}_{i}^{M}\}$ for the $i$-th subject, by setting the window size as $L$ and the step of the sliding window as $s$, resulting in $M=\lfloor\frac{T-L}{s}\rfloor+1$ sub-segments and the corresponding dynamic FCs. In each training iteration, S-A first randomly select one sub-segment (e.g., $m$-th sub-segment) for the first view $G^{a}=\{\mathbf{X}^{m},\mathbf{A}^{m}\}$, and then regard its neighboring sub-segment (e.g., $(m\pm 1)$-th sub-segment) as the second view $G^{b}=\{\mathbf{X}^{m\pm 1},\mathbf{A}^{m\pm 1}\}$.

Multi-scale window augmentation (M-A). Considering that FCs generated by different window sizes contain the relevant information, M-A takes two different scales of sliding windows as two related views. In this case, the $\mathbf{S}_{i}$ is divided into multiple $M$ sub-segments depending on the window size $L$. For the $m$-th sub-segment ($m\in[M]$), we sample BOLD signals using two different window sizes $l_{a}$ and $l_{b}$ to calculate FCs, forming two views of graph $G^{a}=\{\mathbf{X}^{m,l_{a}},\mathbf{A}^{m,l_{a}}\}$ and $G^{b}=\{\mathbf{X}^{m,l_{b}},\mathbf{A}^{m,l_{b}}\}$. These two views are fed into each training iteration of GATE, as shown in Fig. 2 ($a$).

It is worth mentioning that it is also possible to jointly consider S-A and M-A, for example, by randomly selecting one in each training iteration. In addition, S-A and M-A can also be easily combined with randomly drop augmentation [47] which is commonly used (see Sec. V-C). In fact, data augmentation plays an essential role in SSL, however not all data augmentations have a positive effect on SSL [22]. Hence, it is important to investigate a suitable data augmentation for SSL considering the essential role of data augmentations in SSL as well as the special properties and data structure of medical data compared with natural data. Compared with [48], our method directly operates on the raw data rather than the graph representation. Furthermore, our proposed augmentation methods for fMRI to generate multi-view data is a versatile solution that can be combined with any model architectures.

With the graph data from two views ($G^{a}=\{\mathbf{X}^{a},\mathbf{A}^{a}\}$ and $G^{b}=\{\mathbf{X}^{b},\mathbf{A}^{b}\}$), we perform an encoder to learn the patterns of the data. Since GCN [49] and its variants have been widely used to capture the semantic information (i.e., $\mathbf{X}$) and the structural information (i.e., $\mathbf{A}$) in the graph data, in this study, we adopt a GCN model as the encoder $f(\cdot)$ to get the embedding of all subjects in each view. It is worth noting that the graph encoder $f(\cdot)$ is versatile for applying other powerful graph learning encoders (e.g., GAT [50] and GIN [51]). For fair comparison and easily application, we apply the commonly used GCN model as the encoder in this study. More specifically, the GCN operation on the $l$-th hidden layer is defined as

where $\mathbf{D}$ is the diagonal matrix of $\mathbf{A}$, $\mathbf{H}^{(l)}$ is the output features of all subjects at the $l$-th layer, $\mathbf{\Theta}^{(l)}$ is the trained weight matrix, and $\sigma(\cdot)$ represents an activation function. The input features of Eq. (1) is the feature vector fMRI features (i.e., $\mathbf{X}$ obtained by S-A or M-A). To make the embedding comparable, we use the same GCN encoder (weight sharing) to project two views into the same embedding space. Consequently, we obtain the normalized embedding matrices of two views, i.e., $\mathbf{Z}^{a}=f(\mathbf{X}^{a},\mathbf{A}^{a})$ and $\mathbf{Z}^{b}=f(\mathbf{X}^{b},\mathbf{A}^{b})$.

To optimize the model, contrastive-based SSL methods [41, 42] apply contrastive loss on positive pairs and negative pairs, which is not suitable with limited number of samples and small number of classes. Reconstruction-based SSL methods [16] apply reconstruction loss (i.e., MSE) to reconstruct input context, which would be impractical due to fitting the low signal-to-noise ratio fMRI features may result in overfitting to spurious features. Therefore, we propose GATE to avoid sampling negative samples or reconstructing fMRI time series. Instead, we design novel augmentation methods together with the corresponding similarity loss to leverage unlabeled data, while prevent feature collapse in the embedding space. Once the embedding matrices of two views are computed, the next step is to maximize the correlation of these two embedding matrices, as the CCA does (see Sec. IV). Furthermore, we define an input-consistency regularization loss as

where $\left\langle\cdot,\cdot\right\rangle$ is the dot product operator, and $\gamma$ is the trade-off coefficient. And $v$ is one of the view (e.g., view $a$ and view $b$) and $Z^{v}$ is the embeddings matrices of this view. In Eq. (2), the first term can be regarded as a regularization of the embedded features, ensuring that the low-dimensional features can maintain the representation capability. The second term in Eq. (2) ensures that the individual dimensions of the feature are uncorrelated to avoid collapsed solutions, i.e., all outputs of the model are equal.

Performing downstream task. After the GCN encoder is trained, there is no need to keep the large graph structure in fine-tuning step for versatile and efficiency in clinical practice because many GNN encoders (e.g., GCN and GAT) belong to transductive learning which means the $\mathbf{A}$ needs to be reconstructed after new samples involved. Thus, we fine-tune the encoder without graph information (i.e., by replacing $\mathbf{A}$ with identity matrix $\mathbf{I}$), followed by a linear layer with ELU activation function (denoted as $\psi(\cdot):\mathcal{Z}\mapsto\mathbb{R}$) on a small portion of labeled data using cross-entropy loss for a particular clinical prediction task. All sliding windows are entered into the model during inference processes. To this end, we give the pseudo code of GATE.

The quantitative comparison is shown in Table I. We can observe that: GATE achieves the best performance on the quantitative metrics across two datasets, followed by CCA-SSG, SAC-GCN, GAT, BGRL, Vanilla GCN, MVGRL, and DGI. Meanwhile, we find that the improvements are significant (with $p<0.05$ via significance testing) compared with state-of-the-art graph SSL methods. Particularly, GATE improves by $12.73\%$ ($p=4.83e^{-4}$) and $4.04\%$ ($p=4.16e^{-3}$) on average, compared to the baseline comparison method DGI and the best graph SSL comparison method CCA-SSG. This indicates the superior performance of our proposed GATE. Moreover, we can observed that the DGI and MVGRL cannot outperform the baseline vanilla GCN. The possible reason is that the limited number of samples and a small number of classes in the dataset (the problem of class collision [62]). We emphasize that the ultimate success of SSL is leveraging unlabeled data, which means SSL methods can gain more benefits and achieve higher results with more unlabeled data.

We further conduct experiments under different ratios of labeled data ($10\%$ to $80\%$). Remarkably, labeled data is only used in the fine-tuning stage under the SSL setting. As shown in Fig. 3, GATE always outperforms the vanilla GCN, especially with a large gap in performance on limited labeled data. Specifically, compared to vanilla GCN, GATE achieves an average improvement of $11.4\%$, $9.9\%$, $4.3\%$, and $2.0\%$ in terms of $10\%$, $20\%$, $50\%$, and $80\%$ percentage of labeled data, respectively. These results confirm our claim that GATE obtains satisfying prediction performance with a small portion of labeled data.

First, in the augmentation S-A, it includes two parameters, i.e., the length of the sliding windows $L$ and the step of sliding windows $s$ which control the total number of the sliding windows. Figure 10 shows the results obtained by traversing $L$ from $20$ to $70$ and $L$ from $10$ to $35$, with fixing the other settings. We could observe that GATE is insensitive to the parameters $L$ within a certain range of values and $s$. Specifically, GATE can consistently obtain promising results except in the case of a tiny value of $L$, and the suggested value ranges of $L$ and $s$ are $[30,60]$ and $[10,35]$, respectively. Second, in the augmentation M-A, we investigate the effect of gaps between two different scale sliding windows (i.e., $\left\|l_{a}-l_{b}\right\|$). Figure 11 shows the results obtained by traversing gaps from $5$ to $40$ with fixing the other settings (e.g., the value of $L$ and the value of $s$ are set to $30$ and $15$, respectively). As a result, the variation of the performance of GATE is relatively stable with the value ranges of gaps is $[10,40]$. This observation demonstrates that GATE is insensitive to the parameters settings of M-A in our experiments. In summary, the reasons of the above observation could be that GATE maximizing the similarity of the two representations (i.e., embeddings of two different sliding windows) to reduce the perturbations on different sliding windows.

In recent years, fMRI analysis has revolutionized with deep learning, while training a powerful model usually requires a large amount of labeled fMRI data. To visually explain the needs of large labeled data for fMRI analysis, we depict the results of vanilla GCN under different label rates in Fig. 3. We can observe that the accuracy of the vanilla GCN decreases significantly with the decreasing labeling rate (e.g., 10% to 50%). However obtain a sufficient number of labeled images for fMRI data could be challenging due to the cost of obtaining neuro-disease annotations. Therefore, designing a method for limited labeled fMRI data is crucial [63]. To solve this limitation, we propose a graph SSL method GATE, which can yield comparable results under label-efficient data (e.g., the accuracy of GATE under 20% labeling rate is comparable to the accuracy of vanilla GCN under 50% labeling rate). The main reason could be that GATE can help extract robust features on label-efficient data. Consequently, our experiments show promising results to tackle the issue of limited labeled fMRI data.

Deep graph learning models, like GCNs, have become popular approaches for populations disease analysis, including neuroimaging [5, 3]. Modeling subject-wise neuroimaging signals (and auxiliary phenotype features) could leverage the representation of similar subjects. Our work investigates how adding similarity regularization affects GCN-based representation learning on fMRI signals. As shown in Fig. 8 and Tab. I, graph information is essential to get good results and GATE is significantly better than other GCN-based methods. The reason could be better extracting associations between subjects and maximizing the correlation of the representations from multiple coupled views. Hence, we can conclude that graph learning can obtain embedding matrices of different views to characterize the representations of neuroimaging.

We further analyze the technical contributions of GATE. Firstly, GATE utilizes an SSL-based strategy to capture vital relationship for generation multiple coupled views of a fMRI BOLD signal. We can achieve good performance of pre-training on unlabeled fMRI data and fine-tuning on small labeled data. Then, we design a GCN encoder for limited labeled fMRI data to extract associations between subjects. Our GCN-based method effectively solves the issue of earning on label-efficient dataset. Moreover, GATE is a versatile method which can be further improved by considering a more powerful GNN encoder (i.e., $f(\cdot)$) and combining a more powerful classifier (i.e., $\psi(\cdot)$) for real application. In general, self-supervised tasks provide informative priors that can benefit GCN in generalizable target performance.