Learning biologically relevant features in a pathology foundation model using sparse autoencoders

TL;DR

This paper demonstrates that sparse autoencoders trained on pathology foundation model embeddings can extract interpretable, biologically relevant features that correlate with cell types and improve robustness, aiding clinical interpretability.

Contribution

The study shows that sparse autoencoders can uncover biologically meaningful features from pathology model embeddings, which are unique to pathology models and generalize across datasets.

Findings

SAE features correlate with cell type counts like plasma cells and lymphocytes.

Biological features emerge across the model's depth and improve robustness.

Biologically-grounded features generalize to out-of-domain datasets.

Abstract

Pathology plays an important role in disease diagnosis, treatment decision-making and drug development. Previous works on interpretability for machine learning models on pathology images have revolved around methods such as attention value visualization and deriving human-interpretable features from model heatmaps. Mechanistic interpretability is an emerging area of model interpretability that focuses on reverse-engineering neural networks. Sparse Autoencoders (SAEs) have emerged as a promising direction in terms of extracting monosemantic features from polysemantic model activations. In this work, we trained a Sparse Autoencoder on the embeddings of a pathology pretrained foundation model. We found that Sparse Autoencoder features represent interpretable and monosemantic biological concepts. In particular, individual SAE dimensions showed strong correlations with cell type counts such as plasma cells and lymphocytes. These biological representations were unique to the pathology pretrained model and were not found in a self-supervised model pretrained on natural images. We demonstrated that such biologically-grounded monosemantic representations evolved across the model's depth, and the pathology foundation model eventually gained robustness to non-biological factors such as scanner type. The emergence of biologically relevant SAE features was generalizable to an out-of-domain dataset. Our work paves the way for further exploration around interpretable feature dimensions and their utility for medical and clinical applications.