Technical Summary: GeoSAE

1. Problem Statement

Brain MRI foundation models (FMs), pretrained via self-supervised learning on tens of thousands of scans, produce high-dimensional, entangled representations that are difficult to interpret. In Alzheimer's disease (AD) research, this opacity is particularly problematic because aging is both the primary risk factor and a pervasive confound; features that appear to predict AD progression may merely reflect normal aging, comorbidities (e.g., cardiovascular disease, diabetes), or scanner artifacts.

Existing interpretation methods face specific limitations:

• Post-hoc attribution (e.g., GradCAM, Integrated Gradients): Produces spatial saliency maps but fails to decompose internal representations into discrete, nameable concepts.
• Linear probing and Concept Activation Vectors (TCAVs): Treat the latent space as a monolithic vector, unable to identify which specific dimensions encode which information.
• Disentanglement methods: Require predefined labels during training and architectural modifications, making them inapplicable to frozen, self-supervised FMs.

Furthermore, applying standard Sparse Autoencoders (SAEs) to deep transformer layers in neuroimaging FMs introduces three specific challenges:

1. Feature Mortality: Standard SAEs suffer severe feature collapse in deep layers, losing >98% of dictionary elements because no geometric prior guides regularization.
2. Age Confounding: Raw feature-outcome correlations conflate neurodegeneration with normal aging.
3. Comorbidity Confounding: Features must disentangle disease-specific signals from shared risk factors associated with comorbidities.

2. Methodology: GeoSAE

The authors propose GeoSAE (Geometry-guided Sparse AutoEncoder), a three-stage framework designed to extract interpretable, age-deconfounded features from frozen brain MRI FMs without modifying the base model.

Stage 1: Geometric Prior Analysis

Before training the SAE, the framework analyzes the geometry of the FM's learned manifold for each transformer layer.

• Activation Selection: The analysis evaluates geometric properties (e.g., angular vs. radial class structure, dimensionality homogeneity) to select the optimal SAE activation function. For the BrainIAC model, TopK activation was selected as it best matched the data's angular-dominant structure.
• Manifold Graph Construction: A $k$-nearest-neighbor (k-NN) graph is constructed on the FM embeddings. Samples considered similar by the FM (neighbors in representation space) are connected with Gaussian kernel weights. This graph serves as a geometric prior for the subsequent training stage.

Stage 2: Manifold-Regularized SAE Training

GeoSAE introduces a manifold regularization term to the standard SAE loss function to address feature mortality.

• Loss Function: The total loss combines the standard reconstruction loss ($\mathcal{L}_{SAE}$) with a manifold smoothness penalty on the pre-activations ($a$) before the sparsity gate:
  $$\mathcal{L} = \mathcal{L}_{SAE} + \lambda \sum_{(i,j) \in \mathcal{N}_k} w_{ij} \|a_i - a_j\|_2^2$$
• Mechanism: By penalizing dissimilar pre-activations for samples that are neighbors on the FM's manifold, the regularization provides gradient signals to all dictionary elements, including those zeroed out by the sparsity gate (TopK). This prevents the collapse of features into a small active set, a common failure mode in deep layers.

Stage 3: Age-Deconfounded Feature Annotation

To assign clinical meaning to surviving features, GeoSAE employs a partial correlation approach.

• Partial Correlations: Instead of raw correlations, the method computes partial Spearman correlations ($\rho_{jc \cdot a}$) between each feature $j$ and clinical variables $c$ (diagnosis, sex, APOE4, comorbidities), controlling for age ($a$). This removes the variance shared with age.
• Category Assignment: Features are assigned to a clinical category (e.g., AD-related, sex-related, genetic, comorbidity) based on the strongest significant partial correlation. Features with significant raw age correlation but no specific clinical association are labeled "aging," while others are "non-specific." Hierarchical False Discovery Rate (FDR) correction is applied to control for multiple testing.

3. Experimental Setup

• Datasets: Training and primary evaluation on ADNI (13.2k T1-weighted scans); external validation on AIBL (1.2k scans).
• Foundation Model: BrainIAC, a 12-layer Vision Transformer (ViT) pretrained on 85k+ scans.
• Tasks:
  1. Enrichment Testing: Validating feature categories against 31 secondary clinical variables.
  2. Selective Prediction: Predicting MCI-to-AD conversion using logistic regression on specific feature subsets.
  3. Cross-Cohort Replication: Applying the ADNI-trained SAE to AIBL without retraining to test stability.

4. Key Results

Feature Decomposition and Survival

• Feature Count: At the optimal layer (Layer 9), GeoSAE produced 161 alive features, which is 7× more than a standard SAE (23 alive features) trained without manifold regularization.
• Layer Trends: Alive features peaked in early layers and consolidated with depth. Clinical specificity increased with depth, with Layer 9 showing the highest concentration of AD-related, sex-related, and genetic features.

Predictive Performance

• MCI-to-AD Conversion: A subset of the top 16 features (selected by activation frequency) achieved an AUC of 0.746.
  • This outperformed the raw 768-dimensional FM embedding (AUC 0.687).
  • It surpassed a supervised 3D ResNet-18 baseline (AUC 0.729) and standard SAEs (AUC 0.737).
  • Crucially, this performance was achieved using only 2% of the embedding dimensions (16 dimensions).
• Specificity: Features annotated as "AD-related" drove the prediction performance. Features annotated as "comorbidity" achieved only chance-level performance (AUC ~0.507), confirming the framework successfully disentangled disease signals from comorbidities.

Replication and Localization

• Cross-Cohort Stability: The top 16 features replicated across ADNI and AIBL without retraining with a correlation of r=0.97 for activation consistency. This is a 32% improvement over standard SAEs (r=0.74).
• Anatomical Localization: Attention rollout maps localized features to neuroanatomically distinct regions consistent with Braak staging:
  • Early Braak stages: Medial temporal lobe (hippocampus, parahippocampal gyrus).
  • Advanced stages: Subcortical nuclei and temporo-parietal networks.
• Enrichment: AD-related features were significantly enriched for AD medication usage (e.g., memantine, donepezil), validating the biological relevance of the extracted features.

5. Contributions and Significance

The paper claims three primary contributions:

1. GeoSAE Framework: A novel method using geometric analysis of FM representations to guide activation function selection and introducing manifold regularization on pre-activations to prevent feature collapse in deep layers.
2. Deconfounded Annotation Pipeline: A rigorous annotation strategy using age-deconfounded partial correlations and hierarchical FDR correction to assign clinical labels to features, effectively separating aging, comorbidity, and disease signals.
3. Empirical Validation: Demonstration that a compact set of 16 interpretable features from a single layer of a frozen FM can predict AD conversion better than the full representation, replicate across independent cohorts, and localize to biologically plausible brain regions.

Significance:
The authors assert that GeoSAE demonstrates that the geometric structure inherent in foundation model representations can be leveraged to extract clinically transparent, mechanistically interpretable features from black-box embeddings. This approach addresses the critical need for reliable, interpretable biomarkers in neuroimaging while controlling for the pervasive confounds of aging and comorbidity. The work is modestly scoped to a single FM (BrainIAC) and a single domain (AD), noting that generalizability to other architectures and diseases requires further investigation.