MMIL: A novel algorithm for disease associated cell type discovery

We wish to train a classifier to predict whether a given cell is healthy or diseased. Typically, we would (1) gather a dataset of cells where we know which cells are diseased and which are healthy; and (2) train a model to predict cell disease status (using e.g. logistic regression, random forest, neural network). This workflow relies on knowing which cells are diseased and which are not. In our setting we have no such luxury: we do not know which cells sampled from sick patients are diseased, we only know that cells sampled from healthy people are healthy. MMIL is an algorithm that trains classifiers in this setting.

MMIL is an iterative process that relies on the following intuition. If we had a trained classifier, we could use it to predict which cells are healthy and which are not. And if we had a prediction of which cells were healthy and which were not, then we could use these predicted labels to train a classifier. MMIL fits a classifier by alternating between these two steps, and is described in Algorithm 1 (and in more detail in Appendix B).

We now consider a publicly available dataset of 800 thousand cells collected from 13 Acute Myeloid Leukemia (AML) patients and three healthy bone marrow controls [5] (see Methods). Each AML sample is a mixture of cancer cells and healthy cells, and this cohort of patients represents a wide range of leukemic blast counts. Typically, a highly-trained hematopathologist must analyze a blood or bone marrow specimen to diagnose and track the progression of disease over time. Hematopathologists do this by manually enumerating the proportion of leukemic blasts in the sample based on each cell’s size, shape, and expression of leukemia-associated markers measured by microscopy or flow cytometry. Particularly in the context of myeloid leukemias, blast enumeration is challenging and error-prone even among expert pathologists [5, 6].

Using MMIL, we sought to automate the identification of cells specific to AML without expert cell-level annotation. Importantly, each cell in the AML dataset has a gold-standard annotation indicating whether it is a leukemic blast as evaluated by a board-certified hematopathologist. These annotations allow us to evaluate MMIL’s performance in identifying leukemic blasts compared to expert clinical judgement.

Each cell was analyzed for the presence of 32 surface proteins and 11 intracellular proteins [5] using mass cytometry (CyTOF), a high-dimensional cytometry platform similar to clinical cytometers commonly used to analyze leukemic tissue specimens in clinical laboratories [7, 8]. However, CyTOF analysis allows a high-dimensional, single-cell characterization of each sample, extending beyond the capabilities of conventional multicolor flow cytometers and making it ideal for deep phenotyping of leukemic cells in a research environment [9].

Using this dataset, we trained 3 models: (1) the “optimal” model, a lasso logistic regression model trained using gold-standard annotations (usually unknown), (2) the ”naive” model, a lasso model trained using patient labels instead of cell labels, and (3) Mixture Lasso, a lasso model trained with MMIL. When training Mixture Lasso, we assumed that $75\%$ of cells in cancer patients are healthy, based on the clinical guideline that patients with at least $25\%$ of blasts in their bone marrow receive a leukemia diagnosis[10]. We estimated the performance of each of these models using leave-one-(patient)-out cross-validation (LOOCV). In each fold of our cross-validation procedure, models were trained using all 3 healthy control marrows and 12 of the 13 AML marrows; then, the performance of each model was evaluated by comparing predictions on the held-out patient’s cells to the pathologist’s gold-standard labels (Figure 1b).

At the cell-level, Mixture Lasso achieved a mean area under the receiver-operator characteristic curve (AUROC) of 0.751 across all held-out patients during the cross-validation. By comparison, the naive model achieved worse performance with an AUROC of 0.658, and the optimal model achieved superior performance with an AUROC of 0.945 (Figure 1e). At the patient-level, when a probability threshold of 0.5 was used to classify cells as leukemic blasts (see Methods), the blast percentages that Mixture Lasso assigned to each patient in the dataset correlated strongly with pathologist-assigned blast percentages (Pearson $\rho$ = 0.88 compared to the optimal model’s $\rho$ = 0.95). The naive model’s blast percentage assignments correlated less with the pathologist gold standard (Pearson $\rho$ = 0.65).

