Scientist Stories: Kexin Huang, Sequential Optimal Experimental Design of Perturbation Screens

Today we'll talk about perturb-seq. I think most of you probably know, but I'll give a very brief overview of what is that. Basically, given a set of gene of interest, we do some perturbation. It can be CRISPRi, CRISPRa, and we then do single-cell RNA readout. And in the end, we'll get a set of outcome, like perturb-seq outcome, where each row is a cell and each column is a gene. And then each value is a gene expression upon perturbation for that cell. And as we all know, perturbation screens are powerful. They answer central questions in biology, such as target discovery, cellular programming. And if we are doing, like, combinatorial perturbation, then it can also ask questions like genetic interaction, like synthetic lethality. But with this great power, it also comes with great limitation. That is, the PerturbSeq is very expensive. And to answer many of these interesting questions, we require a large set of perturbations. And for example, if you want to, as an example, want to do like the entire perturbation space, then it will be like, it's massive. Like if we look at a single gene perturbation, it's 20K. And if we do like double genes, become 40 million and so on. And also we're interested in perturbation in different, like various type of perturbation, like genetic, chemical, and also we're interested in various contexts of perturbation, like different cell contexts, different disease state and things like that. So with this combinatorial explosion, even the largest facility cannot do this kind of thing in the wet lab. So now the question is if we can predict these transcription outcomes of perturbation not seen before. So recently, there's a wide range of work that have been proposed to tackle this problem. So basically, these models, what they do is they taking a pseudobulk of non-targeting cells, and then they pick whatever perturbation that you want to perturb on, and it will predict the post-perturbation gene expression. But the big challenge of these models is that they just simply use the existing training dataset, which is a set of perturbations that is used to answer some specific biological questions. And these biological questions are not designed for machine learning model learning. 'Cause for example, if the dataset is used to answer some question, so the perturbation will be focused on several pathways, it's hard to imagine for the machine learning model to generalize to very distant pathways. So here we're saying that if we're interested in predicting the entire perturbation space, then maybe that's not the way to go. So in this world, we propose a procedure we call iterative perturb-seq or perturb-seq in a loop. Which basically we want to do things iteratively. And we want to let models select perturbations that are most needed for learning instead of just do the perturbation and once for all. Basically, the step is very simple. In each round, we pick a set of perturbations and we do the wild experiments. And then we train the model using all the possible perturbations. And then given the models, we develop a selection strategy that inspects the model's confidence for each perturbation. Of perturbation and each thing and unthink and select what are the promising set of perturbation in the next round to do experiments. And then we can do this in an iterative fashion. And in the end, hopefully we can do as little experiments as possible. And of course, it comes with some realistic constraints. Like for example, each round takes weeks to complete. So we can only do very few rounds. And also another thing is for each round, the number of batch experiments also pretty limited because, as I mentioned, perturb-seq is very expensive. So here, just to give a more concrete taste of what we're doing, basically, we set up an evaluation that simulates the real-world scenario. So we leverage a dataset from RepLogo 2022. There are like 2,000 essential-wide single-gene perturbations. So basically, what we do is for each round, we have a selection strategy to pick what is the most informative perturbation for the model. So we will do like 5 rounds, and then each round we take 100 new perturbations. And we have a— because we're interested in the entire perturbation space, so we have a holdout set, which is randomly picked up around the entire perturbation space. And we use that for each round to measure how the model can predict these holdout perturbations. So we can see a curve as the round goes by, how the model prediction on this holdout set. So if the curve goes really fast, then it means our model is able to learn super quickly and vice versa. Cool. So that was the problem setup. Now I'll talk about our method. So our method, this basically falls into the classic machine learning literature called batch mode active learning. So in batch mode active learning, there are basically 3 principles called informativeness, diversity, and representativeness. So here, just to illustrate, imagine these gray points are all the perturbation that is not being perturbed on. And these red points are the ones that are already perturbed. So to pick informative points, you probably want to pick regions that are really far away from these red points because they are space that the model has not explored before. And for diversity, because we're looking at batch mode active learning, So we don't want to pick repetitive ones. We want to include some notion of diversity into the model. So probably want to pick models that are like here, here, here, here. And the third one is the representativeness. Basically, we want to prefer points that has more higher density compared to points that has lower density. So that's the three principles from batch mode active learning. And luckily, we find this really nice reasoning framework by Holzmüller et al. They propose— they decompose any state-of-the-art active learning method into 3 modules, which is called kernel and transformation and selection rules. For example, kernel is basically defined as similarities between the perturbation that is selected and unselected. So here, for example, one simple way is just to take the predicted outcome and take the dot product. Which measures the similarity between perturbations. And then transformation is in case you want to make it rigorous uncertainty, then you can convert this kernel matrix to an updated kernel matrix that has things you want. For example, if you want to have uncertainty, you can simply do Gaussian process posterior coherence. And the last step is selection rule. So based on this matrix, we can just basically select which perturbation to do next. Here I just give an example. We basically want to greedily select points with maximum distance to all selected points. So this way, in each iteration, we are picking the points that are most informative. And also we're doing it greedily because this way we account for some diversity. So that was just an example, but all of these state-of-the-art active learning methods can be decomposed into some variants of individual modules. And cool. So next, we use this approach and we just test it out. But then we realize there's a significant difference with the classic active learning setting and our active learning setting. That is, in classic active learning setting, where people want to do computer vision, they want to label some image, there are usually tens of thousands of initialized samples where the model is already really good. And also, they are looking at tens of hundreds of rounds. And each round will consist of large samples. But in our case, we're doing extremely limited budgeted regime where we only have hundreds of initialized samples and only have very few rounds. And for each round, we can also only do hundreds of samples. So this has a very bad consequence. That is, the model kernel estimation is really bad because basically it failed to capture the similarities between perturbation, which is hard for the later step to select because if the base kernel is bad, it highly biased the downstream selection step. So we also have an independent data analysis to support this claim. Basically what we did is we basically simply do a vanilla active learning method, which is like uniform sampling. And in each round, we compute the model kernel on the same set of perturbations. And then we also have the ground truth kernel because we have the ground truth for several sets of several purposely dataset. And then we can just calculate the alignment between them. And we find that as round goes by, of course, the kernel alignment goes larger as the model gets better. But the thing is, at the first few rounds, the model kernel is not that good. So now we want to say, maybe you can— how do we tackle this part? So this is our method of motivation. So as we all know, when under resource constraint, prior knowledge is very useful. So because we can really look at vast amount of prior knowledge about perturbation relationship, and hopefully we can enrich this model kernel insufficiency. For example, if we look at there are other perturbistic data in other cellular contexts, then maybe we can guess some perturbation relationship are conserved across contexts. Or there's a reason to achieve genome-wide optical pull screens, then we can hypothesize that maybe the perturbations that have similar morphological changes could have similar perturbation relationship. Or there's also vast literature around each gene and how it affects the downstream effects. And there's knowledge graph capturing about protein-protein interactions, like physical interactions, because for example, if two proteins that are in the same protein complex, then we have a strong prior that they have potentially similar perturbation outcome. And there's also single-cell RNA-seq from vast amount of data and also biochemical structures, like proteins that are similar potentially should have similar functions. Okay, now, okay, we know this. Then how do we incorporate this into the axoneme method and how do we account for this multimodality? Because here we are talking about matrices, tags, graphs, images, and I don't know, geometric objects. So our intuition is we don't care actually about the individual data point. We care about their relationship. We care about the perturbation similarities. So in this sense, we can kernelize everything. We can kernelize, for example, the perturbed sequence, take the readout, and then take the dot product. And for optical pore screening, we can visualize it using a cell profiler, which is an established morphological feature generator. And we can also derive the kind of similarities. And similarly for single-cell RNA-seq atlas, we can do, we basically get a feature from POPS embedding, which is a kind of a method that collects lots of single-cell RNA-seq experiments. And from protein structure, we can do the ESM to get the protein structure. Protein representation. And for the knowledge graph, we can use NootVag. And for literature, we can use DialGBT to capture all the information about this gene. So after we kernelize everything, now we want to incorporate into the active learning methods. Oh, sorry. Before that, we also have a data analysis that showcased that indeed, like, these prior knowledge have really useful information. So we're doing the same analysis, but now we kind of compare against the ground truth with each individual prior knowledge kernels. And we can see, especially in the first few rounds, there are abundant information where the model kernel does not capture, but there's information in the prior information that we can leverage. So using this motivation, we kind of develop a simple method to incorporate them. Basically what we have is we first kernelize everything in the same space and we fuse them with the model kernel at each step. And then we get an updated kernel matrix where it basically captures the better perturbation similarities. And then we can do kind of the similar transformation rules, transformation selection rules and whatever. So here now I'll quickly show some results. So first of all, because this is the first time we try to do active learning for perturbistic data, we also build a first benchmark for that. So we can see that indeed the red one is the random uniform selection strategy, and all the previous active learning methods indeed show some improvement. And then the orange line is our method. We can see that we further surpass this performance. For example, just to give a sense, we achieved a similar accuracy with the uniform sampling with, I think it takes 5 rounds to to achieve the similar accuracy for the random, but we only use less than 1 round. And similarly, for the best active learning baseline, we have more than 3 times speedup. And also, just to zoom in a bit, we also find that the fastest speedup actually happens in the first round, which also matches our data analysis showing that in the first 2 rounds, the information is most useful. And we also check other biological meaningful metrics and see that the trend is conserved. I think due to the sake of time, I'll just quickly walk past this result. And we also did some ablation where we find that if we remove the model kernel, we just use the prior only. And we see that this is slightly worse than integrating both model and prior. But surprisingly, prior, just using prior without the model kernel, it also has really good performance. So we also disentangle each of these prior information. And we see that different prior has different level of usefulness. I think the most useful one are the perturbed speech dataset because as expected, they are potentially very conserved perturbation relationship across different cellular contexts. And the OPS also has lots of rich signals. And then one last thing that we look at is because in reality, when we deploy this system, there will be strong batch effects because in each round, we're doing different experiments. So we did a simulation to account for this batch effect. What we did is, because in the Replog dataset there are 48 lanes, and we can basically, in each round, we just focus on learning or training from preselected A lanes so that each round we're looking at different lanes. And we can see that the model performance also kind of conserved. Cool. So that's all for the talk. And this is our paper. And I want to thank again for the amazing team. Thank you.

