MAESTRO: Masked Encoding Set Transformer with Self-Distillation
Authors: Matthew Eric Lee
Jaesik Kim
Matei Ionita
Jonghyun Lee
Michelle L. McKeague
YONGHYUN NAM
Irene Khavin
Yidi Huang
Victoria Fang
Sokratis Apostolidis
Divij Mathew
Shwetank
Ajinkya Pattekar
Zahabia Rangwala
Amit Bar-Or
Benjamin A Fensterheim
Benjamin A. Abramoff
Rennie L. Rhee
Damian Maseda
Allison R Greenplate
John Wherry
Dokyoon Kim
Venue: ICLR 2025 Poster (2025)
Keywords: self-supervision
representation learning
immunology
biology
single-cell
cytometry
set
set representations
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We developed a self-supervised set representation learning model for vector-sized representations of single-cell data

Abstract: The interrogation of cellular states and interactions in immunology research is an ever-evolving task, requiring adaptation to the current levels of high dimensionality. Cytometry enables high-dimensional profiling of immune cells, but its analysis is hindered by the complexity and variability of the data. We present MAESTRO, a self-supervised set representation learning model that generates vector representations of set-structured data, which we apply to learn immune profiles from cytometry data. Unlike previous studies only learn cell-level representations, whereas MAESTRO uses all of a sample's cells to learn a set representation. MAESTRO leverages specialized attention mechanisms to handle sets of variable number of cells and ensure permutation invariance, coupled with an online tokenizer by self-distillation framework. We benchmarked 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 cell-type proportions and capturing clinically relevant features for downstream tasks such as disease diagnosis and immune cell profiling.

Published as a conference paper at ICLR 2025
MAESTRO: MASKED ENCODING SET TRANSFORMER WITH DELF-DISTILLATION
Matthew E. Lee Jaesik Kim 1:2 .atte Ionita' .Jonghyun Lee" .Michelle McKeague
Yonghyun Nam 1.2 Irene Khavin', Yidi Huang" 1.2 Victoria Fang 1.3 SSokratis Apostolidis
Divij Mathew Shwetank', Ajjnkya Pattekar' Zahabia Rangwala", Amit Bar-Or',
Benjamin A. Fensterheim Benjamii Abramoff' Rennie Rhee' Damian Maseda
Allison R. Greenplate E. John Wherry Dokyoon Kim 1.2,*
Institute for Immune Health and Immunology
DDepartment of Biostatistics Epidemiology and Informatics
'Department of Systems Pharmacology & Translational Therapeutics University of Pennsylvania
Philadelphia, PA 19104, USA
{mat thet.hew.aa. leel peeeeeeeeenmemicinne.... upenn.edu
ABSTRACT
The immune system is a complex network cells, orchestrating coordinated re- sponses throughout the human body a person's lifetime. Cytometry enables profiling this network, but current approaches focus on enumerating and phe- notyping immune cells, failing to quantify the system as 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 representa- tions, MAESTRO uses all of a sample's cells learn 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 ap- plied in cytometry. Our model outperforms existing approaches in retrieving cell- type distributions and capturing clinically relevant features for downstream tasks such as disease diagnosis, age, sex. Code available: https://gitubb.commmatteee IeEESROO
INTRODUCTION
Cytometry, a single-cell technology. has revolutionized biomedical research by enabling high- dimensional profiling of individual cells within biological systems (Bendall et al., (20TT). Cytometry provides information about cell-types, states, and functions, which is crucial for understanding im- mune responses, disease mechanisms, and therapeutic interventions (Spitzer & Nolan 2016). One of the primary challenges in analyzing cytometry data is the number of cells and features of the data (Wang et al. 2023). Hundreds of thousands cells are collected for each sample, where each cell is characterized by measurements across tens of proteins, resulting in complex datasets that are difficult to process and interpret using conventional methods. Processing large-scale single-cell datasets demands substantial resources and efficient algorithms capable of scaling to millions of cells without compromising performance (Kimball et al., 2018). Variable set sizes-due to differences in the number of cells collected from each sample- complicate the application of traditional deep learning models that expect fixed-size inputs. Previous methods mitigate this issue by randomly sampling cells a fixed number, limiting the use of all available data (Hu et al.). Furthermore, cytometry data exhibits permutation invariance, meaning the ordering of cells within a sample is arbitrary and carries no intrinsic information. Models must be robust to such permutations to avoid
"These authors contributed equally this work.

---

Published as a conference paper at ICLR 2025
Cells
Set Input
Teacher
SSd
Inpur
Student
Random Sample
L LinchorOT
Linew
Predicted Reconetruction
Liner
Mask Rate A, B.
LTotal
LinkhornoTT
Lislt-Distillation
Figure 1: MAESTRO Architecture. N cells are sampled due to computational constraints. The sampled subset is fed into student model, while the full set fed to teacher model. The teacher model only requires encoder to provide stable target of the full input set for the student model. The teacher encoder parameters are the exponential moving average the student's s, therefore it does not have the same computational constraints. For the student model, apply masking at various rates to promote diverse learning and use Set Transformer backbone for encoding, followed by a ionmuutaaioniivvriian decoding and reconstruction mechanism. We non-linear projection heads on the latent embeddings to align representations of the full set from the teacher and the masked subset from the student. ISAB, PMA, and SAB blocks defined in 3.1
spurious correlations. Existing methods may fail under these demands, highlighting the need for more scalable solutions.
Self-supervised learning (SSL) is powerful paradigm for learning representations of large, high- dimensional datasets (Chen al. 2020 Radford al. 2021 Kramer 1991). In single-cell analysis, SSL methods have demonstrated success in capturing cellular states, trajectories, and functional characteristics, outperforming traditional gating methods, dimensionality reduction techniques, su- pervised approaches, and standard statistical modeling (Cui et al., (024)) De Donno et al., [023] Bunne et al. 2024). Although current SSL models excel at understanding individual cells, there significant gap in their application to cell populations, as previous methods were designed for individual cells without incorporating sample membership information, resulting in the loss cell population-level insights. This limitation hinders the extraction of meaningful, compact, and compu- tationally practical representations of the entire cellular milieu in samples, making the discovery of multicellular interactions diffieult. Consequently, set representation learning approaches needed to effectively encode and represent cell populations.
Current set representation learning methods inadequate for cytometry data because they either struggle to manage the number of cells inherent to such datasets, or rely on supervised learning approaches instead of self-supervised methods requiring extensive labels (Van Gassen al. 2016). Specifically, some approaches maintain permutation invariance but struggle with the size of cy- tometry data, limiting their effectiveness (Lee et al. 2019 Mialon et al. 2022). Conversely, other methods can handle the size of cytometry, but do not preserve permutation invariance, making their representations sensitive to the ordering of cells within data matrices (Lecun al.) 2013). Ad- ditionally, many existing techniques are supervised, requiring labeled data, which is often scarce in single-cell studies and restricts their ability to leverage large, unlabeled datasets effectively (Hu al.). These combined limitations prevent current set representation learning methods from effec- tively processing and representing cytometry data in single-cell analyses.
We propose MAESTRO (Masked Encoding Set TRansformer with Self-Distillatinn), self- supervised architecture for generating vector representations set-structured data (see Figure MAESTRO employs attention mechanisms tailored for sets, including multi-heas attention pooling (Lee et al. 2019). to handle large, variable-sized datasets while ensuring permutation invariance. By integrating masked encoding strategy and self-distillation framework, it learns informative rep-
2

---

Published as a conference paper at ICLR 2025
resentations without extensive labels. We validate MAESTRO through experiments both global (set-level) and local (clement-level) representations, demonstrating its ability to encode sample- wide features and individual elements that distinguish phenotypes (Chan et al. 202I). Using linear probing, we predict sample-level features like disease phenotypes from the embeddings and as- sess cell-level representations by retrieving cell-type distributions without labels. This showcases MAESTRO'S 's proficiency in encoding fine-grained, high-dimensional cellular information relevant immune profiles.
To the best of our knowledge, MAESTRO is the first attention-basee self-supervised set representa- tion learning architecture capable of handling sets on the order of hundreds of thousands of elements and variable- sets in the context of single-cell data. Our key contributions are as follows:
We introduce MAESTRO, a self-supervised set representation learning architecture that effectively generates vector representations of immune profiles from cytometry data.
We Mww.elopeeMMEETEEEE that addresses the size of biological datasets that were previously too large to handle. Our method works with sets containing hundreds of thousands of elements, expanding what's possible in machine learning and data analysis.
We evaluate MAESTRO against the state-of-the-an cytometry analysis method, manual gating, well as against novel approaches like Deep Sets, Set Transformer, and OTKE (mentioned below) that have not previously been applied to cytometry data.
2 RELATED WORK
Current methods for analyzing cytometry data as set include manual gating. DeepCyTOF, and clustering algorithms to identify cell populations (Mathew et al. (2000) Huu et al.) Van Gassen et al.) 2013). Manual gating a method where experts visually identify and select cell populations of in- terest scatter plots or histograms based on their characteristic protein expressions. Manual gating is labor intensive and subject to operator bias due to manual thresholding of expression levels. Deep- CyTOF uses a 1-dimensional convolutional neural network with pooling, which, due to architecture design is effectively equivalent to Deep Sets (Zaheer et al., 2017). Automated clustering methods offer scalability, but may not capture complex cellular heterogeneity (Van Gassen al. 2015 Traag et al. 2019) and are sensitive to parameter settings, such as the number of clusters (Pierson & Yau 2015). For immunology research, manual gating is the ofate-of-the-art for analyzing cytometry data at the cellular level as demonstratedy by ground truth labels in (Kim et al.) [2024] et al. 2022
Arvaniti Claassen 2017). There are no
methods that exist beyond manual gating and clustering to represent cytometry set data as compact vector, therefore these stand as state-of-the-are for set representation.
Recent machine learning models designed for set data address permutation invariance and variable set sizes inherent in single-cell datasets. Deep Sets (Zaheer et al. 2017) made foundational con- architecture capable
tributions to set representation learning by introducing of handling unordered sets. They established a universal approximation theorem for set functions, demonstrating that any permutation-invariinn function can be represented by applying function to each element individually followed by an aggregation step, such as summation. The Set Trans- former (Lee al. 2019) incorporates attention mechanisms designed process unordered sets. Empirically, we find that size is still too large and other mechanisms must also be incorporated alongside ISAB. The Optimal Transport Kernel Embedding (OTKE) (Mialon et al.) 2022) intro- duces a novel approach by aligning set representations through optimal transport distances while embedding them into a feature space using kernel techniques. OTKE ensures permutation invari- ance by treating sets as distributions, enabling robust and meaningful comparisons regardless of element order. The Differentiable Expectation-Maxiinzatoo (DIEM) (Kim) (022) integrates the Expectation- Matiattiion-Maiiizziiiiiooo (EM) algorithm into differentiable framework, enabling end-to-end training of neural networks on set-structured data. Although the previously mentioned methods lay the foundational framework for MAESTRO, they have not been applied to cytometry data and present limitations. Deep Sets and Set Transformer are supervised approaches, restricting their use to labeled datasets. Deep Sets and Set Transformer struggle to manage the number of cells in cy- tometry data. OTKE requires fixed-size input sets and similarly cannot handle the size of cytometry data. Lastly, DIEM is constrained by the absence of public codebase.
crrmuutationinvaaiia
3

