MAESTRO: Masked Encoding Set Transformer with Self-Distillation

The immune system is a complex network of cells, orchestrating coordinated responses throughout the human body over a person's lifetime. Cytometry enables profiling of this network, but current approaches focus on enumerating and phenotyping immune cells, failing to quantify the system as a whole. We present MAESTRO, a self-supervised set representation learning model that generates vector representations of set-structured data, which we apply to learn immune profile representations. Unlike previous studies that learn cell-level representations, MAESTRO uses all of a sample's cells to learn a set representation. MAE-STRO leverages attention mechanisms to handle sets of variable number of cells and ensure permutation invariance, coupled with a self-distillation framework. It is capable of reconstructing immune profiles (cells) even when 90% are hidden as its training objective. We benchmark our model against existing cytometry approaches and other existing machine learning methods that have never been applied in cytometry. Our model outperforms existing approaches in retrieving celltype distributions and capturing clinically relevant features for downstream tasks such as disease diagnosis, age, sex. Code available: https://github.com/matthew-lee1/MAESTRO