Learning Robust and Task-Invariant Functional Representation from fMRI through Siamese Self-Supervised Learning

Technical Summary: Learning Robust and Task-Invariant Functional Representation from fMRI through Siamese Self-Supervised Learning

Problem Statement

Functional magnetic resonance imaging (fMRI) is a critical tool for investigating brain function, yet its application in machine learning is hindered by high data acquisition costs, small sample sizes, and the inherent noise of psychiatric rating scales. These factors, combined with the high dimensionality of fMRI data, create a significant risk of model overfitting in supervised learning settings. Furthermore, existing supervised frameworks often rely on binary classification labels that collapse diverse disease subtypes, obscuring meaningful differences and reducing interpretability. While large-scale foundation models offer a potential solution, their computational requirements for pretraining and fine-tuning are often prohibitive. Additionally, defining generalizable "negative" examples for contrastive learning in fMRI is challenging, as treating different subjects as negatives may discard shared functional patterns, while using task labels to construct negatives can entangle features with specific task spaces, limiting cross-task generalization.

Methodology: BrainSimSiam

The authors introduce BrainSimSiam, a lightweight, data-efficient self-supervised representation learning framework designed to learn robust and task-invariant functional representations from task-based fMRI. The framework utilizes a Siamese architecture with a stop-gradient mechanism, eliminating the need for negative sample pairs.

Key Architectural Components:

1. Dual-View Input: The model processes two complementary views of the fMRI data for each instance:
   • Functional Brain Graph View: The fMRI series is parcellated into Regions of Interest (ROIs) using standard atlases (Shen for HCP, Desikan–Killiany for Biopoint). Node features are derived from Pearson correlation matrices, while edges represent the strongest partial correlation coefficients.
   • Spatial Activation Image View: A 3D mean activation map is generated by voxel-wise time-averaging of the fMRI time series, preserving absolute signal intensity.
2. Encoder Architecture: The encoder consists of two parallel branches:
   • A GNN branch using Graph Attention Network (GAT) layers with global average and max pooling.
   • A CNN branch utilizing 3D convolutional layers.
   • Outputs from both branches are concatenated and passed through a projection MLP to produce the final embedding.
3. Augmentation and Masking:
   • ROI-Aligned Graph-Image Masking: A novel operator that jointly occludes voxels in the 3D image corresponding to dropped nodes in the graph. This ensures consistent perturbation across both views.
   • Task Invariance Loss: To learn representations that generalize across tasks, the framework introduces a loss term that encourages the model to preserve spontaneous, task-invariant functional patterns. This is achieved by sampling a scan from the same subject under a different task and minimizing the cosine similarity loss between the original view and this cross-task view.
4. Training Objective: The model minimizes a symmetric cosine-similarity loss between the outputs of the encoder and a predictor head, utilizing a stop-gradient operation on one branch to prevent model collapse without negative samples.

Key Contributions

1. Framework Development: The authors developed BrainSimSiam, a self-supervised framework for task-based fMRI that outperforms both supervised and existing self-supervised baselines on multiple datasets.
2. Unified Multimodal View: The work unifies voxel-wise fMRI and graph-based functional views via a joint ROI masking scheme. This mechanism supports potential multimodal fusion by aligning graph-structured networks with registered structural (T1-weighted) or diffusion MRI.
3. Task-Invariant Learning: By leveraging contrasts across diverse task stimuli as natural augmentations, the framework learns to disentangle evoked responses from spontaneous brain dynamics, resulting in task-invariant representations.

Experimental Results

The framework was evaluated on two datasets: the Human Connectome Project (HCP) (1,005 subjects, 7 task paradigms) and the Biopoint dataset (75 ASD subjects, 43 healthy controls).

• Performance: BrainSimSiam consistently outperformed supervised baselines (GCN, GAT, GraphSAGE with CNN) and self-supervised baselines (CGL, CGL-CNN) across classification (biological sex, ASD) and regression (cognitive scores, age) tasks.
  • On the HCP dataset, the end-to-end fine-tuned BrainSimSiam model achieved 91.9% accuracy and 0.980 AUC for biological sex classification, competitive with larger foundation models like NeuroSTORM but with significantly fewer parameters (52 million).
  • On the smaller Biopoint dataset, the benefits of self-supervision were amplified. Notably, MLP probing (freezing the pretrained encoder) surpassed end-to-end fine-tuning, indicating that the learned representations are robust and mitigate overfitting in data-limited scenarios.
• Ablation Studies: Removing the projection MLP, the task invariance loss, or either the GNN/CNN branch significantly degraded performance, confirming the necessity of non-linear integration of graph and image features and the value of task-invariant constraints.
• Qualitative Analysis: Zero-shot correlation analysis showed that features from the learned embeddings had higher correlations with biological and behavioral phenotypes (e.g., age, sex, cognitive scores) than traditional functional connectivity (FC) metrics.
• Interpretability: Using GNNExplainer, the authors demonstrated that while pretraining learns from distributed whole-brain signals, downstream task-specific importance maps become localized to salient regions (e.g., right thalamus for sex classification, superior frontal gyrus for ASD), aligning with prior neuroscientific literature.

Significance and Claims

The paper claims that dataset-specific self-supervised pretraining yields robust, generalizable representations for fMRI, enabling accurate prediction of heterogeneous phenotypes even with limited data. Unlike large-scale foundation models that require massive computational resources and harmonization across diverse public datasets, BrainSimSiam demonstrates that a lightweight, data-efficient approach can approach the performance of large-scale models while outperforming purely supervised baselines. The authors emphasize that their method effectively addresses the challenges of label noise, small sample sizes, and the need for task-invariant features in neuroimaging, offering a practical pathway for data-limited applications.