In their original study, Tsai et al. identified two sets of markers that could be used to distinguish between different stages of myeloid development in healthy and AML cell populations. Among these were the cytoskeletal proteins $\beta$ actin, Lamin A/C, and Lamin B1; the granule-associated proteins SerpinB1, VAMP-7, ribosomal RNA (rRNA), Lactoferrin, and lysozyme; CD45; and the cell size marker wheat germ agglutinin lectin (WGA). Two of the most important markers for the pathologists’ AML blast enumeration were rRNA (highly expressed in AML blasts due to their rapid growth) and Lactoferrin (a granule protein only expressed in mature myeloid cells that remains lowly expressed in undifferentiated AML blasts). Figures 1b-c show that Lactoferrin, Lamin A/C, CD45, and rRNA were selected by both the Mixture Lasso and optimal lasso models, whereas SerpinB1 and WGA were uniquely selected by Mixture Lasso. However, Lamin B1 and $\beta$ actin were uniquely selected by the optimal lasso model, and only the optimal model selected rRNA with a positive coefficient. Thus, Mixture Lasso closely follows, but does not entirely reproduce the pathologist’s decision-making process, suggesting an ability to augment (not just replicate) existing clinical decision processes.

To visualize Mixture Lasso’s predictions, we constructed several uniform manifold approximation and projection (UMAP) [11] plots using all cells and markers in the dataset (Figure 2a-e). We find that cells close to one another in UMAP space tend to have similar Mixture Lasso probabilities, indicating that Mixture Lasso tends to identify phenotypically coherent cell populations. Likewise, regions of UMAP space with high Mixture Lasso probabilities tend to correspond to regions containing large numbers of pathologist-identified AML blasts. However, Mixture Lasso does not simply assign high probabilities to all cells collected from cancer patients; rather, it assigns high probabilities to cells in areas of high-dimensional phenotype space preferentially occupied by cells from cancer patients, but not from healthy patients. Intuitively, this makes sense: leukemic blasts should map to regions of phenotype space with little to no overlap with phenotypes present in healthy samples. This observation is quantified in (Figures 2d,f).

Protein-specific UMAP plots further indicate that Mixture Lasso assigns high disease-association probabilities to cells with unique marker expression patterns–generally, those low in Lactoferrin and high in SerpinB1, CD16, and/or CD56. Furthermore, while rRNA was not explicitly selected as a disease-associated feature in the Mixture Lasso model, cells with high probabilities also tend to express high levels of rRNA (Figure 5 and Figure 6).

Next, we sought to evaluate MMIL’s performance in the setting where both labeled and unlabeled cells are present–e.g. a scenario in which most leukemic patients in a dataset do not have cell labels, but a small number of them have been annotated. We hypothesized that a model capable of leveraging both the expert labels of a clinician and a larger corpus of unlabeled data would outperform models trained using only labeled or only unlabeled data [12].

However, using expert labels to train a model on a difficult classification task can be risky: clinical misidentification of leukemic blasts is a known challenge in hematopathology, with a high degree of documented interobserver variability and methodological disagreement (Supplemental Figure 7) [6]. This is particularly concerning for training machine learning models, as the degradation of model performance due to label error is a widely characterized problem [13, 14]. Thus, we also hypothesized that models trained on both labeled and unlabeled data would be more robust to label error than models trained on labeled data alone, as they would be less susceptible to over-fitting on the imperfect labels.

To test these hypotheses, we developed a semi-supervised version of MMIL (“1-shot MMIL”) using the training procedure illustrated in (Figure 3a). Briefly, 1-shot MMIL is trained identically to MMIL with one change: during model training, a single AML patient (termed the “1-shot patient”) is chosen and their gold-standard cell labels are used for supervision instead of probabilistic labels. Like the cells from healthy patients, the 1-shot patient’s cell labels remain fixed throughout all iterations of the EM. Otherwise, model training and LOOCV are carried out as before.

