PaperlayerView exact revision on arXiv

Source-linked AI summary

EXPOSE: Explainable and Domain-Robust Embeddings from Pathology Vision Foundation Models using Sparse Autoencoders

Anja Witte, Maximilian Lennartz, Jan Baumbach, Guido Sauter, Stefan Bonn, Patrick Fuhlert, Marina Zimmermann

arXiv:2608.28191v1cs.CVcs.LG

TL;DR

Computational pathology VFMs are vulnerable to domain shifts, partly because their embeddings entangle biological and domain-specific information. EXPOSE uses an SAE bottleneck and a domain classifier to identify and mask domain-specific latent features before relapse prediction, and this selective masking improves cross-domain robustness across unseen shifts.

  • Problem

    VFM performance can degrade under domain shifts, while biological and domain-specific factors are entangled in representations, making domain-specific dimensions difficult to identify.

  • Method

    EXPOSE learns sparse VFM representations with an SAE, identifies domain-specific latent features using a linear domain classifier, and masks them before training relapse prediction.

  • Results

    Selective masking of domain-specific sparse features improved cross-domain relapse prediction and robustness across additional unseen domain shifts.

  • Takeaways & Limitations

    Sparse representations contain identifiable domain-specific variation that can be selectively masked to improve cross-domain prediction.

  • Takeaways & Limitations

    Evaluation used a single VFM and ReLU-based SAE, and domain-specific features were identified using only one pair of acquisition domains.

Abstract

from arXiv · show

Vision Foundation Models (VFMs) are widely used in computational pathology but remain sensitive to domain shifts arising from variations in staining, tissue preparation, and scanner hardware. A key limitation is that VFM embeddings entangle biological with domain-specific information, hindering cross-domain generalization. We propose Explainable Probing of Cross-Domain Sparse Embeddings (EXPOSE), a framework that uses Sparse Autoencoders (SAEs) as an explainable bottleneck to identify and suppress domain-specific components in VFM embeddings. We train a sparse representation of VFM features, use a linear classifier to identify domain-specific latent dimensions, and mask these features prior to downstream relapse prediction without retraining the backbone model. Experiments on a large prostate cancer dataset with multiple acquisition domains show that SAE features capture both domain- and task-specific information, which are partially disentangled in the latent space. Removing domain-specific features improves cross-domain performance and increases embedding robustness as measured by the Domain Robustness Index (DoRI). Code is available at https://github.com/imsb-uke/expose .

1 Introduction

EXPOSE addresses domain shifts in computational pathology by using sparse autoencoders to expose and suppress domain-specific components in VFM embeddings. The framework identifies domain-related latent dimensions and masks them before relapse prediction, improving cross-domain robustness while partially separating biological and domain-specific information.

  • Motivation: Domain shifts from tissue preparation, staining, and scanner hardware can substantially degrade VFM performance in external domains.VFM representations may entangle biologically relevant information with domain-specific factors, complicating reliable feature selection.
  • Motivation: SAEs transform dense VFM embeddings into sparse, more explainable representations whose latent dimensions can be individually analyzed and suppressed.Prior work suggests sparse features may align with semantically meaningful concepts and dataset-consistent biological structures.
  • EXPOSE framework: EXPOSE trains an SAE, uses a linear domain classifier to rank latent dimensions, and masks the top-k domain-specific features before downstream relapse prediction.The framework uses SAE features as an explainable bottleneck rather than retraining the VFM backbone.
  • Findings: Masking domain-specific features improves cross-domain robustness and provides evidence that biological and domain-specific information are partially separated in the sparse representation.The optimal masking level differs between in-domain and out-of-domain data.

2 Related Work

Related work establishes SAEs as tools for explainable sparse representations and highlights the unresolved question of whether domain-specific information can be automatically localized and suppressed. EXPOSE addresses this gap by learning domain importance directly in the sparse latent space and evaluating robustness after masking.

  • SAEs for explainability: SAEs are used in foundation models to learn sparse features that may activate consistently for distinct semantic concepts.This property is often described as monosemanticity and supports analysis of model representations.
  • Domain robustness: Prior VFM robustness methods include ensembling, distillation, and embedding alignment, but they do not explicitly identify domain-specific feature dimensions.This motivates feature-level domain analysis in sparse representations.
  • SAEs in pathology: Computer-vision studies report that SAE features can capture medically meaningful concepts and recurring biological factors across datasets.These findings suggest sparse representations may separate stable semantic information from dataset-specific variation.
  • Feature intervention: Sparse feature interventions have been explored for visual steering, causal analysis, and masking visually irrelevant artifacts such as ink or staining.ProtoMIL uses expert-driven artifact identification and activation masking to focus downstream models on biologically relevant cues.
  • Open question: It remains unclear whether domain shortcuts are automatically encoded as identifiable sparse features and whether suppressing them improves robustness under constrained domain training.EXPOSE studies this hypothesis with a linear domain classifier and feature masking before cross-domain evaluation.

