Scientist Stories: Rebecca Boiarsky, Single Cell RNA Sequencing Foundation Models

Hi, I'm Rebecca Boyarsky, and I'll be presenting our work, a deep dive analysis into single-cell RNA sequencing foundation models. I worked on this together with Nalini Singh and Alejandro Buenzia and was advised by Gaddy Getz and David Sontag at the Broad Institute and MIT. Recently, there's been a boom of foundation models for single-cell RNA sequencing data, such as SCBERT, SCGPT, and GeneFormer. These models take inspiration from the success of foundation models in other areas of machine learning, like natural language processing and machine vision, and the big idea is that they train transformers to learn rich gene and cell embeddings using vast amounts of unlabeled single-cell RNA-seq data, and then fine-tune the models to excel at downstream tasks of interest. These models have the potential to help us discover rare cell types, learn gene regulatory networks through the transformer's learned attention weights, overcome issues like batch effects, and excel at downstream supervised tasks like cell type annotation, and perturbation response prediction, even when limited training data is available. In the work I'll present here, we probe the capabilities of these models to deliver on some of these promises. Before I jump into our analysis, I want to give a quick high-level overview of how transformer models are being used to model single-cell gene expression data. So each cell is represented in the model as a collection of genes shown here. Each box represents a gene, and there are two pieces of information that the model needs to know about each gene. First, which gene is it represented in green and called gene identity embeddings throughout this talk. And second, how highly expressed is that gene, represented in blue. And the expression of each gene is binned as a preprocessing step. These gene inputs are fed into transformer layers, which are the major building block of recent large language models. The transformer layers apply attention in order to learn contextualized gene embeddings, in which each gene's representation is influenced by its cellular context, where context refers to the expression of all the other genes in that same cell. So how are these transformer layers trained? The details vary between methods, but some sort of masked pre-training is employed. During pre-training, the expression of some portion of genes is masked, meaning hidden from the model, and the model's trained to use the context of all the other observed genes in order to predict the expression of the masked genes. In theory, this should force the model to develop a good understanding of gene-gene relationships in order to be able to excel at this task. And the weights learned during pre-training should provide a good initialization of the model for all the downstream tasks we might be interested in. Each cell can then be represented either by aggregating the contextualized gene embeddings or through training an additional CLS token. Pre-training is often followed by fine-tuning, in which a small amount of labeled data is used to further train the model to excel at a particular downstream task of interest. This is shown here, where the learned contextualized gene embeddings are used as input to a smaller neural network, which is trained to predict some cell-level label, such as cell type or response to perturbation. The full network, including the transformer, can be backpropagated through during this learning, updating all weights in order to improve the at the particular task of interest. Pre-training and fine-tuning allow us to leverage massive amounts of unlabeled data to help improve performance on downstream tasks, but these models come at a cost. In particular, taking SCGPT as an example and comparing it to a simple model like logistic regression for the fine-tuning task of cell type annotation based on 3,000 input genes, we see that SCGPT has over 51 million trainable parameters, while logistic regression has just 3,001. GPT required training on a GPU, which can be a barrier to entry for many researchers. And for the particular dataset profiled here, training took over 300 times longer and took 17 gigabytes of GPU memory. And this is just for an input of 3,000 genes, but it would be reasonable to want to use all 20,000 protein coding genes in your model, in which case these numbers would all significantly increase. Furthermore, this slide describes fine-tuning, but pre-training is even more expensive in all of these regards. And there are also, of course, environmental costs associated with the energy consumption needed to train these models at a large scale. Because of these costs, it's really important that we understand where these models excel, what they should be used for, and what they're not necessary to be used for. Here in this work, we took a deep dive into two such foundation models, SCBERT and SCGPT. Because of time constraints, I'm not going to explain the differences between these 2 models, but you can refer to the original publications or to our preprint for more details. In particular, we investigated the benefit of the pre-trained then fine-tuned paradigm for single-cell RNA-seq data with a focus on the fine-tuning task of cell type annotation. We noticed that the papers that introduced these models didn't compare their cell type annotation results to one of the most simple baselines. Regularized logistic regression. When we compare against this baseline, we find that more often than not, logistic regression is comparable or outperforms the results of the foundation model. Still, one promise of foundation models is that the representations learned during pre-training might allow you to perform better at few-shot learning, or training with only a limited amount of labeled data. With this in mind, we also tested whether the foundation models would excel as we reduce the amount of training data available. But we still found that logistic regression was often superior in the few-shot setting. Let's take a look at these results. There are 4 main datasets used for the task of cell type annotation in the SBERT and SCGPT papers. SBERT reported results for a 68,000-cell PBMC dataset. Blood cells. SCGPT reported results on 3 different datasets, multiple sclerosis, which contains neuronal cell types from healthy donors and patients with MS. A human pancreas dataset, and a dataset of tumor-infiltrating myeloid cells from 9 different tumors. The size of the train and test sets for each dataset are shown here on the plot. And it's worth noting that the multiple sclerosis dataset is split in such a way that it represents a case of dataset shift where the test distribution is significantly different than the training distribution because all the healthy donors were put into the train set and the MS patients were put in the test set. So this represents a particularly challenging case for doing prediction. Okay, here we have our first set of results. SCBERT compared to logistic regression for cell type annotation in the PBMC data. It's hard to see, but there are multiple dots on the plot for each model representing runs over different random seeds. So models trained from scratch, or fine-tuned from scratch with different random seed settings. We report the macro F1 score, which is the harmonic mean of precision and recall, and which is a good metric for classification when you have unbalanced classes, which is the case in all of the datasets here. As you can see in the plot, and perhaps more clearly in the table, logistic regression outperformed SCBERT for annotating blood cell types. Next, we ran a logistic regression baseline for all the datasets evaluated in the SCGPT paper. Starting with the multiple sclerosis data, we found that SCGPT did in fact outperform the logistic regression baseline. One possible explanation for this is that the multiple sclerosis data is the dataset that suffers from distribution shift. The pre-trained model might actually be uniquely suited to handle this kind of situation since it's seen many different clinical contexts during pre-training, while logistic regression was only trained on a single clinical context that is no longer present in the test set. In other words, it's possible that the pre-training serves as a form of regularization that protects the foundation model from overfitting to the training distribution during fine-tuning and helps it to generalize. For the pancreas data, the logistic regression is competitive with SCGPT. We find that the metrics are fairly well matched between the two models. And finally, for the myeloid data, logistic regression slightly outperforms SCGPT. For our few-shot experiments, we subsampled each training dataset to just 10, 25, 50, or 75% of its original size, as shown on the x-axis of these plots. For the most part, we found that the model which outperformed with the full training data also outperformed in the few-shot setting. For the pancreas data, where logistic regression and SCGPT were very closely tied in the— with the full data, we found that actually logistic regression became significantly better than SCGPT in the few-shot setting. So unlike our hypothesis that the pre-trained models would have a unique advantage for few-shot learning, we found that logistic regression with fewer parameters actually fared better when training data was limited. While logistic regression is competitive for the fine-tuning task of cell type annotation, foundation models still have other potential benefits over logistic regression. One being that they learn gene and cell embeddings, which might capture meaningful biology and can be useful for other tasks. With this in mind, we ran some ablation studies to better understand the representation learning capabilities of the foundation models. First, we tested the effect of removing pre-training altogether to directly evaluate how much pre-training helps with the fine-tuning task of cell type annotation. Here's a visual of the experiment. You'll recall what I explained in the foundation model overview, that in these models, the weights in the transformer layers are initialized via pre-training. Instead, in this experiment, we initialize the transformer to have random weights. And we then proceeded directly to training the network to predict cell types in the different datasets. For SCBERT, we found that there was no difference in performance for annotating the PBMC cell types with or without pre-training. For SCGPT, we found that the pre-training did benefit the model. We need some further follow-up analysis to understand what accounts for the differences in the efficacy of pre-training between SCBERT and SCGPT. Finally, we present a cautionary tale showing that good pre-training performance can be deceiving and does not always mean that meaningful data representations are being learned. To introduce this experiment, I'm going to give a little more detail about the way that input data is processed in the SCBERT model. As we said earlier, each gene is input as a sum of a gene identity embedding representing which gene is which, and an expression embedding representing how highly each gene is expressed. The expression embedding actually represents binned expression data, and so it maps to a categorical value. In the case of SCBERT, they bin their data according to the following 6 bins. And note that these bins are pretty coarse. All gene expression is categorized into one of 6 bins, where UMI stands for unique molecular identifiers, which is a measure of single-cell gene expression. During pre-training, the transformer layers are used to learn nuanced representations of each gene given the context of its cell state. Each input gene is unique in that it's a unique combination of a gene identity embedding and an expression embedding, and it ends up being represented by a unique contextualized gene embedding after passing through the transformer layers. Even the masked genes shown here in black at the bottom, end up having different contextualized representations since the gene identity embeddings differ for the different masked genes. This allows the model to make a unique prediction for the value of each masked gene, as shown here at the very top of the schematic. In our ablation experiment, we removed the gene identity embeddings completely from the input to the model so that the only so that only expression embeddings remain. In this setting, the model has absolutely no notion of which gene is which. And recall that transformers are permutation invariant, so the model can't even use positional information to deduce which gene is which. The only information that the model has about each gene is which expression bin it fell in, one of the 6 expression bins. And therefore, every gene with the same binned expression will have the same hidden embedding output by the transformer. Similarly, every masked gene will necessarily have the same representation output by the transformer, and the model will be forced to make the same prediction for every masked gene in a cell during pre-training. It seems like this should be an impossible situation for the model to do well at pre-training, since it doesn't even know which gene it's making a prediction for. But what we found is that the model actually still does pretty well. While the regular SCBERT model achieves 78% accuracy on the masking task, this severely limited model still achieves 72% accuracy. So we found ourselves wondering, how is it possible, given that the model does not know which gene is which, that it still does fairly well at predicting what the masked expression was? And digging in, we found that in this setting, the model learns to always predict that the masked value is equal to the most common bin for that cell. Since there are only 6 bins, and some bins are more common than others, this leads to a pretty high accuracy of 72%. This ablation study is meant to serve as a cautionary tale that if the pre-training task is too easy, the model can learn simple heuristics to do well at the pre-training without actually learning the rich gene representations that we're after. This underscores the importance of designing challenging pre-training tasks that encourage the model to learn representations that will be useful and informative in the ways that we'd like. Like, and also highlights the importance of carefully introspecting into your model to make sure that it's actually learning what you think it's learning. With these results in mind, I want to wrap up with a few key takeaways from our work. First, it's important to compare deep models to simple baselines to understand for which tasks they're truly needed. Second, since simple baselines are often comparable or better than foundation models at cell type annotation, we need to move beyond cell type annotation to showcase the value of foundation models, which SCGPT, GeneFormer, and other papers are already starting to do. Foundation models might really show their worth when it comes to more challenging fine-tuning tasks, which simple baselines can't do well at, zero-shot capabilities, or in elucidating gene-gene interactions and gene regulatory networks through their learned attention weights. We also showed here that strong pre-training or fine-tuning performance does not necessarily mean that rich data representations are being learned, and this calls for careful introspection into models as well as for careful crafting of pre-training objectives. Finally, we didn't show these results here, but in our preprint, we show how hyperparameter settings and random seeds can have a really big effect on model performance, and therefore, detailed and transparent reporting of all of these settings is critical in order for the community to be able to reproduce and build off of each other's work. Finally, I'd like to thank the following funding sources and people who helped shape this work through feedback and discussion. All of our analysis code is available on GitHub, and we hope you'll check it out.