Large Connectome Model: An fMRI Foundation Model of Brain Connectomes Empowered by Brain-Environment Interaction in Multitask Learning Landscape

TL;DR

The paper tackles the challenge of data-scarce brain imaging tasks by introducing the Large Connectome Model (LCM), a decoder-only Transformer trained with multitask learning over brain-environment interactions (BEIs) using functional connectomes. By tokenizing BEIs and leveraging cross-attention to FC matrices, LCM achieves strong performance on sex prediction, behavior recognition, and early diagnosis across eight fMRI datasets, with notable gains when pretraining includes diverse BEIs. The authors demonstrate scalable improvements with model depth, effective few-shot generalization, and interpretable attention patterns that align with known neural networks. This approach promisingly advances neuroimaging foundation models toward clinically applicable, robust, and scalable brain analysis tools, with public data and pretrained weights provided for reproducibility.

Abstract

A reliable foundation model of functional neuroimages is critical to promote clinical applications where the performance of current AI models is significantly impeded by a limited sample size. To that end, tremendous efforts have been made to pretraining large models on extensive unlabeled fMRI data using scalable self-supervised learning. Since self-supervision is not necessarily aligned with the brain-to-outcome relationship, most foundation models are suboptimal to the downstream task, such as predicting disease outcomes. By capitalizing on rich environmental variables and demographic data along with an unprecedented amount of functional neuroimages, we form the brain modeling as a multitask learning and present a scalable model architecture for (i) multitask pretraining by tokenizing multiple brain-environment interactions (BEI) and (ii) semi-supervised finetuning by assigning pseudo-labels of pretrained BEI. We have evaluated our foundation model on a variety of applications, including sex prediction, human behavior recognition, and disease early diagnosis of Autism, Parkinson's disease, Alzheimer's disease, and {Schizophrenia}, where promising results indicate the great potential to facilitate current neuroimaging applications in clinical routines.

Figures (9)

• Figure 1: Learning strategies of previous brain foundation models and LCM. The pretraining in (a) previous brain foundation models is a reconstructive representation learning based on the information (info.) bottleneck, while (b) in LCM it is a multitask learning for multiple brain-environment interactions (BEI) token embeddings by a Transformer decoder, where MHSA is multi-head self-attention, MHCA is multi-head cross-attention, and FFN is a feedforward network. The finetuning in (c) previous studies is training a relatively small head, e.g., a multilayer perceptron (MLP), for the downstream task. (d) LCM finetunes the BEI tokens along with new tokens representing the downstream task.
• Figure 2: Scalability is demonstrated by model size vs. training loss, where the training is supervised by arbitrary non-brain-imaging phenotypes as BEIs in our multitask learning.
• Figure 3: Our LCM surpasses other foundation models, demonstrating outstanding efficiency, on our biggest downstream application, ABIDE ($n$=1,025), as an example. Even the smallest LCM (147M), achieves comparable performance while being significantly efficient in both parameters and resource usage.
• Figure 4: The architecture of one layer of the LCM.
• Figure 5: Pre-training and finetuning of LCM use a two-stage learning strategy: (1) Getting momentum by computing loss for all layers, and (2) adaptive training for the best layer. Note that ground truth (GT) could be a pseudo-label of the BEI, e.g., subjects are healthy by default in HCP datasets.