Scientist Stories: Kasia Kedzierska, Limits of zero shot foundation models in single cell biology

Hello, my name is Kasia Kanjirska, and today I'll be talking about assessing the limits of zero-shot foundation models in single-cell biology. This work was carried out during my internship at Microsoft Research in England under the mentorship of Lauren Crawford, Ava Amini, and Alex Lu. When we dissect a tissue, we get a set of different cells, and before the advent of single-cell technologies, we'd only be able to get a readout across average of those cells. But nowadays, we can get a more detailed picture. For example, we can run a single-cell RNA-seq experiment and get an expression profile of each of the cells separately. Let's look at this matrix that is an output of the single-cell RNA-seq. We have cells in rows and genes in columns, and each position of that matrix will tell us about the level of expression of a given gene in a given cell. However, this matrix is quite an idealized version of a real-life one. In real life, the data is much noisier and sparser, and the challenge lies in uncovering the patterns in the data. We could imagine having a model that would take that matrix and create an embedding space where we can project the data. And what we would see, we would see, for example, those 4 clusters, and we can get excited about potential exploring the differences between those clusters and learning about the biological differences. However, what usually happens is that batch effects compound the data and And if we would not know about the batch effect in this data, we would focus on the technical differences rather than the biological ones. And the goal is to learn about the biological, about cells and their identities. So what do we want? We want a model that would create an embedding space, a different embedding space such that would reflect the biological differences in our data. And recently, due to popularization and success of systems such as ChatGPT, a new type of model became popular within the single-cell biology community. The foundation models claim to learn universal patterns from observing vast amounts of data during pre-training—tens of millions of cells spanning different tissues, tissue types, organs, and cell types. Those models claim to serve as a starting point for various analyses such as labeling of cell types, gene prediction, gene expression prediction, and perturbation analysis, and so on. Those models include scBERT, GeneFormer, scGPT, scFoundation, and many more. They are quite popular within the field because of their potential automatization of tasks such as cell type labeling, comparing different datasets, and so on. However, because those models are quite new, there does not exist an established evaluation setup, and most of the work so far has focused on fine-tuning. What does fine-tuning mean? It means that when we take the model and either retrain all of it or parts of it on a smaller dataset, whereas where we take zero-shot, we take the model as is. Evaluating the fine-tuned model also comes with potential setbacks. First of all, we need a set of clear labels, which does not always come easily in biological exploratory work. And secondly, it requires a complicated setup or extensive GPU resources, whereas zero-shot is accessible. Additionally, if we assume that the models learn the biological patterns during their pre-training, those should be accessible in a zero-shot setting. And that's why we decided to evaluate those models in zero-shot scenarios. We used— we took two models. We evaluated two models, Geneformer and SCGPT. They both draw inspiration from the foundation models in natural language processing, but they have some differences. For example, they differ in how they take the input. Geneformer takes the ranked list that is created based on the expression for each cell, and SCGPT bins the expression. The binned expression in SCGPT allows comparison across the cells directly. When they embed the values, Geneformer only embeds the gene IDs and SCGPT embeds the gene IDs and the expression value directly, and then sums those embeddings in embedding space. When it comes to architecture, SCGPT is bigger. And when it comes to training objective, Geneformer uses masked language modeling. What it predicts is the gene at a given position. And SCGPT predicts the expression of masked genes from unmasked genes and cell embedding in one pre-training objective. And the other part of the pre-training objective is the prediction of the gene expression from cell embedding. And when it comes to the cell embedding, in Geneformer, that's average over gene embeddings in XCell, and SCGPT learns the cell embedding directly as a first token. In our evaluation setup, we looked at 5 datasets. We looked at the pancreas that spans one tissue type and several batches. We also looked at PBMCs that have one or two technical batches. And then we looked at more complex datasets such as cross-tissue immune and Tabula Sapiens, which span multiple tissues and multiple technical batches. We also evaluated the models on 3 tasks: cell type clustering, when we ask the question how well the cell types are separating within the cell embedding space. We also asked how well the batches are integrated within the same cell embedding space, and we investigated the pre-training objective. And now moving on to the results. First, I'm going to talk about how well the cell embedding space separates the cell types. And for that, I will use two scores: average silhouette width and average BIO scores. Silhouette score is going to tell us about how distinct and how separate the cell types clusters are in the cell embedding space. And average BIO is a composite score that takes silhouette score and two other scores and also talks about how well the clustering can be performed within that space and how well those clusters align with the cell type labels. For both scores, the higher the value, the better the performance. And as I said, we will evaluate them on 5 different datasets. And to give us some perspective, we will evaluate and compare GeneFormance and GPT-Human performance to a simple heuristic method of choosing highly variable genes. And also to using SCVI, a variational autoencoder method that is commonly used in single-cell RNA-seq data analysis. Then we calculated the median across all the datasets. And what we see is that SCVI outperforms foundation models in both scores, in silhouette score and in the average BIO score. Additionally, what we see is that highly variable genes outperform the foundation models in the average BIO. Composite score. That is the same for the pancreas, which follows, which largely follows the median. However, when we look at the PBMC datasets, that is a little bit more complicated. And let's focus on the PBMC scores for SCVI and SCGPT-Human. So what we see is that when it comes to the silhouette score in both datasets, the SCGPT outperforms the SCVI. However, when we look at the average BIO score, we see that SCGPT-Human outperforms SCVI in one of the datasets, but underperforms SCVI in the other dataset. And also what we can notice is the spread of the values. SCVI is quite concise within the values, and SCGPT has quite a big spread of those values for PBMC datasets, suggesting that it's somewhat unreliable for the dataset coming from the same tissue but different data sources. The more complex datasets also largely follow the medians. And generally what we can see here is that SCVI and highly variable genes outperform foundation models in cell type separating. Now let's talk about batch integration. Here we're going to focus on the pancreas dataset. And as I told you, it spans one tissue type and several different batches. In a perfect batch integration scenario, we would see clusters separating based on cell types and lack of organization due to the batch type. And when we look at the highly variable gene, we see an improvement over the space generated by all genes. We see some clusters based on the cell type and some mixing of the, of the batches. However, when we look at the SCVI, we see very nice pronounced clusters based on the cell type and However, we still see some organization with respect to batch within the clusters. However, when we look at the GeneFormer, we see that the primary driver of the organization in our data is actually batch type. And then there is a lot of mixing within the cell types. In scGPT, we see clusters based on batches, but then some additional organization based on cell types. Overall, we see that foundation models did not perform well at batch integration. However, this is just a UMAP projection. This is just evaluating this and giving us intuition. But to formally evaluate the performance of those models, we looked at the batch integration score. Batch integration score is calculated on the embeddings directly from the models or the methods. Batch integration score is the modified silhouette score where the higher value would signify a better mixing of the batches. And we evaluated the models on 4 datasets: pancreas, PBMC, immune, and tabula sapiens. And again, we calculated the median across all the datasets with the dashed line. And first thing we see is that highly variable genes outperform other models especially the foundation models. And second thing that we see is that, for example, in the case of the GeneFormer, there is unreliability of the score on different datasets. There is a high difference between— there is a big difference between the lowest score and the highest score within the 4 datasets, meaning that we can't really be sure that if it performed well for one dataset that it will also perform well in another. Overall, highly variable gene selection outperforms more complex models such as foundation models. So we saw that the foundation models that we evaluated underperformed in our previous evaluation metrics, and we wanted to see what happens in the pre-training. So first, we're going to look at SCGPT. And as I mentioned earlier, SCGPT has two training objectives: gene expression prediction when the masked gene expression is predicted from unmasked genes and cell embedding. And we're going to look at the input expression bin and predicted expression bin. And what we see that in this first pre-training objective, SCGPD actually outputs an estimate similar for all of the input expression bins. That improves when we take the gene expression prediction from the cell embedding, GPC, with the more variability across the predictions. However, to get a more formal evaluation, we looked at the mean squared error (MSE), and to give us a baseline, we calculated the mean input expression bin for each gene across all cells and used that as a reference. So here we plotted the— in a dashed line, we plotted the median across all the datasets for the mean bin prediction. And if the model will perform better, those values are going to be below the line. And if the model is going to perform worse, the values are going to be above the line. And what we see, as expected, SCGPT random, which is a random initialization of the model, performs worse, with GBC performs worse than GEP. However, when we look at the SCGPT kidney and SCGPT blood, which are variation of SCGPT pre-trained on kidney dataset and blood datasets, We see that those values improve over a random model, and we also see that the GEPC objective outperforms GEP, which is an inverse of what we saw in SCGPT-random, suggesting that the model learns the cell embedding. We also see that only GEPC objective for BLAT and HUMAN outperforms mean, and only in several datasets, not all. When thinking about Geneformer, because of the nature of the output of the model, we will be comparing the input expression rankings with the output, with the predicted expression rankings. And what we see is that we have a lot of points on the diagonal, suggesting that Geneformer performs well at predicting those values for several, for quite a lot of genes. But we also see a lot of points on the bottom of the plot and on the left, which suggests that Geneformer also tends to predict genes that were not present in the input ranking or not predicting the genes that were present in the ranking. To formally evaluate it, we will be looking at the Pearson's correlation between the input and the predicted expression. And to give us a baseline, we calculated the average ranking. So for each gene, we calculated an average position in the ranking rankings present in the input. And then in the dashed line, we have the median for all those datasets in the average ranking. And again, and here when we're going to see the values above the line, that means that the model performs better than the average ranking. And if the values are going to be below the line, that means that the model performs worse. And what we see is that Geneformer outperforms the average ranking in PBMCs, but underperforms in other datasets. In summary, the cell embeddings generated by the foundation models struggle with cell type separation in zero-shot scenarios. They are susceptible to batch effects, and that the current approach to pre-training may offer limited advantages for model performance. I invite you to read more about our research in our preprint, Assessing the Limits of Zero-Shot Foundation Models in Single-Cell Biology, available on bioRxiv. And I would encourage you to see the version 2, where we updated quite a bit of things. And I would like to take this moment to thank, uh, to thank Alex, Ava, and Lauren. I loved working with you, uh, and I would like to thank the whole of the MSR New England lab, uh, and the interns cohort of 2023. And to you for listening, thank you very much, and I'm happy to take your questions.

