Uncovering cognitive taskonomy through transfer learning in masked autoencoder-based fMRI reconstruction

TL;DR

The paper addresses the challenge of reconstructing neural signals and understanding relationships among cognitive tasks from fMRI data. It adopts a masked autoencoder (MAE) framework with brain- and time-masking to reconstruct resting-state and task-based fMRI, and uses a transfer-learning protocol to quantify pairwise task relationships across 23 subtasks, forming a cognitive taskonomy. The study demonstrates robust cross-subject fMRI reconstruction, reveals task taxonomy features such as motor subtask clustering and similarities among emotion, social, and gambling tasks, and shows how the taxonomy can guide selection of source tasks to improve neural decoding on target tasks. This approach provides a data-driven mechanism to address signal loss and to inform transfer learning strategies in neuroimaging applications, with potential extensions to other modalities like EEG.

Abstract

Data reconstruction is a widely used pre-training task to learn the generalized features for many downstream tasks. Although reconstruction tasks have been applied to neural signal completion and denoising, neural signal reconstruction is less studied. Here, we employ the masked autoencoder (MAE) model to reconstruct functional magnetic resonance imaging (fMRI) data, and utilize a transfer learning framework to obtain the cognitive taskonomy, a matrix to quantify the similarity between cognitive tasks. Our experimental results demonstrate that the MAE model effectively captures the temporal dynamics patterns and interactions within the brain regions, enabling robust cross-subject fMRI signal reconstruction. The cognitive taskonomy derived from the transfer learning framework reveals the relationships among cognitive tasks, highlighting subtask correlations within motor tasks and similarities between emotion, social, and gambling tasks. Our study suggests that the fMRI reconstruction with MAE model can uncover the latent representation and the obtained taskonomy offers guidance for selecting source tasks in neural decoding tasks for improving the decoding performance on target tasks.

Figures (7)

• Figure 1: Framework for fMRI reconstruction and transfer learning: (A) fMRI preprocessing, which involves preprocessing fMRI data and parcellating the brain using the MMP atlas. (B) The MAE model for fMRI reconstruction. Masked fMRI data are first processed through patch embedding, then passed through the transformer-based encoder and decoder, and finally linearly projected to reconstruct the fMRI signals. (C) Transfer learning is utilized to estimate the relationships among the source and target tasks. Initially, the MAE model is trained on the source task using masked fMRI data from the source task. In the second step, the MAE model, trained on the source task, is transferred to the target task. During this specific training phase, which utilizes masked fMRI data from the target task, the encoder is frozen to prevent updates, while only the decoder is tuned.
• Figure 2: Results of resting-state fMRI reconstruction using the MAE model. (A) Different masking strategies. Original data, brain-masked data, time-masked data, both brain and time masked data are shown from left to right. (B), (C), and (D) show the reconstruction results of the MAE model with varying mask ratios using the brain mask strategy, the time mask strategy, and the brain and time mask strategy, respectively. (E) Illustration of fMRI reconstruction by the MAE model of the brain and time mask strategy. The original signal, the masked BOLD signal, and the MAE reconstruction results are shown from top to bottom.
• Figure 3: Results of brain mask testing. (A) The MAE reconstruction results for 360 brain regions (sorted according to brain networks). (B) The MAE reconstruction results for each brain network.
• Figure 4: Results of time mask testing. (A) The MAE reconstruction results for masked fMRI data, with each time frame individually masked. (B) The MAE prediction results for fMRI signals at various future time scales. (C) The MAE reconstruction results of different brain networks for masked fMRI data, with each time frame individually masked. (D) The MAE prediction results of different brain networks for fMRI signals at various future time scales.
• Figure 5: Results of fMRI reconstruction across cognitive tasks using the MAE model. The tasks are sorted by the correlation coefficient of the reconstruction results, from highest to lowest. The colors correspond to the cognitive categories to which each subtask belongs.