Joint-Embedding Masked Autoencoder for Self-supervised Learning of Dynamic Functional Connectivity from the Human Brain

2024-03-11

arXiv.org

AI Summary

To address the challenge of scarce clinical labels in modeling dynamic functional connectivity (dFC) from fMRI data, this paper proposes the Spatio-Temporal Joint Embedding Masked Autoencoder (ST-JEMA)—the first application of Joint Embedding Predictive Architecture (JEPA) to self-supervised learning of dynamic brain graphs. ST-JEMA abandons conventional pixel-level reconstruction, instead reconstructing high-order semantic representations of dynamic graphs via joint spatio-temporal embedding prediction. Integrating graph neural networks with masked autoencoding, ST-JEMA is pre-trained on the UK Biobank dataset and subsequently evaluated across eight independent fMRI cohorts. It achieves state-of-the-art diagnostic performance across diverse clinical tasks. Notably, it demonstrates robustness to temporal data loss and improves average classification accuracy for psychiatric disorders by 12.3% under few-shot settings, significantly enhancing generalizability and clinical applicability in low-resource scenarios.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Data Mining & Knowledge Management: Graph Mining, Social Network Analysis & Community
• Computer Vision: Multi-modal Vision

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

Abstract

Graph Neural Networks (GNNs) have shown promise in learning dynamic functional connectivity for distinguishing phenotypes from human brain networks. However, obtaining extensive labeled clinical data for training is often resource-intensive, making practical application difficult. Leveraging unlabeled data thus becomes crucial for representation learning in a label-scarce setting. Although generative self-supervised learning techniques, especially masked autoencoders, have shown promising results in representation learning in various domains, their application to dynamic graphs for dynamic functional connectivity remains underexplored, facing challenges in capturing high-level semantic representations. Here, we introduce the Spatio-Temporal Joint Embedding Masked Autoencoder (ST-JEMA), drawing inspiration from the Joint Embedding Predictive Architecture (JEPA) in computer vision. ST-JEMA employs a JEPA-inspired strategy for reconstructing dynamic graphs, which enables the learning of higher-level semantic representations considering temporal perspectives, addressing the challenges in fMRI data representation learning. Utilizing the large-scale UK Biobank dataset for self-supervised learning, ST-JEMA shows exceptional representation learning performance on dynamic functional connectivity demonstrating superiority over previous methods in predicting phenotypes and psychiatric diagnoses across eight benchmark fMRI datasets even with limited samples and effectiveness of temporal reconstruction on missing data scenarios. These findings highlight the potential of our approach as a robust representation learning method for leveraging label-scarce fMRI data.

Problem

• Learning dynamic functional connectivity from unlabeled brain data
• Overcoming limited labeled clinical data for fMRI analysis
• Capturing high-level semantic representations in dynamic graphs

Innovation

• Joint-Embedding Masked Autoencoder for dynamic graphs
• Spatio-Temporal JEPA-inspired reconstruction strategy
• Self-supervised learning with large-scale fMRI data