---

Published as a conference paper at ICLR 2025
In self-supervised learning, methods like Masked Autoencoders (MAE). DINO, and iBOT have achieved remarkable success, particularly in computer vision tasks (He et al. 2022 Caron et al, 2021 Zhou et al., 2022). MAE learns representations by reconstructing masked portions of imput data, capturing essential features without relying on explicit labels. DINO employs a student-teacher framework for self-distillation, where the student model learns from the teacher's outputs to achieve high performance, enhancing the efficiency of representation learning. iBOT extends these ideas by introducing online tokenizer, facilitating the learning of meaningful token-based representations. By drawing on the strengths of these frameworks, MAESTRO addresses the unique challenges of set representation learning cytometry data.
METHOD
Let S be set in the form of matrix of cytometry data (where rows single cells and columns are features of cells). We aim to learn a function from the set of all finite sets cells to R", mapping variable-sie sized input sets to fixed-dimensional representations. Standard transformers are unsuitable for unordered sets due to their reliance on positional information. To overcome this, we employ specialized attention blocks lacking positional encodings, ensuring permutation invariance and accommodating variable cardinality. The initial layers use permutation equivariant operations compute attention and extract features elementwise, allowing the attention mechanism adjust appropriately with any input order. Finally, we apply pooling by attention to aggregate set informa- tion into a permutation invariant output. We take advantage of an online tokenizer teacher model to provide an encoded target of the full set of cells for student model that encodes a masked subset of cells, thereby learning a representation of all cells in a set.
3.1 PRELIMINARIES TRANSFORMER ATTENTION BLOCKS
3.1.1 MULTI-HEAD ATTENTION BLOCK
The core component of model is the Multi-Head Attention (MHA) mechanism (Vaswani al., 2017], which enables the model to attend to information from multiple representation subspaces si- multancously. Given input set S {I 22,.......... ,,,,,..,,,,,. where each z; E Rd we first project into queries Q, keys K, and values V using learnable weight matrices Wq. WK. and Wy respectively: Q == SWQ K SWK, and V SWy. The attention mechanism for a single head is defined
(QK' Vale
Attention (Q, K, V) = softmax
V, where d: = dimensionality of K
(1)
To capture diverse features, MHA employs Il parallel attention heads, each computing its own atten- tion output: head,, Attention(Qu, Kr. Vv)) for (l...., H}. These outputs are concatenated and projected through a linear transformation with weight matrix Wo E R Hd. xd. xd,
MHA K, V) Concat(head,,.... ..., head W where head, = Attention (Q,, Kj, Vi)
(2)
The Multi-Head Attention Block (MAB) integrates the MHA with additional components such as a cell level feed-forward neural network denoted by FF to enhance cell-level representations. Specifi- cally, given set Sn and set S,, the MAB is defined as:
MAB(S,, S,) = LayerNorm (H + +F(H))
(3)
where,
H LayerNorm (So + MHA(S SI,
(4)
The Multi-Head Attention (MHA) mechanism exhibits differ- permutation properties depending on which of its inputs permuted, crucial for understanding how our model handles shuffled input data.
Permutation Properties of MHH.
Theorem 1 (Permutation Properties of MHA). Let Q, K,V be the query, key, and value matrices respectively, and Pq. Pkv be permutation matrices. The Multi Heaa Attention mechanism satisfies:
MHAA(PQQ KK PRvV) = PQ MHA(Q,K, V).
Proof. See Appendix B.1
4

---

Published as a conference paper at ICLR 2025
3.1.2 INDUCED SET-ATTENTION BLOCK
The Induced Set-Attention Block (ISAB) extends the MAB efficiently handle large input sets. Based on our own studies, ISAB's alone are insufficient for the number of cells in cytometry data. Given an input set S Raxd and a set of learnable inducing points Rmxd where n is the number of elements in set, m is the number of inducing points, and m << n. the ISAB is formally
defined as a two-step process:
(5) The first step transforms / into H by attending to the input set S. The second step uses H. which encodes information about S, to produce the final output. This structure reduces the computational complexity from Onn²) to (nm)), where n is the size of the input set. This formulation enables efficient processing of large sets assisting in the processing of large cytometry data sets. We find that ISAB's s are also permutation equivariant:
ISAB (S) MAB(S, H) E Rund where H MAB(I, E R
Theorem 2 (Permutation Equivariance of ISAB). The Induced Set-Attention Block ISABm. (S) is permutation equivariant. Formally, for any permutation matrix P, holds that:
ISAB,, m (PS) = P mS).
Proof. See Appendix B.2
3.1.3 POOLING MULTIHEAD ATTENTION
The Pooling by Multihead Attention (PMA) block provides a learnable aggregation mechanism for set-structured data. Given set of input elements S = [x1, (x1,X2 where each E IR. and learnable token E TRd the PMA is formally defined as:
PMA(S) MAB(s, S)
This aitention-based pooling allows the model assign varying importance to different elements during aggregation, which is beneficial for identifying and focusing on the most relevant cell popu- lations or features in cytometry data. This module is crucial to the permutation invariance property of MAESTRO, allowing for the collapse of hundreds of thousands of cells into single, learnable token.
Theorem 3 (Permutation Invariance of PMA). The Pooling by Multihead Attention operation PMA(S) is permutation invariant. Formally, for any permutation of the input set S, holds that:
Proof. See Appendix B.3
3.1.4 SELF-ATTENTION BLOCK
The Sell-Attention Block (SAB) is specialization of the MAB where the input set attends to itself. Given input set S {I1, I2, ... Φn) where each x; E Rdd the SAB is formally defined as: SAB(S) = MAB(S,
The SAB performs self-atten colf-attention among the elements in the set, attending to itself, resulting in an output set the same size as the input. This operation captures pairwise interactions between set elements, and stacking multiple SABs allows the model encode higher-order interactions.
Theorem 4 (Permutation Equivariance SAB). The Self-Artention Block SAB(S) permutation equivariant. Formally, for any permulation mairix P, it holds that:
SAB(PS) P SAB(S).
Proof. See Appendix B.4
3.2 MAESTRO ARCIITECTURE
3.2.1
NON-RANDOM BLOCK MASKING
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 NNon-Radom Block Masking (NRBM)
We introduce ing (NRBM), a technique designed to en- hance the model's learning of meaningful patterns by considering semantic relation- ships within the input data. Given in- put set S = {XI..... j with X; E IRd and desired mask ratio p [0, 1J, NRBM first randomly selects a reference element r from S. then computes cosine simi- larities between r and all other elements, sorting the elements based on these sim- ilarities. This ordering ensures that se- mantically related clements are grouped together. A block mask is then created,
Block Mask-
Input: Set (x),...,x,) F. mask ratio p 00
Output: Masked set S', mask vector M
1. Select random reference x, from S
2. Compute similarities: 8; sim(x;; for all
X; E.S.
Sort S by descending s; to get Ssoned
4. Compute mask size: k = [pn]
Create S' by masking contiguous elements in
Swited
Shufile S'
7. Generate mask vector M corresponding to S
Return: S', M
with |np] masked elements contiguous blocks within this ordered set. The mask is applied by replacing masked elements with special token t. Finally, to maintain the model's permutation invariance, the resulting set is randomly shuf- fled. This approach destroys uninformative relationships due redundant elements, challenging the model to infer broader patterns and ensures the model doesn't 't rely on any particular ordering of the input. We outline NRBM in Algorithm 0 For detailed formulation of the NRBM process, including the specific steps and equations, please refer to Appendix C.T
3.2.2 ENCODER AND DECODER
The encoder f defined as series of nested operations, starting with a linear layer, followed by three consecutive ISABs 3.1.2). and ending with PMA (3..3) operation:
f(S) PMA(ISAB,(IAAB) (ISAB, (Lincar(S3)))) (IALL((cr = where z = f(8) E R is the latent representation.
(8)
The decoder takes the latent representation and applies PMA B.1.3) to unpool the vector to the size of S, this is followed by three consecutive SAB 3.1.4) operations and linear layer to reduce the dimensionality back to the original size.
gz)) Linner(SSB22555 (SAB, (SAA,(PMA(2)) S
where is the reconstructed set.
Reconstruction Loss We employ method called Sinkhorn Optimal Transport (SOT) for reconstruction loss (Cuturi). This approach al- lows us to compare the original and recon- structed sets without relying on specific order- ing of elements. We can think of this method as finding the best way to "match" clements from the original set to elements in the recon- structed set. The Sinkhorn- Linkhorn-Knopp algorithm it- eratively calculates this matching. Algorithm 2 describes the distance calculation from input to output. We define our reconstruction loss using the Sinkhorn Distance:
Algorithm 2 Sinkhorn Optimal Transport Dis-
tance
Input: Sets S, S, Iterations T Output: Sinkhorn Distance d
1. Compute pairwise distance matrix Dij distance(x; E S, x, S)
for i,j, where A is an 2. Initialize AG a
x matrix
3. for = 1 to T do
4. Normalize each row of A such that
Σ, Aij =1
5. Normalize each column of A such that
Σ, Aij =
6. Update Aij = Aij -Das
Cree = SinkhornDistance(S,e
(10)
7.
Compute E., Aij Dis
where S = {xi,..., ) } is the original set and
Return: d
S (8., S,,) } is the reconstructed set.
This loss function ensures that the reconstruction quality is evaluated in a manner, as it only depends how well match the original and reconstructed sets, not their order. For a more detailed mathematical formulation of the SOT and the reconstruction loss, refer to Appendix C.2

---

Published as a conference paper at ICLR 2025
3.2.3 ONLINE TOKENIZER VIA SELF-DISTILLATION
MAESTRO employs a seff-istillation framework with an online tokenizer to tackle the computa- tional challenges of encoding large eytometry datasets. In this context, the teacher model acts as an online tokenizer, providing a token distribution target of the full set of cells (elements) for the student model, which operates on a masked subset. Let S X2,, Xn } be the input set of cells, where each X; E R2 represents cell's nigh-dimensional measurement. We define the student model fa and the teacher model ft using the encoder structure:
fs(mm) ((NBB(((NN)
(11)
f((S) f(S)
(12)
where p the masking rate and S. S is subset of size < n.
The teacher's parameters 0, are an exponential moving average (EMA) of the student's parameters 0g:
0: ad; -1- 0.)0.
(13)
where is the EMA decay rate. The online tokenizer (teacher) provides stable target for the student model. The self-distillation process is achieved by minimizing the KI. divergence between the softmax distributions of the latent representations produced by the student and teacher models:
SD (L(softmax(fx(xx)/)) || softmax((f(x.)/7))
(14)
where is temperature parameter that controls the softness of the probability distributions, and X; E Sm are the elements of the masked subset. This loss encourages the student model to pro- duce latent representations that are similar those of the teacher model, effectively distilling the knowledge from the full dataset into the student model that only sees a masked subset. Notably, this approach feasible because the teacher model, operating on the full set, does not require kack-propagation or gradient calculations, which would be computationally prohibitive for large datasets. The student model, working on smaller subset, can perform gradient calculations and back-propagation, enabling it to learn and update its parameters.
3.2.4 OVERALL
Using all of the above components, describe the MAESTRO algorithm here:
Algorithm 3 MAESTRO Model Overview
Input: Input set mask ratio p. EMA decay rate a, temperature
Output:
Reconstructed set
1. Initialize student parameters 0g, and teacher parameters 0g < 8,
2. Randomly sample subset Sa S
3. Apply NRBM to SM to obtain masked set S' and mask vector M: (S) = NRBM(SM. p)
4. Encode masked set using student encoder: z, = (S': 0.)
5. Decode latent representation to reconstruct the set: S g (zy 0.)
6. Encode original set using teacher encoder: 2, = fi(S; 0z)
7. Compute reconstruction loss: Liee Dinthor SM: $)
8. Compute self-distillation loss: LSD =
Σ
Sarl 1=
KL.
softmax (z)
(2')
softmax
(z)) 2/
SSa
(
9. Compute total loss: - + LSD
10. Update student parameters: 0, + 0. 77 Va, L
11. Update teacher parameters using EMA: 02 t 02 + (1 a)
12. Return reconstructed set $

---

Published as a conference paper at ICLR 2025
4 EXPERIMENT
To evaluate the effectiveness of MAESTRO in representing and analyzing cytometry set data, conducted experiments using large cohort of cytometry samples. We assessed the model's perfor- mance in reconstructing masked cells, latent embedding representativeness, and linear probing tasks to evaluate the effectiveness of embedding important immunological concepts the representation.
4.1 DATA
We utilized a dataset of 1,514 whole blood cytometry samples spanning 14 cohorts and 11 pheno- types (Appendix [E.]] Data were generated at three locations over various time points (Appendix E.2). with raw data showing batch effects (Appendix E.3.I). We employed the technical control sample BatchControlHD2, which exhibits minimal batch effects in learned representations (Figure 3) compared to raw data (Appendix E.3.2). Disease diagnostic and meta data were provided by the primary clinician teams for each study. Appendix E.. displays the distribution of cell counts per 1,,xmum==1338,5222
sample iminimum=11,825
highlighting dataset variability. Each sample is represented matrix with cells rows and proteins columns. Notably, cell types obtained through manual gating are used only to evaluate the representations learned by MAESTRO.
4.2 PRE-TRAINING: RECONSTRUCTION OF MASKED CELLS
All
Unmasked
Input Cells
Masked
Cells
Predicted
UMAAP2
Cells
UMAPI
Figure 2: Set reconstruction. UMAP of the unmasked input cells (yellow, row two), the original masked cells (red, row three), and the predicted cells (blue, row 4). The top row shows all of these overlaid. Each column represents random held-out test sample. Dotted boxes highlight areas where very few no cells exist in the unmasked input set but are still reconstructed by MAESTRO.
To demonstrate MAESTRO's capability in accurately reconstructing masked cells, one of the objec- tives during training. Figure showcases the reconstruction results for eight randomly selected sam- ples from a held-out test set during pre-training. In these plots, the unmasked input cells are depicted in yellow (row two), the original masked cells in red (row three), and the predicted reconstructions by MAESTRO are shown in blue (row four). The visualization shows a projection of the cells onto two-dimensional space using Uniform Manifold Approximation and Projection (UMAP) (Melnnes) et al. 2022). The close alignment between the original and reconstructed cells indicates that MAE- STRO effectively captures the underlying structure the data. Due to the NRBM procedure, there certain regional islands (as denoted by black, dashed boxes) of the UMAP where few masked input cells exist, but are still accurately reconstructed. This demonstrates the model's ability to infer complex patterns and relationships within the data, even for cells that are distinct from other input populations. Pre-training details and implementation including runtime and memory usage can be found in Appendix F2
un-
8

---

Published as a conference paper at ICLR 2025
4.3 VECTORIZED IMMUNE PROFILES
a]
MAESTRO Latent Space
bj
Netrest Neighter Label Matching
We evaluated the quality of MAESTRO's la- tent representations by projecting the 1,024- dimensional embeddings into two dimensions using UMAP for visualization, colored by each sample's 's diagnosis (Figure 3 The resulting plot shows that samples cluster by diagnosis, indicating that MAESTRO effectively captures meaningful patterns related to disease pheno- types. Samples denoted as BatchControlHD2 are technical controls used in our data genera- tion process, which we see cluster among them- selves despite exhibiting batch effects the raw data F.3.2 We quantitatively assessed the representations by generating a contingency ta- ble of nearest neighbor diagnosis proportions (Figure 3). This reveals that many samples have nearest neighbors with matching diagnoses, as shown by the high values along the diago- nal, supporting the model's ability to represent meaningful information.
UMAP1
Haanod fuszhher Razznaa
Figure 3: Vectorized immune profiles. a) Sample-level (set) latent embeddings projected into a low-dimensional space. Each diagnosis has unique shape-color combination for its point. Samples tend to cluster by diagnosis. b) The y- axis is true diagnosis of each point, the x-axis is the diagnosis of each point's 's nearest neighbor. We show the proportion of nearest neighbor diagnosis to each point's true diagnosis. High values along the diagonal demonstrate the majority of samples have a nearest neighbor with matching diagnosis.
4.4 LINEAR PROBING
To evaluate the discriminative power of the learned representations, we performed a linear probing task where used the latent representations inputs to basic regression model for predicting diagnosis, age, and sex. Figure [4] summarizes the performance of MAESTRO compared to other methods (see Appendix D for tabular performance metries). We compare MAESTRO gating and clustering which are the gold standard for sample representations in biomedical research due the lack of set representation methods that can handle the idiosyncrasies of cytometry data. We also provide benchmarks against previously mentioned methods (Deep Sets, Set Transformer, and OTKE) that have never been used on cytometry data before and are suitable for set analysis. These methods are unable to handle the number of cells in a sample or require fixed size input (further highlighting MAESTRO's capabilities), so we have a random subset of 10,000 cells as input for these models. Details implementations for each of these methods can be found in Appendix F3 MAESTRO achieved superior accuracy, area under receiver operator curve, and Fl score in classifying diagnosis and sex, and achieved the lowest MAE and highest R-squared values for age regression. This indicates that the representations learned by our model are more informative downstream prediction tasks. Lastly, we perform ablation study the same task using our MAESTRO architecture, indicating the importance of various design choices (Table The proposed MAESTRO model outperforms other variants of the same model.
manual
Classification (Diagnosis and Sex)
3d Mean Absolute Emrr (MAE) o un
Regression (Age)
1.0
1.0
0.8
0.8
Metric
044
Manual Gating
0.0
Kmeane ( (62A)
OTKE ( 1024)
Doep Sets (supervisd)
Set Transfomeer (supervised) MAESTRO (ours)
Manual Gating
Kmwans (ka 1024)
OTKE 1024)
Deep Sats (supervised)
Set Tranaformer (supervised)
MAESTRO (ours)
0.0
Figure 4: Linear probe. Classification and regression benchmark of sample diagnosis, sex, and age. MAESTRO outperforms all benchmarked methods across all metrics. Our study provides benchmarks against OTKE, DeepSets, and Set Transformer. These methods are suitable for set representation learning but have never been applied to cytometry.

---

Published as a conference paper at ICLR 2025
Model
Accuracy
AUROC
F1
Random Masking
0.887 0.029
0.987 :: 0.007
0.837 :: 0.052
Multi-rate Masking
0.898 1: 0.023
0.987 1: 0.004
0.861 1: 0.036
Masking
0.909 0.025
0.991
0.004
0.873 4: 0.042
Self- Distillation (full model)
0.923 + 0.014
0.992
+ 0.003
0.897 + 0.029
RRandm Masking
0.721 0.023
0.955 0.014
0.485 = 0.029
Table 1: Ablation study. Each subsequent row builds upon the module of previous rows above denoted by +/-. Notably, the masked modelling and self-distillation modules are imperative aspects to the performance of the model. The removal of masked modelling has a major impact on the performance of the model. Details of ablation implementation be found in Appendix F.4
4.5 RETRIEVAL OF CELL-TYPE DISTRIBUTIONS
Cell Proportion Retrinun!
Tele Sww.eaadaa
Many SSL studies that operate at the cell-level aim to predict cell-type or use it as a metric for evaluation (Cui al., 2024 Kim et al. 2024). Our study operates at the sample-level, there- fore we predict entire cell-type distributions, much harder task. Accurately predicting cell- type distributions of samples is of significant biological importance, as it can provide insight into the immune landscape and disease mech- anisms. We assessed MAESTRO's ability to retrieve these distributions by using the latent representations input to single-layer neu- ral network to predict known cell-type distribu- tions obtained through traditional manual gat- MAESTRO's perfor- mance in predicting cell-type distributions av- eraged across all cells (Figure (5) and against 16 individual cell-types (Figure 5p). We bench- mark MAESTRO against k- k-means clustering, OTKE, Deep Sets, and Set Transformer. De- tails on each implementation can be found in MAESTRO outperforms all benchmarks methods, suggesting that MAESTRO not only cap- tures global patterns pertinent to diagnosis, sex, and age but also encodes local, cell-level informa- tion critical for understanding cellular heterogeneity.
6.75
Figure 5: Cell-type distribution retrieval. Color scheme is shared between plot and b. a) From each sample representation we predict the distri- bution of 16 various cell-types. We average across all 16 cell-types to show general performance by mean absolute error (MAE). b) Radar plot show- ing the performance of each model in predicting 16 different cell-types. Each line-color represents a different model, each cell-type is represented by a spoke, values denote MAE. As we move away from the center the MAE gets higher, less cover- age across the map means lower error.
ing.
F5
CONCLUSION
MAESTRO is a learning architecture that generates representations of immunc profiles from high- dimensional cytometry data a clinically relevant modality due to its direct measurement of pro- tein expression on immune cells, providing more immediate functional insights. MAESTRO effec- tively handles variable set sizes, permutation invariance, and computational constraints in single-cell datasets. It produces holistic sample representations that contain immunologically meaningful infor- mation, complementing the standard manual gating approach, as demonstrated by probe experiments evaluating age, sex, and strong diagnosis predictions, highlighting their clinical and immunological relevance. Additionally, the ability to infer cell type distributions from embeddings shows that the embeddings contain biologically meaningful information. MAESTRO's 's holistic representation of the immune system is essential for predicting outcomes, identifying health trajectories, and advanc- ing precision medicine- medicine-ccpaabllites that are difficult to achieve with previous approaches. Further potential applications include treatment response prediction, early disease detection, differentiating vaccine response across diseases, assessing patient protection levels from vaccines at specific time- points, and disease subtyping. Limitations, such as single- sagple batch processing and challenges in detecting rare populations, discussed in Appendix [.6 Future applications include adapting MAESTRO to other single-cell modalities, computational histopathology, and modti repre- sentations further advance biomedical and immunological discoveries.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Eirini Arvaniti and Manfred Claassen. Sensitive detection of rare disease-associated cell subsets via representation learning. Nature Communications, 2017.
Sean C Bendall, Erin F Simonds, Ping Qiu, David A El-ad, Peter O Krutzik, Rachel Finck, Ryan V Bruggner, Raphael Melamed, Alejandro Trejo, Oleksandr I Omatsky, et al. Single-cell mass cytometry of differential immune and drug responses across human hematopoietic continuum. Science, 63266060 2011.
Charlotte Bunne, Yusuf Roohani, Yanay Rosen, Ankit Gupta, Xikun Zhang, Marcel Roed, Theo Alexandrov, Mohammed AlQuraishi, Patricia Brennan, Daniel B. Burkhardt, Andrea Califano, Jonah Cool, Abby F. Dernburg, Kirsty Ewing, Emily B. Fox, Matthias Haury, Amy E. Herr, Eric Horvitz, Patrick D. Hsu, Viren Jain, Gregory R. Johnson, Thomas Kalil, David R. Kel- ley, Shana O. Kelley, Anna Kreshuk, Tim Mitchison, Stephani Otte. Jay Shendure, Nicholas J. Sofroniew, Fabian Theis, Christina V. Theodoris, Srigokul Upadhyayula, Mare Valer, Bo Wang, Eric Xing, Serena Yeung-Levy, Marinka Zitnik, Theofanis Karaletsos, Aviv Regev, Emma Lund- berg, Jure Leskovee, and Stephen R. Quake. How to build the virtual cell with artificial intelli- gence: Priorities and opportunities. 2024.
Mathilde Caron, Hugo Touvron, Ishan Misra, Herve Jegou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties self-supervised vision transformers. Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 9650-9660 2021.
Adam Chan, Wei Jiang, Emily Blyth, Jean Yang, and Ellis Patrick. treekor: identifying cellular-to- phenotype associations by elucidating hierarchical relationships high-dimensional cytometry
data. Genome Biology, 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint arXiv:2002.05 20..000...5703., 2020.
Lijun Cheng, Pratik Karkhanis, Birkan Gokbag, Yueze Liu, and Lang Li. Dgcytof: Deep learning with graphic cluster visualization to predict cell types of single cell mass cytometry data. PLOS Computatational Biology, 2022.
Haotian Cui, Chloe Wang, Hassan Maan, Kuan Pang, Fengning Luo, Nan Duan, and Bo Wang- segpt: toward building a foundation model for single-cell multi-omics using generative ai. Nature Methods, 21:1470-1480, 2024.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport.
Carlo De Donno, Soroor Hediych-Zadeh, Amir Ali Moinfar, Marco Wagenstetter, Luke Zappia, Mohammad Lotfollahi, and Fabian J. Theis. Population-level integration single-cell datasets enables multi-scale analysis across samples. Nature Methods, 20:1683-1693 1692, 2023.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16000-16009, 2022.
Zicheng Hu, Alice Tang, Jaiveer Singh, and Atul J. Butte. A robust and interpretable end-to-end deep learning model for eytometry data.
Jaesik Kim, Matei Ionita, Matthew Lee, Michelle L. McKeague, Ajinkya Pattekar, Mark M. Painter, Joost Wagenaar, Van Truong, Dylan T. Norton, Divij Mathew, Yonghyun Nam. Sokratis A. Apos- tolidis, Cynthia Clendenin, Patryk Orzechowski, Sang-Hyuk Jung, Jakob Woerner, Caroline A.G. Ittner, Alexandra P. Turner, Mika Esperanza, Thomas G. Dunn, Nilam S. Mangalmurti, John P. Reilly, Nuala J. Meyer, Carolyn S. Calfee, Kathleen D. Liu, Michael A. Matthy, Lamorna Brown Swigart, Ellen L. Burnham, Jeffrey McKeehan, Sheetal Gandotra, Derek W. Russel, Kevin W. Gibbs, Karl W. Thomas, Harsh Barot, Allison R. Greenplate, F. John Wherry, and Dokyoon Kim. Cytometry masked autoencoder: An accurate and interpretable automated immunophenotyper.
Cell Reports Medicine, 5, 2024.
Minyoung Kim. Differentiable expectation expectation-maxiiizzei maximization for set representation learning. In Pro- ceedings of the International Conference Learning Representations, 2022.
11

---

Published as a conference paper at ICLR 2025
Abigail K. Kimball, Lauren M. Oko, Bonnie L. Bullock, Raphael A. Nemenoff, Linda F. van Dyk, and Eric T. Clambey. A beginner's guide to analyzing and visualizing mass cytometry data. The Journal of Immunology, 2018.
Mark Kramer. Nonlinear principal component analysis using autoassociative neural networks.
AlChE Journal, 37:233-243 243, 1991.
Yann Lecun, Yoshua Bengio, and Geoffrey Hintong. Deep learning. Nature, 2015.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Richard Kosiorek, Seungjin Choi, and Yee Whye permutation-invartee neural networks. In International Conference on Machine Learning, pp. 3744-3753. PMLR, 2019.
Teh. Set transformer: A framework for
Divij Mathew, Josephine R. Giles, Amy E. Baxter, Derek A. Oldridge, Allison R. Greenplate, Jennifer E. Wu, Cécile Alanio, Leticia Kuri- Kuri-Cervantes, M. Betina Pampena, Kurt D' Andrea, Sasikanth Manne, Zeyu Chen, Yinghui Jane Huang, John P. Reilly, Ariel R. Weisman, Caroline A. G. Ittner, Oliva Kuthuru, Jeanette Dougherty, Kito Nzingha, Nicholas Han, Justin Kim, Ajinkya Pattekar, Fileen C. Goodwin, Elizabeth M. Anderson, Madison E. Weirick, Sigrid Gouma, Clau- dia P. Arevalo, Marcus J. Bolton, Fang Chen, Simon F. Lacey, Holly Ramage, Sara Cherry, Scott E. Hensley, Sokratis A. Apostolidis, Alexander C. Huang, Laura A. Vella, The UPenn COVID Processing Unit, Michael R. Betts, Nuala Meyer, and E. John Wherry. Deep im- mune profiling of covid-19 19 patients reveals distinct immunotypes with therapeutic implications. Science, 369, 2020.
Leland Mclnnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. In International Conference on Learning Representations, 2022.
Grégoire Mialon, Dexiong Chen, Alexandre d' Aspremont, and Julien Mairal. A trainable optimal transport embedding for feature aggregation and its relationship to attention. In Proceedings International Conference of Learning Representations 2022, 2022.
Emma Pierson and Christopher Yau. Zifa: Dimensionality reduction for zero-inflated single-cell gene expression analysis. Genome Biology, 16(1):1-10, 2015.
Alec Radford, Jong-Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Stskever. Learning transferable visual models from natural language supervision. arXiv preprint arXiv:2103.00020, 2021.
Noah Shazeer. Glu variants improve transformer. arXiv, 2020.
Matthew H Spitzer and Garry P Nolan. Mass cytometry: single cells, many features. Cell, 165(4): 780-791, 2016.
V.A. Traag, Waltman L., and N.J. van Eck. From louvain to leiden: guarantecing well-connected communities. Scientific Reports, 9, 2019.
Sofie Van Gassen, Celine Vens, Tom Dhaene, Bart N. Lambrecht, and Yvan Saeys. Floremi: Flow density survival regression using minimal feature redundancy. Journal of Quantitative Cell Sci- ence, 89, 2016.
Sven Van Gassen, Bart Callebaut, Marylene Van Helden, Bart N Lambrecht, Piet Demeester, Tom Dhaene, and Yvan Saeys. Flowsom: Using self-organizing maps for visualization and interpreta- tion of cytometry data. Journal of Quantitative Cell Science, 9777))63-6655 2015.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31st In- ternational Conference on Neural Information Processing Systems, pp. 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc.
Kaiwen Wang, Yuqiu Yang, Fangjiang Wu, Bing Song, Xinlei Wang, and Tao Wang. Comparative analysis of dimension reduction methods for cytometry by fime-of-flight data. Nature Communi- cations, 14, 2023.
12

---

Published as a conference paper at ICLR 2025
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póczos, Ruslan Salakhutdinov, and Alex Smola. Deep sets. In Advances in Neural Information Processing Systems, volume 30, pp- 3391-3401, 2017.
Jinghao Zhou, Chen Wei, Huiyu Wang, Huiyu Shen, Wei Wang, Weicheng Cai, and Nuno Vascon- celos. iBOT: Image BERT pre-training with online tokenizer. In International Conference on Learning Representations, 2022.
13

---

Published as a conference paper at ICLR 2025
A ACKNOWLEDGMENTS
We would like to acknowledge that Figure was created with BioRender.com. Further, we acknowl- edge that Large Language Models (LLMs) were used throughout the study.
14

---

Published as a conference paper at ICLR 2025
B PROOFS
B.1 PROOF OF MHA PERMUTATION PROPERTIES
We need prove three properties of Multi-Head Attention (MHA):
1. Permutation Equivariance: When all inputs are permuted. 2. Permutation Equivariance for Q: When only the query permuted. 3. Permutation Invariance for K and V: When only keys and values are permuted.
Proof. Let Q E R K, V E R xd be the query, key, and value matrices respectively, and Pg R Pkv R X be permutation matrices.
Permutation Equivariance (all inputs): When ng = nku and = Pg = Pke:
(PK)
MHA(PQ. PK, PV) = softmax
PV
Vdg PQK P
= softmax
PV
vdk
QKT
Psoftmax
V
Vdg
PMHA(Q. K, V)
Permutation Equivariance for Q: When only Q is permuted:
(P,
MHA(P, Q, V)
softmax
Vdh
P2QK
softmax
VTl
QK
P Pysoffmax
V
dd
P.M
Permutation Invariance for and V: When only K and V are permuted:
Q(PkoK)
MHA(Q. (Q, PR. K V) = softmax
Pku V
Vdl. QK
PT
softmax
kv
PEo V
Vdl
QK
softmax
V
Vdz
MHA(Q. K, V)
These properties together demonstrate that: a) MHA is permutation equivariant when all inputs are permuted equally. b) MHA is permutation equivariant with respect when only permuted. c) MHA is permutation invariant with respect to K and V when only K and V are permuted.
15

---

Published as a conference paper at ICLR 2025
B.2 PROOF OF ISAB PERMUTATION EQUIVARIANCE
To prove that the Induced Set-Attention Block (ISAB) is permutation equivariant, we need to show that for any permutation matrix P. I (PS) ISA (S).
Proof. Let S 1 be the input set, I Rm be the set of learnable inducing points, and PE E Raxx be any permutation matrix.
Recall that ISAB is defined as two-step process:
H MAB(I, (I, S, ISAB m (S) MAB H, H)
1) First, we examine how permutation affects H:
H' = MAB(I, PS, PS)
LayerNorm(I MHA(I, PS,PS)
From our above MHA proof, we know that MHA is invariant when only K and V are permuted.
Therefore:
H = H
2) Now, we examine the second step:
ISAB (PS) = MAB(PS, H,
LayerNormmPS + MHA(PS, H. H))
3) To show equivariance, we need to prove that this is equal to PISAB..(S)
PISAB, (S) PMAB(S, H. H) PLayerNorm(S MHA(S, H. H))
4) The equality holds because:
(AyerNorm(PS.. + MHA(PS, H, H)) LayerNorm PS PMHA(S, H, H)) PLayerNorm(S MHA(S, H. H))
This equality holds because: a) LayerNorm is permutation equivariant b) MHA is permutation equivariant in its first argument (as shown in the revised MHA proof) c) H remains unchanged under permutation of S
Therefore, we have shown that ISAB.. (PS) = PISAB,, (S) for any permutation matrix P, proving that ISAB is permutation equivariant.
16