To benchmark 1-shot MMIL’s performance, we compare it to 0-shot MMIL (the usual MMIL), the 0-shot naive model (a naive model trained as described before), the 1-shot naive model (a model in which the 1-shot patient’s gold-standard cell labels are used during model training, but all others use their inherited patient labels), and the 1-shot optimal model (a model trained on only the 1-shot patient using their gold-standard cell labels). Furthermore, to benchmark 1-shot MMIL’s robustness to imperfect labeling, we again fit each 1-shot model after randomly permuting 25% of the 1-shot patient’s cell labels–a proportion chosen to match the variability observed between pathologists when enumerating leukemic blasts in AML (Figures 7 and Methods). All models were lasso-regularized logistic regression classifiers.

We performed 13 1-shot experiments: one for each AML patient in the dataset to serve as the 1-shot patient. The results of these experiments are summarized in Figure 3b-e. In the “0-shot” case, we once again observed that Mixture Lasso outperformed the naive model (Figure 3b, left). In the 1-shot case, we found that Mixture Lasso’s performance improved substantially compared to 0-shot Mixture Lasso (average AUROC increase of 0.058) while maintaining its superior performance to the 1-shot naive model (Figure 3b, middle). We also found that training an optimal model using only the 1-shot patient outperforms 1-shot Mixture Lasso; however, after training labels were permuted, 1-shot Mixture Lasso outperforms both the naive model and the optimal model, whose performance degrades substantially after training on the imperfect labels (Figure 3b, right). Together, these results suggest that MMIL is capable of improving its performance using even a small number of gold-standard labels, while remaining more robust to noisy labeling than alternative models.

Based on these results, we wondered how the inclusion of the 1-shot patient during model training affected the feature sets selected by Mixture Lasso. To interrogate this, we examined the 1-shot experiment with the largest average AUROC improvement between the 0-shot and 1-shot models across all patients (patient AML-5An was the 1-shot patient). In this experiment, rRNA was selected as the most cancer-associated feature by the optimal model (Figure 3c). Interestingly, rRNA was not identified as a cancer-associated marker by the 0-shot Mixture Lasso model, but was identified by the 1-shot model (Figure 3d). In addition, several markers (CD64, CD3, CD8, and CD19) that were included in the 0-shot Mixture Lasso–but that were not used in Tsai et al.’s manual gating of leukemic blasts–were removed in the 1-shot model. This held as a general trend across all 1-shot experiments, across which rRNA was never selected as a cancer-associated marker by 0-shot Mixture Lasso, but was selected as a cancer-associated marker in 10/13 of the corresponding 1-shot Mixture Lasso. At the same time, markers CD3, CD5, CD15, and CD11c were often selected by 0-shot Mixture Lasso but are not canonically considered markers of AML blasts, and these were selected less frequently by 1-shot Mixture Lasso. Taken together, these results demonstrate that MMIL can adapt its feature selection process via semi-supervision, leveraging even a small number of expert-labels to prioritize clinically relevant markers over non-informative ones.

Finally, we applied MMIL to a dataset of 1.1 million blood and bone marrow cells collected from three Acute Lymphoblastic Leukemia (ALL) patients and three healthy controls. This dataset contains samples collected from multiple tissues (blood and bone marrow) and multiple timepoints throughout treatment (diagnosis, Day 8 post-treatment initiation, and Day 15 post-treatment initiation). Using this dataset, we wanted to evaluate the ability of an MMIL model trained at diagnosis to accurately identify and track residual leukemic cells during treatment progression.

The identification of residual leukemic cells throughout treatment is crucial for assessing a patient’s risk of relapse and therefore for clinical decision making. Thus, MMIL’s ability to maintain strong performance despite distributional shift–i.e. a shift in the distribution of leukemic blast phenotypes over time due to treatment effects or changes in disease state–would be paramount for its application to disease monitoring in cancer.

