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

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.

• The paper introduces GeoSAE, a geometry-guided sparse autoencoder that extracts interpretable brain MRI features from foundation models.
• It employs k-NN manifold regularization and age-deconfounded annotation to prevent feature collapse and ensure clinical relevance.
• Experimental results highlight improved MCI-to-AD conversion prediction and robust cross-cohort replication, enhancing model trustworthiness.

GeoSAE: Geometric Prior-Guided Sparse Autoencoder Annotation for Brain MRI Foundation Models

Introduction

Interpretability in brain MRI foundation models (FMs) is a critical challenge, particularly in the context of Alzheimer's disease (AD) research. High-dimensional, entangled representations learned by state-of-the-art visual transformer (ViT)-based FMs are powerful but opaque; they can inadvertently encode confounds such as aging and comorbidities, leading to unreliable clinical insights. The paper introduces GeoSAE, a geometry-guided sparse autoencoder (SAE) framework that leverages the manifold structure of FM representations to induce interpretability and robustness while explicitly addressing confounding variables through deconfounded annotation.

Figure 1: Overview of GeoSAE. (a) Geometric prior guides SAE activation and $k$-NN manifold graph construction. (b) Manifold regularization prevents feature collapse. (c) Age-deconfounded annotation assigns features to clinical categories.

Methodology

GeoSAE's core contributions are threefold: (1) a geometry-informed regularization framework that mitigates feature collapse common in deep transformer layers; (2) a post-hoc, partial correlation-based feature annotation scheme that controls for age—a dominant confounder in neuroimaging; and (3) cross-layer, cross-cohort validation demonstrating interpretable, stable feature extraction from foundation models without model retraining.

Geometric Analysis and Manifold Regularization

Feature mortality—where most SAE dictionary elements never activate—is acute in brain MRI transformers due to the absence of geometric priors. GeoSAE addresses this through layer-wise geometric analysis, which informs the choice of the TopK activation function and the construction of a $k$-nearest-neighbor (k-NN) graph to capture the latent data manifold. By introducing a manifold regularization loss on encoder pre-activations, gradients are propagated even to otherwise zeroed dictionary elements, resulting in a substantial increase in the number of alive, non-redundant features.

Deconfounded Feature Annotation

Raw feature–outcome correlations in neuroimaging are heavily confounded by age. GeoSAE annotates features by computing partial Spearman correlations with clinical variables, controlling explicitly for age and applying hierarchical FDR corrections. Each alive feature is assigned to the clinical category (e.g., AD, sex, genetic, comorbidity, non-specific) with which it shows the strongest, significant age-deconfounded association. This annotation protocol enables the construction of clinically meaningful feature sets suitable for downstream tasks and avoids conflation of disease-related and confounding signal.

Experimental Results

GeoSAE was evaluated on a large-scale T1-weighted MRI dataset from ADNI (13,218 scans) and externally validated on AIBL (1,266 scans). The foundation model backbone is BrainIAC, a 12-layer ViT.

Cross-Layer Analysis

A cross-layer analysis reveals that feature consolidation parallels depth in the transformer: early layers yield broad, low-level feature sets, while deeper layers synthesize features exhibiting increasing clinical specificity. At layer 9, feature annotation and predictive performance peak.

Figure 2: (a) Alive features by category and clinical annotation rate across layers; depth brings consolidation and specificity. (b) MCI-to-AD conversion AUC peaks at layer 9 before declining as scanner artifact features dominate.

Selective Prediction and Annotation

At the optimal layer, GeoSAE delivers strong selective prediction for MCI-to-AD conversion: the top 16 features (comprising solely AD, sex, and genetic categories) achieve 100% annotation and an AUC of 0.746, exceeding both the full FM embedding (AUC 0.687, 768D) and supervised conventional baselines. Notably, comorbidity features attain only chance-level classification.

GeoSAE's feature set presents minimal inter-feature redundancy compared to standard SAEs, and the high clinical annotation rate demonstrates the effectiveness of manifold regularization in combating collapse while focusing on interpretable signal.

Figure 3: Annotation heatmap for layer 9 GeoSAE features with age-controlled partial correlations, highlighting clear separation by clinical category.

Neuroanatomical Localization

Attention rollout through the transformer, aggregated over maximally activating subjects and mapped to standard brain atlases, reveals that the most predictive AD-conversion features localize to neuroanatomically plausible regions, aligning with established Braak staging of AD progression. Features are spatially segregated, each attending to a unique subset of brain regions.

Figure 4: Top-4 conversion-predictive features show distinct localization—medial temporal lobe, subcortical nuclei, and temporo-parietal networks—matching known AD progression.

Cross-Cohort Replication

Applying the ADNI-trained GeoSAE to AIBL without any retraining demonstrates strong cross-cohort replication (Pearson $r=0.97$), outperforming both standard SAE and variational autoencoder baselines. The annotation assignments and feature activation patterns generalize robustly, indicating the extraction of stable, biologically relevant signals.

Figure 5: Strong cross-cohort annotation replication of GeoSAE features between ADNI and AIBL, colored by clinical category.

Ablation Studies

Comprehensive ablation confirms that manifold regularization is pivotal—removing it drastically reduces the number of active features and their clinical annotation rate. Sensitivity analysis across regularization strengths, dictionary sizes, and activation sparsity demonstrates robustness. Removing AD-related features substantially reduces predictive AUC, while removing comorbidity or non-specific features has minimal or even positive effects, substantiating that predictive gains stem from clinically relevant annotations.

Implications and Future Directions

GeoSAE closes critical interpretability gaps for clinical deployment of brain MRI FMs by extracting orthogonal, clinically specific features that generalize across cohorts. Its methodology sets a precedent for geometry-informed interpretability in biomarker discovery from black-box models, enabling actionable decomposition of FM representations into medically meaningful dimensions. Practically, this advances the trustworthiness and clinical admissibility of AI-driven neuroimaging analytics.

Future work should explore the application of GeoSAE to other foundation model architectures beyond ViT and additional clinical domains. The quadratic scaling of the k-NN construction is a current computational bottleneck but may be mitigated via approximate or distributed neighbor search techniques. Furthermore, manifold-informed regularization could be integrated into end-to-end FM training, potentially enhancing both interpretability and generalization.

Conclusion

GeoSAE demonstrates that the learned geometric structure of large, frozen foundation models can be leveraged to extract sparse, stable, and interpretable clinical features—even in the confounder-rich regime of neuroimaging. The manifold-regularized sparse autoencoder provides a viable pathway for turning opaque high-dimensional representations into actionable neurobiological and clinical insights.