---

Published as a conference paper at ICLR 2025
B.3 PROOF OF PMA PERMUTATION INVARIANCE
To prove that the Pooling by Multihead Attention (PMA) operation is permutation invariant, we need show that for any permutation matrix P applied to the input set S, PMA(PS) PMM(S).
Proof. Let S E Bexd be the input set, E RR 1xd be the learnable token, and P E IR be any permutation matrix.
Recall that PMA is defined as:
PMA(S) MAB(s, S,
1) First, let's consider PMA applied the permuted input:
PMA(PS) MAB(s. PS,
2) Expanding the MAB operation:
MAB(s, PS, PS) = LayerNorm(s + MHA(s, PS))
3) From our above MHA proof, know that MHA is invariant when only K and permuted, which is the case here since is not permuted. Therefore:
MHA(s, PS, PS) = MHA(s,
4) Consequently:
MAB LayerNorm(s MHAIS PS, PS))
LayerNorm(s + MHA(s, S,S)
MAB(s,
5) This shows that:
PMA(PS) PMA(S)
Therefore, have shown that PMA(PS) PMA(S) for any permutation matrix P, proving that PMA is permutation invariant.
17

---

Published as a conference paper at ICLR 2025
B.4 PROOF OF SAB PERMUTATION EQUIVARIANCE
To prove that the Self-Attention Block (SAB) is permutation equivariant, we need to show that for any permutation matrix P, SAB(PS) P SAB
Proof. Let S E Rexd be the input set and P E I x be any permutation matrix.
Recall that SAB is defined as:
SAB(S)= == MAB(S,
1) We start by expanding the MAB operation:
SAB(S) MAB(S, S,
LayerNorm(H + FF(H))
where H LayerNorm(S MHA(S,S,S)
2) Now, let's apply permutation to the input:
SAB(PS) MAB PS, PS)
LayerNormí H' + FF(H"))
where H' LayerNorm(Pa) PS + MS,
3) Using the permutation equivariance of MHA when all inputs are permuted (proven earlier):
H'
LayerNorm (PS PMHA(S,S,S))
LayerNorm (P + MHA(S,S,5))
PLayerNNrm(SS MHA(S,
PH
4) The feedforward network FF is applied elementwise, so it's permutation equivariant:
FF(PH) PFF(H)
5) Combining these results:
SAB(PS) = LayerNorm (P FF(PH))
LayerNorm(PP (PH + PFF(H))
LayerNorm (PH + -FF(I)))
PLayerNorm( H +F(H))
PSAB(S)
Therefore, we have shown that SAB(PS) = P SAB(S) for any permutation matrix P, proving that SAB permutation equivariant.
18

