CytoSAE: Interpretable Cell Embeddings for Hematology
This work addresses interpretability for medical imaging models, specifically in hematology, offering a tool for clinicians to understand model inferences, though it is incremental as it adapts existing SAE methods to a new domain.
The paper tackled the lack of interpretability tools for medical imaging foundation models by applying sparse autoencoders to hematology, resulting in CytoSAE, which identifies morphologically relevant concepts validated by experts and achieves performance comparable to state-of-the-art in AML subtype classification while providing sub-cellular explainability.
Sparse autoencoders (SAEs) emerged as a promising tool for mechanistic interpretability of transformer-based foundation models. Very recently, SAEs were also adopted for the visual domain, enabling the discovery of visual concepts and their patch-wise attribution to tokens in the transformer model. While a growing number of foundation models emerged for medical imaging, tools for explaining their inferences are still lacking. In this work, we show the applicability of SAEs for hematology. We propose CytoSAE, a sparse autoencoder which is trained on over 40,000 peripheral blood single-cell images. CytoSAE generalizes to diverse and out-of-domain datasets, including bone marrow cytology, where it identifies morphologically relevant concepts which we validated with medical experts. Furthermore, we demonstrate scenarios in which CytoSAE can generate patient-specific and disease-specific concepts, enabling the detection of pathognomonic cells and localized cellular abnormalities at the patch level. We quantified the effect of concepts on a patient-level AML subtype classification task and show that CytoSAE concepts reach performance comparable to the state-of-the-art, while offering explainability on the sub-cellular level. Source code and model weights are available at https://github.com/dynamical-inference/cytosae.