TARDRL: Task-Aware Reconstruction for Dynamic Representation Learning of fMRI

The mask autoencoder (MAE) is utilized in functional magnetic resonance imaging (fMRI) analysis to construct brain representation learning models and conduct prediction for various fMRI-related tasks (e.g., disease detection). 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. Although the MAE helps to improve prediction performance, directly applying it to fMRI may lead to sub-optimal results for the following reasons: 1) The reconstruction process is not task-aware, meaning the extracted brain representations are unable to sufficiently consider downstream tasks, thereby affecting prediction performance; 2) Random masking of fMRI data ignores that the varying contributions of different brain regions to different prediction tasks. To address these issues, we propose Task-Aware Reconstruction Dynamic Representation Learning (TARDRL). Different from the conventional sequential design, this approach sets up reconstruction and prediction tasks in parallel to learn robust task-aware representations. Based on the parallelized framework, we leverage attention maps from specific tasks to guide the fMRI time series reconstruction, which in turn helps to learn task-aware fMRI representations and improve disease prediction accuracy. Extensive experiments demonstrate that our model outperforms state-of-the-art methods on the ABIDE and ADNI datasets, with high interpretability. The codes are available in the repository.