---

Published as a conference paper at ICLR 2025
C FORMAL DEFINITIONS
C.I NON-RANDOM BLOCK MASKING
Let & {x1, X2,... X, be an input set where each x; E I. and let p E [0,1] be the desired mask ratio. The Non-Random Block Masking (NRBM) process consists of the following steps:
1. Random Element Selection: Select random element from the input set:
X/o
Uniform(1,.....
This element serves as reference for similarity calculations.
2. Similarity Calculation: Compute the cosine similarity between each element and the ran- domly selected element:
X;
$ cosine cosine www.laiitiiiii.........m
Vi {l,...,
X |||
3. Similarity-based Ordering: Define permutation π : nI - (l....,.)) n) that sorts elements based on their similarity to r:
8t(() 2 8m(j) for all i <
This ordering ensures that similar elements are grouped together.
4. Block Mask Creation: Create binary mask M (M Mn) with the following parameters:
Number of elements to mask: m [mp]
Block size: = max(1, L k//m
Define the mask as:
1, ( mod (2b) < and Ej Mj m
M, = 0, otherwise
This creates contiguous blocks of masked elements in the similarity-ordered sequence.
5. Masked Set Creation: Let t Id be the mask token. Create the masked set S by replacing masked elements with t:
S' = (y1, ... .n), where y, St,
if M, M,
Xx
The complete NRBM function is then defined as:
NRBM : P(Rd) x [0. + P(Rd) x {0, 13" x S,
NRBM(S, = (S',M,m))
where S' is the masked set, M is the binary mask, and π is the permutation used for sorting.
This formulation ensures that:
The masking is based on similarity to randomly chosen element, promoting semantic coherence.
Masks are applied in contiguous blocks, encouraging the model to learn robust patterns.
The original ordering be restored using the returned permutation
19

---

Published as a conference paper at ICLR 2025
C.2 SINKHORN OPTIMAL TRANSPORT
To ensure our reconstruction loss is permutation-invariinn, we employ Sinkhorn Optimal Transport. This method allows to compare the original and reconstructed sets without relying specific ordering of elements.
The key idea to find an optimal 'matching" between the original reconstructed set $ {x..... x This matching represented by a transport plan
{x1 .... Xn } and the
The transport plan P is an x matrix where each element Ps represents how much of element x; from the original set "transported" to element Xj the reconstructed set. Formally:
[..y] R where 0 S P.j S
To ensure that P represents valid matching, constrain it to be doubly stochastic matrix:
Ps = for all and Σ Pa. for
These constraints ensure that each original element is fully 'distributed" across the reconstructed elements, and each reconstructed element is fully 'accounted for" by the original elements.
We define a cost matrix C, where each element Ciy represents the 'cost" of transporting x, to Xy:
Cij = x - 2,
The goal to find the optimal transport plan P that minimizes the total transport cost:
P
arg min
Σ PipCas
subject to the constraints on P
P
1
To make this optimization problem tractable, we add an entropy regularization term:
P
arg min PisCrs - Σ Ps log Pij
bd
where is the regularization strength.
This problem can be solved efficiently using the Sinkhorn-Knopp algorithm, which iteratively up- dates P until convergence.
Finally, we define our reconstruction loss as:
Lrc
Cr
This loss function ensures that the reconstruction quality is evaluated in permutation-invaaiiat manner, as it only depends on the optimal matching between the original and reconstructed sets, not their order.
20

---

Published as a conference paper at ICLR 2025
D APPENDIX TABLES
D.1 DIAGNOSIS
Model
Accuracy
AUROC
F1
Manual Gating
0.381 0.021
0.781 + 0.027
0.190
KMeans (k 1024)
0.563 0.019
0.938 + 0.007
0.334
OTKE (k = 1024)
0.878 0.015
0.988 + 0.003
0.838
Deep Sets (supervised)
0.836 0.016
0.981 + 0.002
0.842
Set Transformer (supervised)
0.884 0.019
0.989 + 0.001
0.851 0.015
MAESTRO (ours)
0.923 + 0.014
0.992 + 0.003
0.897 + 0.029
Table 2: Diagmosis Classification Benchmark. MAESTRO outperforms all benchmarked methods across Accuracy, AUROC, and F1 metries for diagnosis classification.
D.2 SEX
Model
Accuracy
AUROC
Manual Gating
0.528 0.066
0.540 1 0.074
1.515
KMeans (k = 1024)
0.631 0.029
0.687 1 0.036
0.601
OTKE (k = 1024)
0.513 = 0.053
0.521 1 0.072
0.500
Deep Sets (supervised)
0.682 0.058
0.721 + 0.065
0.665 1.667
Set Transformer (supervised)
0.698 0.062
0.756 = 0.059
0.685 0.066
MAESTRO (ours)
0.723 0.061
0.800 + 0.051
0.726 0.062
Table 3: Sex Prediction Benchmark. MAESTRO achieves superior performance across Accuracy, AUROC, and Fl metrics for sex prediction tasks compared other methods.
D.3
AGE
Model
MAE
R²
Manual Gating
19.96
1339.8
-0.670
1 0.118
KMcans (k = 1024)
18.74
+34.58
-0.370
- 0.111
OТКE (k 1024)
8.46
+18.69
0.647
+ 0.099
Deep Sets (supervised)
7.83
0.662
+ 0.082
Set Transformer (supervised)
7.46
= 1.06
0.687
+ 0.054
MAESTRO (ours)
6.70
+ 0.99
0.711 + 0.088
Table Age Prediction Benchmark. MAESTRO outperforms all benchmarked methods with the lowest Mean Absolute Error (MAE) and the highest R² score in age prediction tasks.
21

---

Published as a conference paper at ICLR 2025
E APPENDIX FIGURES
E.1 DATA
a)
Sample Counts by Diagnosis
b)
Cell Counts per Sample
324
300
140
1,386,520
Frequency
250
208
Frequency
120
200
190
100
150
139
130
100
40
Voaculine
Aoute COVID
BatehC ontooonto
Cardiac Ser Sercoldesls
Hidradanttis Suppurativ
Healthy
Long covId
Muttiplo Selerasle
Lympttma
Promatura Infent
Sepale
Skin Concar
12 14
Number of Cells (1e5)
Figure 6: a) We show the counts of samples for each of the 11 different phenotypes (BatchCon- trolHD2 is Healthy) used in our analysis. In total, there are 1,514 samples used. b) We show the distribution of number of cells per sample (the size). The minimum size is 11,829 cells and the maximum set size is 1,386,520 cells.
E.2 DATA GENERATION: TIME AND PLACE
08/20/003
Prepare Date
11/19/2022
02/18/2022
Site 1
05/2/2021
Site
04/28/2211
01/28/2022
10/30/2022
12011/023
05/01/024
Stain Date
Figure 7: The x-axis denotes the date that the sample was stained with antibodies for protein marker detection, the y-axis denotes preparation date, with histograms showing the distribution across these dates. The color denotes the location that the sample was processed which we keep anonymous for double blind review.
22

---

Published as a conference paper at ICLR 2025
E.3 BATCH EFFECTS
E.3.1 ALL SAMPLES
a) Raw Cytometry Data Cells by Study Batch
b)
Raw Cytometry Data Protein Markers by Study Batch
CDta
COllc
CD14
CD16
cD15
CD20
CD25
CD27
CD28
CD38
CD4s
CD56
cbs7
UMAP2
Density
CD6bb
CD123
CD127
cD161
CD1B3
Study
Study
CD185
CD134
CD196
CD197
CD294
Stuty
Stuty
Study
CD4SRA
CD45RO
HLA-DR
bob
TCRgd
Shudy
Study
UMAP1
Figure 8: We randomly sample 1000 cells from all of our samples. a) We plot all of the cells on UMAP colored by their study cohort. While the cells globally cluster by cell-type, we show that there are clear batch effects across the cohorts. b) We take the marker intensity distributions for each study and plot the density. This plot clearly demonstrates batch effects across the studies.
E.3.2 BATCHCONTROLHDD TECHNICAL CONTROL ONLY
a) Raw Technical Control Cells by Study Batch
b)
Technical Protein Markers by Study Batch
CD8
CD4
CD8a
CDIlc
CD14
cD1s
Cbla
Cb26
co2s
CO22
CDOS
CD88
CDiS
cD56
CO57
UMAP2
Density
cbssb
CD1Z3
CD127
CD161
CD183
CouSs
CD194
CD196
CD197
C0294
Study
Sturty
Sturty
CDASRA
CD4SRO
HL
IgD
TCRed
Study
Study
UMAP1
Figure 9: We use BatchControlHD2 as technical control across the denoted studies. a) Here randomly sample 1000 cells from each of these samples and plot them using on a UMAP which clearly show batch effects. We plot the distribution of protein marker expression for each study that BatchControlHD2 is in. Lastly. in Figure 3 we show that BatchControlDD clusters together demonstrating that there less batch effects after processing with MAESTRO.
23