To evaluate MMIL’s ability to detect residual leukemic cells over the course of treatment, we analyzed all cells in the ALL cohort for the presence of 29 proteins as previously described using CyTOF (see Methods). Optimal, naive, and Mixture Lasso models were trained using only bone marrow cells from the diagnostic timepoint, and model performance was once again evaluated using leave-one-patient-out cross validation. Next, we applied our diagnostic models to paired blood samples from the diagnostic timepoint to evaluate each model’s ability to generalize to a new tissue context that is more easily monitored clinically. In addition, we applied each model to paired blood samples taken at days 8 and 15 post-treatment initiation to evaluate its ability to track a patient’s disease burden over time.

Trained without access to gold standard labels, Mixture Lasso demonstrates robust performance and generalizability, with an AUROC of 0.941 at diagnosis and 0.815 at day 15 post-treatment initiation. The naive model begins with AUROC of 0.923 at diagnosis, but this degrades across time to 0.662 at Day 15 (Figure 4). This decline in performance aligns with clinical challenges encountered over the course of a patient’s treatment, including the diminishing proportion of leukemic blasts over time (as low as 0.1-5% of all cells) and phenotypic plasticity in resistant cells post-therapy [15]. Notably, the naive model’s brittle performance in response to these shifts highlights its susceptibility to false positives. Mixture Lasso, however, retains strong performance even as the percentage of leukemic cells becomes low, suggesting its potential for robust disease monitoring and minimal residual disease detection in leukemia.

As before, we analyzed Mixture Lasso’s coefficients to examine which protein markers best distinguish between healthy and leukemic cells. Among the markers with positive (cancer-associated) coefficients identified by Mixture Lasso (Supplemental Figure 9) were CD10, CD19, PAX5, and CD34—each of which has been previously described as aberrantly overexpressed in ALL [16]. For example, CD10, or “Common Acute Lymphoblastic Leukemia Antigen”(CALLA), is a surface protein expressed in primitive, undifferentiated B cell precursors as well as in the majority of ALL cells [17]. Mixture Lasso also identified CD58, a well-known marker of clinical MRD [18]. Conversely, all the proteins with negative coefficients in the Mixture Lasso model are phenotypic markers of mature B cells — such as surface immunoglobulin (IgMs), intracellular immunoglobulin M (IgMi), CD20, and the lambda chain of the B cell receptor (Supplemental Figure 9). Each of these proteins are expressed primarily during stages of B cell development that are not reached by leukemic cells, whose differentiation is blocked at earlier stages of B-cell development  [7].

In multiple instance learning (MIL), data consists of labeled bags, each with many unlabeled instances. Negative bags have only negative instances, and positive bags have at least one positive instance. In our case, the “bags” are people, each with many unlabeled cells; healthy people have only healthy cells (negative instances), and sick people have at least one diseased cell (positive instance). MIL methods can be broadly categorized into three paradigms [21]: instance space, bag space, and embedded space. Bag and embedded space methods aim to label bags; the instance space paradigm, like ours, aims to label instances, and so is our primary focus. For example, multiple instance SVM (mi-SVM) [22] performs SVM with the goal of assigning all negative instances to one halfspace and at least one instance from every positive bag to the other. This method trains an SVM iteratively: first, instances are assigned labels and an SVM is fit. Then, using the fitted SVM, instances are relabeled and a new model is trained. This is repeated until the labels stabilize. Another related approach is multiple instance logistic regression with a lasso penalty (MILR) [23], which treats the instance labels as missing data, and uses expectation maximization to optimize a lasso-penalized binomial log likelihood. This is the same general approach as ours. However, in the expectation step, MILR uses the joint probability of all instances in the same bag; as a result, MILR tends to be biased when the number of instances in each bag becomes large. Our approach, Mixture Modeling for MIL, instead treats instances as the primary objects of study, and does not use their bag-level grouping, thereby avoiding issues related to computing joint probabilities of many instances. This is particularly important in medical settings where we may sample tens or hundreds of thousands of cells per person. Finally, there are many neural network architectures designed for bag space multiple instance learning [24] that can be used to label cells. One common method is to apply attention [25] to simultaneously label patients and identify the cells responsible for the predicted bag label.

