GeoSAE: Geometric Prior-Guided Layer-Wise Sparse Autoencoder Annotation of Brain MRI Foundation Models

TL;DR

GeoSAE introduces a geometry-guided sparse autoencoder that effectively interprets brain MRI foundation models, identifying reliable biomarkers for Alzheimer's disease with high reproducibility and neuroanatomical relevance.

Contribution

It presents a novel geometry-guided SAE framework that prevents feature collapse and enables interpretable, age-deconfounded biomarker annotation from brain MRI foundation models.

Findings

GeoSAE predicts MCI-to-AD conversion with AUC 0.746 using only 2% of features.

Features replicate across cohorts with r=0.97, showing high reproducibility.

Features localize to neuroanatomical regions consistent with Braak staging.

Abstract

Brain MRI foundation models learn rich representations of anatomy, but interpreting what clinical information they encode remains an open problem. Standard sparse autoencoders (SAEs) suffer from severe feature collapse in deep transformer layers, and in Alzheimer's disease (AD) research, aging confounds nearly every clinical variable, making naive annotation unreliable. We propose GeoSAE, a geometry-guided SAE framework that uses the foundation model's learned manifold structure to prevent feature collapse and annotates each surviving feature via age-deconfounded partial correlations. Applied to ~14k T1-weighted MRI scans from the Alzheimer's Disease Neuroimaging Initiative (ADNI) and the Australian Imaging biomarkers and Lifestyle (AIBL) datasets, GeoSAE identifies a compact, fully interpretable feature set that predicts mild cognitive impairment (MCI)-to-AD conversion (AUC 0.746) using only 2% of the embedding dimensions, while comorbidity-annotated features achieve only chance-level performance. The identified features replicate across cohorts without retraining (r=0.97) and localize to neuroanatomically distinct regions consistent with Braak staging. This shows that geometry-guided SAEs can extract interpretable, biomarkers from frozen brain MRI foundation models.