---

Published as a conference paper at ICLR 2025
F APPENDIX INFORMATION
F.I DATA
We use CyTOF (Cytometry by Time of Flight) to generate our cytometry dataset. CyTOF captures high-dimensional data through mass eytometry, technique that simultaneously measures up to 50 or more protein markers at the single-cell level. Unlike traditional flow cytometry, CyTOF uses metal isotopes conjugated to antibodies to bind specific cellular proteins. These metal-tagged antibodies are then detected using time-of-fligh mass spectrometry, which minimizes signal overlap and allows for precise and comprehensive profiling of cellular phenotypes. This method facilitates in-depth analysis of complex biological systems, offering detailed snapshot of cellular states and functions across large cell populations.
F.2
Using 1,514 whole blood cytometry samples, we create train and test set using an 80/20 random split within each diagnosis to ensure that all diagnoses are represented in the test set. This test dataset is never used for training during pre-training, linear probing, ablation, or cell-type propor- tion retrieval. We pre-train MAESTRO on four NVIDIA A A100 GPU's for 156 hours at 29 minutes per epoch (324 epochs). Memory usage varies each sample is of variable size, but average the model consumes 58 GB per GPU. We train on our dataset with Adam optimizer and a batch size of 1. Due to variable size inputs, we can only four samples at once). The learning rate starts at le-4 and using a cosine annealing scheduler to a minimum of 1e-8. We use NRBM to mask four different copies of our samples at masking rates of 20%, 40%, 60%, and 80%. Masking is done on the fly during pre-training and not before-hand. We mask by replacing selected indices with learnable mask token. MAESTRO is configured with four attention heads, hidden size of 2,048, and latent dimension of 1,024. Our input cells start at 30 dimensions (representing 30 different protein markers), which we use a linear layer transform to 1,024 dimensions before using three ISAB (with 2,500 learned points per ISAB) blocks, followed by PMA. After pooling to a single learnable seed (vector), we decode the vector by copying the pooled input to each index of the unmasked cells and use the masking token denote indices where the original cell was masked, We follow this with a PMA block with number seeds equal to the size of the input matrix, and follow this with three SAB blocks. Finally we a linear layer to transform our output to its original 30 dimensions. All attention blocks use a SwiGLU activation function to learn both linear and non- non-linear relationships (Shazeer 2020). Reconstruction is calculated using Sinkhorn Optimal Transport with Euclidean distance for the cost matrix calculation. We align latent representations from the student and teacher model by using non-linear projection head the latent representation. Following the projection head we use softmax activation to convert our repre- sentations probability distributions. We use Kullback-Leibler (KL) - Divergence minimize the difference in distributions. The teacher model is updated with the exponential moving average of the student and momentum value of 0.999. Student and teacher model temperatures are set to 0. and 0.07, respectively.
batch size of one (one for each GPU
F.3 BENCHMARKING METHOD IMPLEMENTATION
MAESTRO is pre-trained with an 80/20 train/test set. For linear probing, the same split is used for training and testing the logistic regression model in all benchmarked methods. The logistic regression model is implemented using scikit-learn LogistieRegression)) We repeat five-fold cross- validation 10 times for each method. Due to imbalances of diagnoses we report the Accuracy, AUROC, AUPRC, and FI-Score.
MAESTRO MAESTRO is pre-trained according to the above. Following pre-training, latent rep- resentation for all samples are derived using the pre-trained MAESTRO model. Each representation is a vector with 1,024 dimensions. Following, the samples are split into the train and test groups as used during pre- training to train and test logistic regression model. The logistic regression model is fit on the training set for 14 various diagnoses, and then inference is done on the held-out test set.
24

