CytoSAE: Interpretable Cell Embeddings for Hematology

TL;DR

CytoSAE introduces an interpretable sparse autoencoder trained on blood cell images that generalizes well, identifies relevant cellular concepts, and aids in disease classification with explainability at the sub-cellular level.

Contribution

This work presents CytoSAE, a novel sparse autoencoder for hematology that provides interpretable, patient-specific cellular concepts and outperforms existing methods in disease detection.

Findings

Generalizes to diverse hematology datasets

Identifies morphologically relevant cellular concepts validated by experts

Achieves state-of-the-art performance in AML subtype classification

Abstract

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.