3 Methods

EXPOSE converts VFM embeddings into sparse representations, identifies domain-associated latent dimensions with a linear classifier, and masks them before relapse prediction. It also measures embedding robustness by comparing biological-label and domain information among nearest neighbours.

  • Sparse representation: EXPOSE trains an SAE to reconstruct dense VFM embeddings as sparse, overcomplete representations with 3,840 latent dimensions.The SAE uses a fivefold expansion of 768-dimensional inputs and ReLU activations for non-negative sparse features.
  • Sparse representation: The SAE objective combines embedding reconstruction with ℓ1 regularization to encourage sparsity.The encoder produces h, the decoder reconstructs z, and training minimizes squared reconstruction error plus a sparsity penalty.
  • Domain-feature identification: A binary domain classifier ranks sparse features by classifier-weight sign and magnitude, selecting ScanB-associated, ScanA-associated, or both-domain features for masking.Positive weights predict ScanB, negative weights predict ScanA, and absolute weights identify features relevant to either domain.
  • Robust relapse prediction: The masked sparse representation h′ = h⊙M is used by a linear sigmoid classifier for 5-year relapse prediction without retraining the backbone.Masking removes features most informative for domain discrimination and leaves the relapse classifier to use the remaining sparse features.
  • Robustness measure: DoRI compares nearest-neighbour agreement in biological labels and domains, normalized against uniform-embedding expectations and bounded from −1 to 1.Positive DoRI means label information has greater influence than domain information; negative values indicate the opposite.

4 Experimental Setup

The experiments use paired ScanA and ScanB prostate-cancer data to train the SAE and domain classifier, while relapse prediction is trained on ScanA and evaluated across six acquisition domains. Patient-level stratification, fixed preprocessing, validation-based tuning, and repeated random seeds define the evaluation protocol.

  • Data and splits: ScanA and ScanB each contain 8,141 registered samples for SAE and domain-classifier training, while relapse classification is trained on ScanA.Both datasets use 70% training, 15% validation, and 15% test splits with patient-level stratification by 5-year relapse label.
  • Data and splits: Four additional subdatasets—ScanA.spot, ScanA.thin, ScanA.thick, and ScanA.long—are withheld from training and used as evaluation domains.These subdatasets represent additional tissue-preparation or digitization shifts.
  • Preprocessing: H0-mini extracts 768-dimensional patch embeddings from centered, downsampled TMA crops that serve as SAE inputs.Images are cropped to 2048 × 2048 pixels at 40× magnification and downsampled to 5×.
  • Training protocol: All three tasks use validation-based hyperparameter tuning and early stopping, with the SAE optimized using AdamW, learning rate 1.5 × 10^-4, and ℓ1 weight 10^-5.The classifiers also use AdamW and batch size 128; reported results are generally means ± standard deviations over five seeds.
  • Evaluation: Relapse prediction is evaluated with AUROC5 on all six subdatasets against a frozen-VFM linear base model.This separates in-domain performance from cross-domain generalization.

5 Results