---

Published as a conference paper at ICLR 2025
K-means We build two version of the K-means representation. One where we set = and the other k = 1024. We build these representations using the same 80% of training samples used for all other methods. We perform k-means clustering all available cells across the samples. Then, for each sample we calculate the proportions of each cluster that exist. This vector of cluster proportions used the input logistic regression model. The logistic regression model is fit on the training set for 14 various diagnoses, and then inference is done on the held-out test set. Due to imbalances of diagnoses we report the Accuracy, AUROC, AUPRC, and F1-Score.
ОТКЕ The OTKE method is unable to take variable sized sets and using large input sets over- loaded the available GPU memory. For cach sample, we use random subset of 10,000 cells (equal to the number of cells input to the student model of MAESTRO). Then, we train the OTKE method, unsupervised, using our training samples. We use the default parameters as described by OTKE (Mialon et al. ).22). We then create representations for all of our samples using the learned kernel. We use the training samples to train a logistic regression model and report metrics on the held-out test set.
Deep Sets As above, Deep Sets is subset to 10,000 cells (same student model of MAESTRO). Deep Sets is inherently a supervised model so we directly train the model to predict the 14 diagnoses using the 80% training set. We traing using hidden dimension of 2,048 (same MAESTRO), and train until convergence or max 300 epochs (same as MAESTRO). We report metrics on the held-out test set.
Set Transformer As above, Set Transformer is subset to 10,000 cells (same as student model of MAESTRO). Set Transformer is inherently a supervised model so we directly train the model to predict the 14 diagnoses using the 80% training set. We traing using the same encoder architecture as MAESTRO (linear layer, three ISABs, and a PMA) using hidden dimension of 2, 2,048 (same as MAESTRO), and latent dimension of 1,024 (same as MAESTRO) and train until convergence or max 300 epochs (same as MAESTRO). We report metrics on the held-out test set.
F.4 ABLATION IMPLEMENTATION
Unless otherwise denoted, all models in the ablation study use the same training parameters as the above section for MAESTRO pre-training. As above, we use the same 80/20 train/test split for our benchmarking studies.
Random Masking We begin with the most simple model possible. We use only the masked autoencoding module of MAESTRO. We randomly subset each sample to 10,000 cells as we would with the Set Transformer. We randomly mask out elements of the set using 50% masking rate. We encode and decode each sample as described above in MAESTRO pre-training. After pre-training derive representations for each sample using the pre-trained model. Using these representations, we use the 80% training set from pre-training to also train logistic regression model to predict our 14 diagnoses. We report metrics using the held-out test set.
Multi-rate Masking We hypothesize that better representations can be learned using various masking rates, forcing the model to reconstruct missing cells using different amounts of informa- tion. Again, randomly subset each sample to 10,000 cells We randomly mask out elements of the set at rates of 20%, 40%, 60%, and 80%. These values are arbitrarily chosen to cover the spectrum of 0% 100%. We encode and decode each sample as described above in MAESTRO pre-training. After pre-training we derive representations for each sample using the pre-trained model. Using these representations, use the 80% training from pre-training to also train a logistic regression model to predict our 14 diagnoses. We report metrics using the held-out test set.
would with the Set Transformer.
Block Masking We hypothesize that our model is able to reconstruct the masked cells because they are randomly masked, and it's able to learn the reundancy of cells in an entire set. We implement NRBM to prevent our model from learning these relationships. Using the same implementation as the above multi-rate masking, we swap out random masking for NRBM. We encode and decode each sample as described above in MAESTRO pre-training. After pre-training derive representations
25

