AG-ViT: Atlas-Guided Vision Transformer for Predicting Cognitive Decline

ICLR 2026 Conference Submission17470 Authors

Keywords: Transformers, fMRI, Alzheimer's disease, self-supervised learning

TL;DR: We develop a transformer for learning fMRI representations from Alzheimer’s patients and predicts the probability of cognitive decline within a given time frame.

Abstract:

The use of resting state fMRI (rs-fMRI) to improve the diagnosis and treatment of neurodegenerative diseases has increased dramatically in recent years. Despite evident progress, producing accurate predictions from rs-fMRI scans remains challenging due to the data's high dimensionality and the limited number of samples. In this work, our aim is to estimate the probability of cognitive decline within a given time frame based on rs-fMRI scans of Alzheimer’s patients. Accurate predictions of disease trajectory can guide medical decision-making and contribute to personalized medicine. To this end, we design a vision transformer to obtain low-dimensional embeddings of rs-fMRI scans. These embeddings are used to train a network that estimates the probability of cognitive decline. By testing our approach on scans from the Alzheimer's Disease Neuroimaging Initiative, we show that models trained on our transformer-based features outperform those trained on handcrafted connectivity features by 15\%–26\% in F1-score.
For interpretability, we develop a simple yet effective method to identify brain regions whose fMRI-derived signal significantly impacted model predictions. The results identified a set of brain regions, some recognized for their early involvement in AD and others for their relative resilience to AD pathology.

Primary Area: applications to neuroscience & cognitive science

Submission Number: 17470