Our setting is also analogous to the presence only or positive unlabeled settings, where some data points have known positive labels, and the remaining labels are unknown. We may instead refer to our data as negative only data: cells from healthy people are healthy, and therefore have known negative labels, and cells from sick people are unlabeled and may have positive or negative labels. In presence only data, unlabeled instances are randomly sampled from the population, and

In our setting, though, unlabeled data have a higher-than-baseline probability of belonging to the positive class. The knowledge that a cell is from a sick person raises the probability that it is diseased:

Ward et al. [26] present a method to train a classifier with presence only data using EM. Mixture modeling for MIL is a modification of their approach to accommodate this structural difference in our data.

Mixture discriminant analysis 2 (MDA2) [20], a clustering method, is related to the problem at hand. MDA2 considers data from $K$ classes. Each class is a mixture of Gaussian distributions, the centers of which are shared across all $K$ classes. Here, we have $K=2$ classes (healthy and sick people). The healthy class is composed of just one distribution: the distribution of healthy cells. Because sick people have both diseased and healthy cells, the sick class is a mixture of two distributions: the healthy and diseased cell distributions. While MDA2 is a natural fit for our setting, it is a generative method, and it assumes the features $X$ are normally distributed in order to model the joint likelihood $P(X,\text{cell labels $y$})$. This is not always appropriate, and we may instead prefer a discriminative method that models the conditional likelihood $P(y\mid X)$. We may also prefer to use classifiers that identify important features (like lasso and ridge regression), or relationships between features (like gradient boosted trees and neural networks).

Alternately, a simple but problematic approach is to train cell-level classifier using patient labels in lieu of cell labels [27, 28]. In our comparison, we term this the “naive” approach. This may work when most cells sampled from sick people are diseased, as the inherited label is correct most of the time. When sick people have a smaller fraction of diseased cells, classifier performance will suffer because of high label noise. Ward et al. [26] further show that using this approach with logistic regression causes coefficients to be biased toward 0, thereby hindering discovery of biologically important features. To directly address label noise, approaches from the denoising literature (e.g. [29, 30, 31]) may be useful. Still, our method leverages the knowledge that cells from healthy people have known labels, which is not assumed by more general methods.

We wish to simulate a dataset with size $n=1000$ and $p=100$ such that $\rho=\zeta=0.5$, and the relationship between $x_{i}$ and $y_{i}$ is given by ${\text{logit }P(y_{i}\mid x_{i})=\beta^{T}x_{i}}$. In each simulation, the first $10$ coefficients of $\beta$ are drawn from a normal distribution with mean and standard deviation 1, and the remainder are 0. Then, $x_{i}\in\mathbb{R}^{100}$ is sampled from a standard multivariate normal distribution with identity covariance; $y_{i}$ is from a binomial distribution defined by $P(y_{i}\mid x_{i})$. Instances with $y_{i}=1$ (diseased cells) necessarily have $z_{i}=1$ (sampled from sick patients), and instances with $y_{i}=0$ (healthy cells) have $z_{i}=1$ with probability $\frac{\rho\zeta}{\rho\zeta+(1-\zeta)}$. We repeat this process until we have $1000$ instances ($500$ train and $500$ test) satisfying the requirement $\rho=\zeta=0.5$.

Then, as before, we trained the optimal model, the naive model, and Mixture Lasso. In addition, we also trained a multiple instance SVM model (mi-SVM, described above). We are interested in both prediction and inference, so we compare the models in terms of AUROC, AUPRC, expected calibration error, and the sum of absolute differences between $\hat{\beta}$ and $\beta$. For the optimal model, calibration is done with Platt scaling using the true labels $y$; this represents the best performance we could reasonably expect. The naive model is calibrated using the incorrect labels $z$, to stay consistent with the fact that it was trained using the incorrect labels. The Mixture Lasso and mi-SVM models are calibrated in two steps: (1) obtain held-out predictions for each point in the training set (or use a separate held out dataset) and (2) fit a Mixture Lasso model using only the held out probabilities to predict the cell labels $y$. This results in a logistic regression model that adjusts the predicted probabilities to be better calibrated. This simulation is repeated $1000$ times. For all objectives, the optimal model (trained using the true labels) performs best. When the true labels are unavailable, Mixture Lasso outperforms the naive approach (Table 1), and performing Platt scaling using Mixture Logistic Regression greatly improves calibration.