---

Published as a conference paper at ICLR 2025
for each sample using the pre-trate pre-trained model. Using these representations, the 80% training set from pre-training to also train logistic regression model to predict our 14 diagnoses. We report metrics using the held-out test set.
Self-Distillation From the above we see minimal, although, improving performances for each added module. Here, we introduce our online tokenizer via self-distillation which allows for the encoding the full set cells. This our full model and details on implementation can found in the above section MAESTRO pre-training.
without Random Masking Finally, to test that the ceff-disillaation method itself not what allows for rich representations in MAESTRO, we test performance using only the self deff-distillation module. We encode sets using student teacher model and align the latents of the full set and subset cells. Notably, there is a large decrease in performance when we remove the masked modelling section.
F.5 CELL-TYPE PROPORTION RETRIEVAL
For each sample, 16 cell-type proportions were derived using manual gating, the gold-standard in immunology. We use these cell-type proportions as the target for retrieval for each representation learning method. Since manual gating used derive these values, it is not benchmarked against in this section.
MAESTRO We pre-train MAESTRO using the above setting. After pre-training we derive each of the latent representations of each our samples. Using the same 80/20 train/test split we train one-layer neural network to simultancously predict all 16 cell-type proportions. This is important because each value depends on the others. To demonstrate that this information is already encoded into the latent representation we only train for three epochs.
K-means K-means clustering is an alternative method to identifying cell-types, but is not consid- ered the gold-standard like manual gating is. Because of this, k-means clustering serves as strong baseline for other methods to outperform. Although we use k values of 512 and 1024 (well above still expect this representation
the expected number of cell-types), proportions. Using the same 80/20 train/test split we train a one-layer neural network to simulta- neously predict all 16 cell-type proportions. This is important because each value depends the others. To test whether information is already encoded into the latent representation we only train for three epochs.
be predictive of cell-type
We derive OTKE sample latent representations as described above. Using the same 80/20 train/test split train one-layer neural network to simultaneously predict all 16 cell-type propor- tions. This is important because each value depends on the others. To test whether this information is already encoded into the latent representation only train for three epochs.
OТKE
Deep Sets Deep Sets supervised and therefore we do not use the one-layer neural network as described above. Instead, use the default Deep Sets model with the 16 cell-type proportions the target values. For each sample we subset their matrix randomly to 10,000 cells due to computational constraints. Using default parameters above, because the information is inherently within the input we only use three epochs to train the model test whether it is able to encode the 16 cell-types.
train the model to predict 16 cell-type proportions. As
Set Transformer Set Transformer is supervised and therefore we do not use the one-layer neural network as described above. Instead, we use the Set Transformer encoder as described above, with the 16 cell-type proportions as the target values. For each sample we subset their matrix randomly to 10,000 cells due to computational constraints. Using default parameters predict 16 cell-type proportions. As above, because the information inherently within the input only three epochs to train the model test whether it is able to encode the 16 cell-types.
train the model to
26

---

Published as a conference paper at ICLR 2025
F.6 LIMITATIONS
MAESTRO faces several limitations that span computational, methodological. and interpretive chal- lenges. The model has high computational demands, including long training times and reliance on advanced hardware such as multiple GPUs with high memory bandwidth, make inaccessible to research groups with limited resources. MAESTRO requires singl-sample batch processing due to variable input sizes limits optimization strategies. resulting in longer convergence times. Mask- ing and reconstruction strategies like Non-Random Block Masking (NRBM) depend on specific assumptions about data redundancy, which may not generalize well to diverse datasets. The model's focus on global representations risks underrepresenting rare cell populations and outliers, which are biologically significant in many studies. Interpretability is another major concern, as the latent repre- sentations encode diagnostically relevant information but lack transparency in how specific features prioritized. complicating biological insights. MAESTRO's dependency fixed protein panels restricts its application to datasets with variable panels, while computational intensity renders unsuitable for real-time analysís. The absence of multi-modal data integration limits its adapt- ability to more complex biological studies involving diverse data types. Additionally, its reliance on manual gating for validation-a labor-intensive and operator-dependen. process- scala- bility for larger datasets. These limitations collectively underscore the need for further refinement of the model, increased validation across heterogeneous datasets, and improved accessibility and interpretability to expand its utility in biomedical research.
27