We have time for a couple of questions. I was wondering if you thought about incorporating those budget variations into your sampling procedure itself, such that, you know, when you have earlier rounds, you can look for things like be a little bit riskier.

That's me.

Where later rounds, because it's later, you know, you want to go for what you're most likely to find. Have you thought about how to go about doing that at all?

Yeah, we actually tried to experiment on account for this. This low sample regime in the selection strategy. But we didn't look at the way you mentioned, but we look at the— because in the initial first rounds, it's more important to select typical samples instead of really uncertain ones in some sense. So we also develop a strategy trying to account for that. So like a sampling strategy that prefer typicality in the first round, and then typicality just goes down as the later rounds. But we end up didn't use that because it's too complicated. We have this new thing and had this as another thing.

So yeah. I also have a question, actually. So I like the way you incorporate different kinds of prior knowledge using these kernels. And it reminds me of what people did for gene function predictions a long time ago. But I'm wondering, so the different prior sources have different quality in terms of data and also different sparsity in the kernel that you're constructing, does that affect your results and your ablation analysis?

Yeah, definitely. So the kernel that we're looking at right now is different similar contexts and things like that. So it's lots of heterogeneity. And so we did investigate a bit on how do we fuse them cleverly. But actually in the end, we find that simple average seems to have the best performance. Mainly— Yeah. I see. Yeah, that makes sense. I think mainly because when we do really complicated ways to do fusion, it's easily overfit. But the simple one doesn't have that issue.