SAE features encode domain- and relapse-related patterns that are partly separated, and masking domain-associated features generally improves robustness. Benefits depend on the selection strategy, masking level, and evaluation domain, with unseen shifts remaining unevenly addressed.

  • 5.1 Identifying Domain-Specific SAE Features: Top positive domain-weight masking is the most consistent strategy across ScanA and ScanB relapse evaluation.Top absolute-weight masking reaches comparable maxima but is less consistent across masking levels.
  • 5.1 Identifying Domain-Specific SAE Features: 0.7159 AUROC5 is achieved on ScanA by masking the top 256 positive-weight features, versus 0.7145 without masking.On ScanB, the same strategy reaches 0.7130 versus approximately 0.7097 for the unmasked representation.
  • 5.1 Identifying Domain-Specific SAE Features: Masking top negative ScanA-associated features provides no consistent benefit and substantially degrades performance at large masking levels.Masking 3,500 positive-weight features partially recovers performance, indicating relapse information is not confined to the masked dimensions.
  • 5.2 Task-Specific Sparse Features: Domain- and relapse-specific features both show cross-domain activation consistency, but domain-specific features more often respond to background or dataset-specific shortcuts.Relapse-specific features predominantly focus on coherent tissue structures rather than acquisition artifacts.
  • 5.2 Task-Specific Sparse Features: Activation differences between scanners align strongly with domain-classifier weights, whereas relapse weights occupy separate regions with little spatial overlap.Relapse features concentrate where absolute domain weights are low, supporting partial disentanglement in the SAE space.
  • 5.3 Masking Domain-Specific Features Improves Domain Robustness: ID performance peaks after masking k = 256 features, while OoD performance peaks after masking k = 3500 features, leaving 340 feature dimensions.Thus, shifted target domains favor substantially stronger masking than the training domain.
  • 5.3 Masking Domain-Specific Features Improves Domain Robustness: On OoD data, masking k = 3500 features changes performance from −0.54 pp for unmasked SAE to +0.75 pp.ScanA.thin remains challenging because domain-feature masking does not improve its performance.

6 Conclusion

EXPOSE uses sparse autoencoders as an explainable bottleneck to identify and mask domain-specific information in VFM embeddings, improving cross-domain relapse prediction without modifying the underlying VFM. The study also finds partial separation of domain- and relapse-specific information, while noting limitations in model and domain coverage.

  • Conclusion: EXPOSE quantifies sparse-feature domain dependence with a linear classifier and masks the most predictive dimensions without retraining the VFM.This provides post-hoc suppression of domain-specific information in the embedding space.
  • Conclusion: Selective masking improved cross-domain relapse prediction over complete embeddings, including robustness to unseen domain shifts despite identifying features from one domain pair.The result supports selectively masking identifiable domain-specific variation rather than removing features indiscriminately.
  • Conclusion: Embedding analysis found domain-specific features aligned with scanner-dependent activation shifts, while relapse-specific features occupied different latent-space regions.Highly activating images and patches showed consistent, explainable activation patterns for several sparse dimensions.
  • Conclusion: Increased DoRI scores after masking suggest that the gains reflect suppression of domain-specific variation rather than simple information reduction.
  • Limitations: The evaluation is limited to one VFM, one ReLU-based SAE, and domain-feature identification from a single acquisition-domain pair.Future work should test other VFMs, advanced SAE architectures, and additional domain shifts.
  • Conclusion: EXPOSE provides an explicit ranking of sparse dimensions by domain or biological relevance, supporting further biological interpretation and robust model development.The authors propose refining this ranking toward highly informative features evaluated by expert pathologists.

A.1 Additional Dataset Information

Table 2 summarizes dataset variants derived from ScanA, contrasting the standard acquisition protocol with controlled changes in scanning, staining, slicing, or sampling.

  • Dataset variants: ScanA is the standard acquisition protocol, while the other variants introduce controlled domain changes.
  • Dataset variants: The variants differ through scanner, staining time, slicing thickness, or sampling location.Differences from the standard protocol are highlighted in the table.
  • Dataset variants: ScanA and ScanB contain registered images of the same TMA spots scanned with different scanners, whereas other tissue-core-1 variants use different slices of the same core.

A.2 Additional Results on Domain Feature Selection Strategies

Table 3 evaluates three strategies for selecting domain-specific sparse features to mask, using ScanA and ScanB AUROC5 performance across validation subdatasets and masking levels.

  • Selection strategies: The table compares three domain-feature selection strategies for masking sparse features.
  • Selection strategies: Results are reported separately for ScanA and ScanB performance in panels (a) and (b).
  • Selection strategies: Performance is evaluated across different numbers of masked features k, with highlighted values marking each strategy’s best and second-best results.

A.3 Additional Results on Domain Feature Masking

Table 4 examines how masking domain-specific sparse features affects cross-domain relapse prediction performance across ID and OoD evaluation sets and varying masking levels.

  • Domain feature masking: The table measures relapse prediction using AUROC5 after masking domain-specific sparse features.
  • Domain feature masking: Results are presented separately for in-domain (ID) and out-of-domain (OoD) evaluation sets.
  • Domain feature masking: Performance is compared across different numbers of masked features k, with highlighted values indicating the best and second-best result for each dataset.
Loading 2608.28191v1…