Reviews:
[{"review_id": "OhpnASWMVz", "reviewer": "Reviewer_4TLg", "summary": "This paper presents MAESTRO, a self-supervised learning method tailored to learn representations of high-throughput cytometry data. The complexity and variability of the data makes it impossible to directly apply many of the previously developed techniques, so the authors come up with a new method using the existing teacher-student architecture to learn representations of immune profiles. The authors present evidence of effective data reconstruction, probe representations in predicting sample diagnosis and cell type proportion, and demonstrate superior performance to existing techniques. In addition, the authors report results of the ablation study to justify the design of MAESTRO.", "questions": "__Contributions__\n\n1. Is MAESTRO tailored to the analysis of immune profiles? How well can it generalize beyond that? What would be the evidence of that? If there are no additional experimental results, please discuss potential applications of MAESTRO to set-structured data outside of immunology, and what modifications, if any, might be needed for such applications.\n2. What is the strongest argument to defend the novelty of this work?\n\n__Figure 3b__\n\n3. How do you explain values for Sepsis, Vasculitis, and two types of COVID?\n\n__Table 1__\n\n4. If manual gating performs so poorly, why is it called the golden standard? Please discuss reasons it remains widely used despite the emergence of more accurate computational methods and consider abstaining from calling it a gold standard.\n5. The table includes 2 methods that are supervised. However, the [cyMAE paper](https://www.biorxiv.org/content/10.1101/2024.02.13.580114v2)  suggests that it is gradient boosting decision trees (GBDT) that achieve top performance among the supervised learning algorithms. Why is there no comparison to GBDT?\n6. Is a single linear probing task enough to evaluate the discriminative power of the learned representations? Is it possible they are biased towards sample diagnosis? Please include other evaluation tasks to provide a more comprehensive assessment of the learned representations.\n\n__Evaluations__\n\n7. Some other methods have been included for comparison despite the fact that they are incapable of handling large datasets. To make it possible, the authors sampled 10k cells for each sample ranging from 11k to 1386k cells in total. How fair and informative is that comparison? Were the other methods optimized to achieve their top performance under such conditions? Is it possible to compare MAESTRO to the other methods on a subset of the large dataset under entirely identical conditions? Please provide a more detailed justification for the comparison methodology.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- I appreciated the detailed and comprehensive method description. Clearly structured and explained in sufficient detail. Although many of the blocks are not novel, such a presentation helps understanding the work.\n- The contribution of the work looks solid, as demonstrated in the experiments. The design of the proposed method is justified with an ablation study. The performance appears superior to previously proposed methods.\n- Arguably, there exist few solutions capable of handling high-throughput cytometry data to this day. MAESTRO seems to make a significant contribution in the domain by tackling this challenge effectively.", "weaknesses": "1. It seems that the work addresses an existing task and tackles it by integrating existing concepts and approaches into a new framework. Therefore, the novelty of this work appears limited and must be further clarified by the authors.\n2. Evaluation is done on a single dataset, which is generally not enough to showcase the effectiveness and robustness of the newly presented method. The cited DeepCyTOF, for example, employed five collections of FCM datasets from FlowCAP-I and three additional collections of CyTOF datasets.\n3. Data and code availability are not discussed. For a method paper, an anonymized repository must be provided for reviewers to verify the soundness and validity of the approach.\n4. The authors cite the paper of [cyMAE](https://www.biorxiv.org/content/10.1101/2024.02.13.580114v2) to claim that manual gating remains state of the art, while this very method was introduced at the NeurIPS 2023 Workshop AI4Science as the first effort to achieve (and, arguably, surpass) this state-of-the-art performance. Comparison to cyMAE is neither presented, nor discussed, which is a questionable choice of the study design.\n5. Only a few concluding remarks are dedicated to the limitations of the approach. More discussion points could follow from the additional evaluations that are currently missing.\n6. References look limited suggesting the authors might not be aware of the other important works in the field. Also, some statements are missing citations (e.g., lines 83-92), which complicates validity assessment.\n7. Minor flaws:\n- line 296: missing bracket typo\n- line 313: double-quote typo\n- line 317: “Algorithm 0” typo", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KTLi87cTvF", "reviewer": "Reviewer_72JE", "summary": "The authors developed a method called MAESTRO (Masked Encoding Set Transformer with Self-Distillation) to effectively capture and summarize the diverse characteristics of immune cells from cytometry data. MAESTRO leverages a specialized attention mechanism and a self-distillation framework within a self-supervised learning setup, enabling it to handle large datasets without information loss. The model generates sample-level representations from the data. The authors evaluated MAESTRO’s embeddings to determine whether they can support downstream diagnostic classification, and enable cell-type proportion prediction.", "questions": "See the weakness section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The manuscript is well-written and easy to follow. The proposed method is effectively designed to handle large datasets without losing sample information. Additionally, the model addresses permutation invariance by using specialized attention blocks that omit positional encodings. This design enables MAESTRO to generate robust, representative embeddings for diagnostic classification and cell-type proportion prediction.", "weaknesses": "(1) The model was evaluated only on datasets from similar experimental settings, which contain minimal batch effects. It is unclear how the method handles batch effects or how the resulting embeddings may be influenced by such variations. \n(2) According to the description, the detected proteins could be different between cells. Currently, the authors select and focus only on the shared detected proteins across all the samples. Could the model be extended to handle all the detected proteins?\n(3) Additionally, the model primarily generates sample-level embeddings, whereas producing cell-level (for each cell) and feature-specific (for each feature) embeddings could be valuable for downstream comparisons. \n(4) Further details on the method's runtime, robustness, and memory usage would also be beneficial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "V9KI6v0Zck", "reviewer": "Reviewer_Jhp9", "summary": "In the paper, the authors proposed MAESTRO, a set-transformer-based method that design for generate the sample embeddings and cell embedding for the cytometry data. The authors compared several baseline method, conducted ablation studies, and further checked the effectiveness of the sample/cell embeddings on sample classification/cell type proportion retrieval.", "questions": "See weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is generally well written and easy to follow.\n\nThe authors considered the specific needs for the cytometry data, such as variable set size, large data scale, and the permutation-invariant problem and designed their model, which is great and add a layer of novelty.", "weaknesses": "The downstream experiments presented in sections 4.3 and 4.4 of the manuscript, which focus on sample classification, are adequately performed. However, the section 4.5 dealing with cell type distribution retrieval does not meet the same standard. The rationale behind fine-tuning the embedding for this task is unclear, given the variability and sample-dependence of cell type distributions. This approach lacks the robustness required for generalization across different datasets.\n\nFurthermore, the manuscript does not convincingly demonstrate the utility of the proposed embeddings in broader cytometry tasks. Downstream applications such as zero-shot cell classification, zero-shot sample characterization beyond disease/health state, and protein representation are notably absent. Incorporating these biologically meaningful experiments would significantly enhance the value and applicability of the research. More rigorous and diverse testing of the embeddings on a range of cytometry tasks is essential to establish their effectiveness and relevance in the field.\n\nI did not get the biological importance on the permutation-invariant module. Any downstream tasks to show the effectiveness of the module?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6PdSTdVgY8", "reviewer": "Reviewer_MoSb", "summary": "MAESTRO introduces a self-supervised set transformer framework for analyzing cytometry data, which is known to be challenging due to its high dimensionality, permutation invariance, and variable sample sizes. Using masked encoding and a self-distillation approach, the model generates vector representations of immune profiles by leveraging attention mechanisms to handle set-structured data. MAESTRO performs better cell-type proportion retrieval and disease phenotype classification than state of the art techniques, improving single-cell analysis in immunology research.\nSince cytometry datasets contain millions of cells per sample, MAESTRO proposes a three -fold strategy :\n- A SSL set representation approach with a permutation-invariant attention mechanisms (using __ISAB, PMA, SAB__); \n- A masked encoder (__NRBM__) that enables to process large cytometry datasets efficiently;\n- A teacher-student model for self-distillation via __EMA__.\n\nThus, MAESTRO learns holistic representations of entire cell sets, capturing both global (sample-level) and local (cell-type) information.", "questions": "Authors, in addition to the __above__ cited __perplexities__, may illustrate whether they have plans for exploring the following points:\n   1. How does the __choice of protein markers__ affect MAESTRO’s performance and generalizability;\n   2. How MAESTRO would perform on __multi-modal__ data types like epigenomic data, e.g. __ATAC-seq__;\n   3. How MAESTRO’s embedding can support unsupervised tasks like __clustering__ or anomaly (__blast__ population) __detection__, in a potential __diagnosis__ scenario.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- __Originality__: The paper presents a novel idea to deal with million-row size cytometry dataset, as single cells ones are, providing a model, MAESTRO, that can capture sample membership information without losing cell population-level interactions rather than uniquely focusing on individual cells. \n- __Quality__: The manuscript clearly explains how to face the challenges of large single-cell cytometry datasets exploiting and combining previous-field ideas not deeply applied by compared models yet.\n- __Clarity__: The paper is quite well written, with no major typos or incomprehension.\n- __Significance__: The manuscript, presents a novel idea to face a common problem in a large cytometry dataset, and may well-impacts research due to its:           \n     - __SSL__ module that doesn’t require any label-acquisition process for training, which is particularly costly for these datasets;\n     - ISAB, PMA, and SAB (__Eqs. 5-7__) that are finely tuned to maintain permutation invariance, a critical feature for handling unordered sets like these; \n     - NRBM (__Alg.1__) rather than random masking to contain cell-populations level information.", "weaknesses": "- **Patient-batch limitations:**  \n   - The manuscript doesn’t address the problem of patient-normalization in scenarios where the model may have to deal with a heterogeneous cohort of patients. Cytometry patients' samples may vary a lot in a heterogenous cohort, and further studies on this generalization process could extend MAESTRO applicability (e.g. https://pubmed.ncbi.nlm.nih.gov/31633883/). For example, authors could specify whether would make sense to __inject__  __patient-level__ information as __prior knowledge__ during the pre-training phase.\n\n-  **Scalability concerns with self-distillation on larger datasets and different batch sizes:**  \n   - This approach may become less effective as datasets start spanning over large patient cohorts since MAESTRO  has been pre-trained on four GPUs at a time, with corresponding  __batch_size=1__, meaning four samples at once have been processed. Under extremely large datasets, feeding the teacher model with complete sets can lead to substantial memory requirements.\n\n-  **Dealing with noisy input:**\n   - It’s not explicitly addressed the robustness of MAESTRO when dealing with noisy inputs, such as debris, dead cells that may be inherited from other cytometry datasets (e.g. flow cytometry ones), and whether this could or couldn’t be taken into account in the SSL strategy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)