Thank you. My question is actually related to Asar's question. Have you tracked the relationship between knowledge kernels? Are there some complementary kernels their knowledge?

You already mentioned that the first 6 themselves are the most important. What would be the next 7 important ones given the 6 that are mentioned? Yeah, definitely. I think we didn't expect the relationship looking at the matrix or something like that, but we did have an experimental analysis and we basically showed that maybe the second best one is the OPS. The optical— Oh, I see. So indeed, there's some synergy. For example, when we integrate everything, it has the best performance. But indeed, it's hard to look at each combination of prior knowledge and try different— try this approach. It takes forever. Yeah, we definitely tried different fusion operators, also weighing different modalities. Yeah, because imagine here we are consider everything similarly, but if we were drawing this iterative perturbsy in another cellular context, ideally the set of weights should be different for each kernel. So we tried some approach, but it didn't work out. So actually, for example, we can utilize information from this already kind of experimented data and we can infer some weights about each contribution to each modality. But we just find, and in the end, the simple taking the mean is the best. But definitely it was investigated in the future.

This one. Okay, very nice talk. So based upon your observation, can we make something that the majority of the perturbation combination is additive? But what do we care about? A very small set of non-additive and cooperative combinations? If so, maybe we can break apart the perturbation back into a sum of marginal, second order, third order, and so forth. Will this significantly reduce complexity, for example?

Yeah, I think that's a good question. I think it's like two different goals. In our approach, we're looking at, we want to have a machine learning model that can predict the entire perturbation space. And so that's usually the active learning regime. And another regime is want to optimize some properties, like want to find maybe like a set of perturbation that has most novel effects. And in that way, I think it's more like a Bayesian optimization scheme where we have a slightly different kind of selection strategies. Yeah. Yeah, you want to see if we can get one here?

Oh, and you send in your app, right? Okay, so hopefully it all