I'm here.

I don't know if you can hear me. Yep.

Can you hear us, Kasia?

Yes, I can. Great.

Questions? That was a really great talk. I guess your talk and the previous talk highlighted how the current pre-training tasks might be too trivial and not super informative. So I was wondering if you had any thoughts about what might be more appropriate tasks, more challenging and more informative.

Yeah.

So generally, I think that first of all, we should actually evaluate those models on the tasks that we presented and what Rebecca presented as well. And it doesn't mean that the models need to perform fantastically, but we need to understand what they actually learn in a pre-training, right? And we need to find out whether they can perform at those tasks and what they cannot do and why is that and what maybe they are missing at the pre-training stage, right? Other I like this thing that Rebecca mentioned that some people are saying that the foundation models are actually hammers without the nail. And I think that the onus is kind of on the creators of those models to tell us what they actually are gonna be good at. I'm not saying that they are not gonna be good at those tasks. It's just, it's hard to imagine what they can do in the future. And I think that, so there are a couple of other foundation models that I would like to evaluate as well, like Universal Cell Embeddings that have been published a couple of days ago. And I think this one has a good chance of performing well in the tasks that we have presented. So I still think I would argue that the tasks that we have selected are still a good evaluation to get a sense of how well those models perform.

One more question.

Hey Kasia, great talk. I had a quick question about the pre-training steps in terms of training on sparse data. Do you think that maybe using pseudo-bulk and using that as a data augmentation step might help learn better embeddings, especially with masked reconstruction?

Yeah, fantastic question. I actually think that additional work to be done with the data will improve those models further because I think in biology, it's a little bit harder to use the same concept that people use in large language models when you just train on a lot of data and then come up with important and interesting patterns. But I think in biology, you need to do some kind of data curation or some kind of data preprocessing and make sure that the data fits better in the task that you— the models or the tasks you want to use on it. So I think, yes, maybe, maybe pseudo-bulking would improve. I think that Improving the quality of the data and making sure that the data is processed in the same way to eliminate some of the potential drawbacks there would be a good first step to see whether that improves for sure.

Okay, great.

Let's thank Kasia again for the great talk.

Thank you very much.

Okay, so for the last talk before our break, we're going to hear from Eli Bogart, from a lead scientist from Deep Genomics. Deep Genomics has actually sponsored MLCD for a number of years now, and we're thankful for their continued support.