EXPOSE: Explainable and Domain-Robust Embeddings from Pathology Vision Foundation Models using Sparse Autoencoders
Computer Vision and Pattern RecognitionMachine Learning
Summary
The authors address the problem that Vision Foundation Models used in pathology can get confused by differences in how tissue samples are prepared or scanned, affecting their accuracy. They created a method called EXPOSE, which uses a special kind of neural network to find and block out parts of the model's learned features that are specific to these differences. By doing this, their method helps the model focus on the biological information instead, improving its performance across different domains without needing to retrain the whole model. Their tests on prostate cancer data showed this approach made the model's predictions more reliable when dealing with varied data sources.
Authors
Anja Witte, Maximilian Lennartz, Jan Baumbach, Guido Sauter, Stefan Bonn, Patrick Fuhlert, Marina Zimmermann
Abstract
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 .