[Paper Review] A Generative Self-Supervised Framework using Functional Connectivity in fMRI Data

Jung‐Won Choi, Seongho Keum+3 more|arXiv (Cornell University)|Dec 4, 2023

Functional Brain Connectivity Studies4 citations

TL;DR

This paper proposes ST-MAE, a generative self-supervised learning framework for dynamic functional connectivity (fc) networks in fMRI data, which reconstructs masked node features using spatio-temporal representations across time. It achieves state-of-the-art performance on downstream tasks—including psychiatric disorder classification—using only unlabeled fMRI data, demonstrating robust representation learning even with limited labeled data.

ABSTRACT

Deep neural networks trained on Functional Connectivity (FC) networks extracted from functional Magnetic Resonance Imaging (fMRI) data have gained popularity due to the increasing availability of data and advances in model architectures, including Graph Neural Network (GNN). Recent research on the application of GNN to FC suggests that exploiting the time-varying properties of the FC could significantly improve the accuracy and interpretability of the model prediction. However, the high cost of acquiring high-quality fMRI data and corresponding phenotypic labels poses a hurdle to their application in real-world settings, such that a model naïvely trained in a supervised fashion can suffer from insufficient performance or a lack of generalization on a small number of data. In addition, most Self-Supervised Learning (SSL) approaches for GNNs to date adopt a contrastive strategy, which tends to lose appropriate semantic information when the graph structure is perturbed or does not leverage both spatial and temporal information simultaneously. In light of these challenges, we propose a generative SSL approach that is tailored to effectively harness spatio-temporal information within dynamic FC. Our empirical results, experimented with large-scale (>50,000) fMRI datasets, demonstrate that our approach learns valuable representations and enables the construction of accurate and robust models when fine-tuned for downstream tasks.

Motivation & Objective

• To address the challenge of limited labeled fMRI data in psychiatric and neurological research.
• To overcome the limitations of existing contrastive self-supervised learning methods in dynamic graphs, which often lose semantic information under perturbation.
• To develop a generative self-supervised framework that jointly models spatial and temporal dynamics in functional connectivity networks.
• To enable effective pre-training on large-scale unlabeled fMRI datasets for downstream tasks with minimal labeled data.
• To improve the accuracy and robustness of models in clinical fMRI applications, especially for psychiatric disorder classification.

Proposed method

• The framework constructs dynamic functional connectivity graphs from fMRI time series using regions of interest (ROIs), forming a sequence of graphs over time.
• Node features are enhanced with learnable spatial encodings and time-encoding vectors generated via a Gated Recurrent Unit (GRU) to model temporal dynamics.
• A masked autoencoder strategy is applied where node features are randomly masked at intermediate time points, and the model reconstructs them using representations from surrounding time steps.
• The model uses a graph neural network encoder to learn contextualized node representations from masked inputs, followed by a decoder to reconstruct original features.
• The loss function minimizes the reconstruction error between predicted and original node features, enabling end-to-end training via backpropagation.
• Temporal information is preserved by conditioning the reconstruction on representations from adjacent time points, enabling the model to learn temporal dependencies.

Experimental results

Research questions

• RQ1 Can a generative self-supervised learning framework effectively capture spatio-temporal dynamics in dynamic functional connectivity networks from fMRI data?
• RQ2 How does the proposed ST-MAE framework compare to contrastive SSL methods in terms of representation quality and downstream task performance?
• RQ3 To what extent can pre-training on large-scale unlabeled fMRI data improve classification accuracy in low-data regimes, such as psychiatric disorder diagnosis?
• RQ4 Does the integration of time-encoding vectors enhance the model's ability to reconstruct masked features and learn meaningful temporal patterns?
• RQ5 Can the framework generalize across diverse fMRI datasets with varying acquisition protocols and phenotypic labels?

Key findings

• ST-MAE achieves state-of-the-art performance on multiple downstream fMRI tasks, including gender classification and psychiatric disorder classification, even with limited labeled data.
• The model demonstrates superior robustness and generalization compared to contrastive SSL baselines, particularly in handling distribution shifts and noisy perturbations.
• On the UK Biobank dataset (n ≈ 40,000), ST-MAE pre-training significantly improves fine-tuning accuracy across benchmark datasets such as ABCD, HCP, ABIDE, and ADHD200.
• The framework shows notable gains in psychiatric disorder classification, highlighting its utility in clinical applications with scarce labeled data.
• The ablation study confirms that both spatial and temporal encodings are essential, with the removal of time encoding leading to a significant drop in performance.
• The generative reconstruction objective outperforms contrastive objectives in capturing meaningful semantic structure in dynamic fc graphs, especially under graph perturbations.