TARDRL: Task-Aware Reconstruction for Dynamic Representation Learning of fMRI
Abstract
Existing studies in fMRI analysis leverage mask autoencoder, a self-supervised framework, to build model to learn representations and conduct prediction for various fMRI-related tasks. It involves pretraining the model by reconstructing signals of brain regions that are randomly masked at different time segments and subsequently fine-tuning it for prediction tasks. Though it has shown improved performance in prediction tasks, we argue that directly applying this framework on fMRI data may result in sub-optimal results. Firstly, random masking is ineffective for highly redundant fMRI data. Secondly, the reconstruction process is not task-aware, ignoring a critical phenomenon: the varying contributions of different brain regions to different prediction tasks. In this work, we propose and demonstrate a hypothesis that learning representations by reconstructing signals from important ROIs at different time segments can enhance prediction performance. Specifically, we introduce a novel learning framework, Task-Aware Reconstruction Dynamic Representation Learning (TARDRL), to improve prediction performance through task-aware reconstruction. Our approach incorporates an attention-guided masking strategy, which leverages attention maps from the prediction process to guide signal masking during reconstruction, making the reconstruction task task-aware. Extensive experiments show that our model outperforms state-of-the-art methods on the ABIDE and ADNI datasets, with high interpretability.