To choose hyperparameters (e.g. the regularization strength in lasso logistic regression), we recommend using cross validation. For each desired choice of hyperparameter, train a model using cells from a subset of the patients, and record the value of the observed data log likelihood (in Section B) on the cells from held out patients. After repeating this across all desired hyperparameter values and all folds, choose the hyperparameters that maximize the cross validated log likelihood.

Importantly, we have assumed that $\rho$ and $\zeta$, the proportion of healthy cells in sick patients and the proportion of cells sampled from sick patients, are known. We make these assumptions because $\rho$ and $\zeta$ are not, in general, identifiable. Ward et al. [26] show that they are identifiable only if we make strong assumptions about the form of $\eta(x)$; even when it they are identifiable, their estimates have high variance. This argument is reaffirmed and expanded upon by Hastie and Fithian [33].

For large datasets, the EM procedure may be slow and it may be more appropriate to optimize the observed log likelihood (Appendix B) through stochastic gradient descent. For smaller datasets, using EM allows us to easily try many different model choices for $\eta$, and to take advantage of off-the-shelf model fitting software in each loop of the algorithm.

As we have described it here, our algorithm relies on the use of classification software that can train a model using soft labels (values between 0 and 1). In Appendix C, we share a modification for classifiers that require hard labels (restricted to be 0 or 1).

This research complies with all relevant ethical regulations. The use of the primary samples was approved by the Institutional Review Board at Lucile Packard Children’s Hospital at Stanford University. Healthy human bone marrow samples (n = 3) were purchased through AllCells. De-identified bone marrow samples from pediatric patients with BCP-ALL were obtained, under informed consent, from the Pediatric Clinic of University of Milano-Bicocca (Centro Maria Letizia Verga, Monza, Italy; n = 3). The use of these primary samples was approved by the Institutional Review Boards at both institutions. Written informed consent was obtained from the parents of the patients or their legal representatives, who agreed to the use of biological material for research and clinical studies. To protect patients’ privacy, samples have been de-identified.

EC is supported by the Stanford Data Science Scholars program and a Stanford Graduate Fellowship. MZ is supported by the Stanford Bio-X Bowes Graduate Student Fellowship. TH was partially supported by grants DMS-2013736 from the National Science Foundation, and grant R01GM134483 from the National Institutes of Health. RT is supported by the National Institutes of Health(5R01EB001988-16) and the National Science Foundation(19DMS1208164). TJK is supported by the National Institutes of Health [grant number 1F31CA239365-01], the Mark Foundation for Cancer Research, the Andrew McDonough B+ (Be Positive) Foundation, and a Point Foundation Graduate Student Scholarship. JS is supported by Associazione Italiana per la Ricerca sul Cancro (AIRC; Start-UP grant no. 27325). GPN is supported by the National Institutes of Health [grant numbers U54CA209971, U54HG010426, and U19AI100627]. KLD is supported by the National Institutes of Health [grant number R01 CA251858-01A1], the Mark Foundation for Cancer Research, the Andrew McDonough B+ (Be Positive) Foundation, the Oxnard Foundation, and the Stanford Maternal and Child Health Research Institute.

Acute Myeloid Leukemia (AML) mass cytometry data are availability on FlowRespository, with Accession ID FR-FCM-Z2E7.

Acute Lymphoblastic Leukemia (ALL) mass cytometry data are available on Dryad, with DOI 10.5061/dryad.8gtht76vw.

An R package to train mixture classifiers can be found on Github in the repo [XXXXXXXX].