Graph Contrastive Learning for Connectome Classification

TL;DR

The paper tackles subject-level connectome classification under limited labeled data by proposing a supervised graph contrastive learning framework built on a multi-task Encoder-Decoder Graph Neural Network that jointly models structural and functional connectivity. It leverages two graph augmentations to create multiple views and employs a two-stage training process: a contrastive pre-training phase followed by fine-tuning a classifier, with a reconstruction objective $\mathcal{L}_{\mathrm{MSE}}(\hat{\Sigma},\Sigma)$ and a classification objective $\mathcal{L}_{\mathrm{CE}}(\hat{y},y)$ balanced by $\lambda$, where the FC is reconstructed as $\hat{\Sigma} = \text{ReLU}(\mathbf{X}_C \mathbf{X}_C^{\top})$. The approach achieves state-of-the-art gender classification on the Human Connectome Project dataset and demonstrates robustness to reduced training data, highlighting its potential for neuroscience applications and precision medicine. The work also emphasizes the importance of decoder-enabled regularization and augmentation strategies in learning informative connectome representations for downstream tasks.

Abstract

With recent advancements in non-invasive techniques for measuring brain activity, such as magnetic resonance imaging (MRI), the study of structural and functional brain networks through graph signal processing (GSP) has gained notable prominence. GSP stands as a key tool in unraveling the interplay between the brain's function and structure, enabling the analysis of graphs defined by the connections between regions of interest -- referred to as connectomes in this context. Our work represents a further step in this direction by exploring supervised contrastive learning methods within the realm of graph representation learning. The main objective of this approach is to generate subject-level (i.e., graph-level) vector representations that bring together subjects sharing the same label while separating those with different labels. These connectome embeddings are derived from a graph neural network Encoder-Decoder architecture, which jointly considers structural and functional connectivity. By leveraging data augmentation techniques, the proposed framework achieves state-of-the-art performance in a gender classification task using Human Connectome Project data. More broadly, our connectome-centric methodological advances support the promising prospect of using GSP to discover more about brain function, with potential impact to understanding heterogeneity in the neurodegeneration for precision medicine and diagnosis.

Figures (4)

• Figure 1: The Encoder-Decoder architecture used in this work processes structural connectivity matrix $\mathbf{A}$ as input, along with nodal attributes $\mathbf{X}^{(0)} = \mathbf{I}_{N}$. The GCN encoder performs graph convolutions and associated information propagation through $L$ layers, producing low-dimensional node embeddings $\mathbf{X}^{(l)}$ at layer $l$. Intermediate representations from all layers are concatenated into one final embedding $\mathbf{X}_C$. The decoder uses an outer-product operation to predict the reconstructed functional connectivity matrix $\hat{\mathbf{\Sigma}}:=\text{ReLU}(\mathbf{X}_C\mathbf{X}_C^\top)$, effectively modeling the SC-FC mapping. Graph-level representations $\mathbf{z}$ are obtained via a pooling operation applied to $\mathbf{X}_C$, summarizing the node embeddings into a single vector, from which a logistic regression classifier predicts the binary label $\hat{y}$.
• Figure 2: A schematic view of the constrastive learning approach. For each input datum $\{\mathbf{A}_k,y_k\}$ we generate two augmented versions $\{\mathbf{A}'_{2k},y_k\}$ and $\{\mathbf{A}'_{2k+1},y_k\}$ using edge dropping and attribute masking. The corresponding graph embeddings $\mathbf{z}_i$ are produced as in Fig. \ref{['fig:encoder-decoder']}. Whereas unsupervised CL only pulls together the representations stemming from the same original graph (represented by green arrows), supervised CL optimizes the contrastive loss $\mathcal{L}^{\mathrm{sup}}$ in \ref{['eq:l_sup_loss']}, which also takes into account the class labels of the subjects. This approach attracts augmented versions originating from subjects with the same label and repels those associated with different ones (represented by green but also red arrows).
• Figure 3: Evolution of the learned embeddings in the validation set for the contrastive learning setup and the Encoder-Decoder architecture, shown after the pre-training step. The first two PCA principal components are visualized. At epoch 500, in the first principal component, a clear separation between embeddings of different genders can be observed. During the fine-tuning step, these learned embeddings are utilized to classify with a logistic regression classifier.
• Figure 4: Comparison of contrastive learning (CL) and baseline classification (Baseline), for encoder and Encoder-Decoder architectures with or without data augmentation (DA), under varying training set proportions. The CL approach consistently outperforms the others and reaches near-optimal results with as little as half of the data, whereas the other frameworks require the full dataset to achieve near-peak performance.