Scientist Stories: Farhan Damani, Beyond the training set and detecting distribution shifts

Thanks so much to the organizers. My name is Farhan Damani. I'm a senior scientist at Dynotherapeutics, and today I'll be talking about detecting distribution shift in model-based optimization. Dynotherapeutics is focused on— is there an echo? I think that was my fault. My sound was on.

Hello?

Yep. Okay, great. Okay, where was I? So Dynotherapeutics is a company focused on gene therapy applications. As we heard earlier today, gene therapy is an emerging therapeutic modality with curative potential through the delivery of genetic medicine. Dyno is focused, Dyno uses machine learning to design AAVs to solve the delivery problem in order to enable gene and gene editing therapy applications. The underlying hypothesis of Dyno is that AAVs can be used to deliver genes to specific tissues and thereby avoiding tissues that you do not want to target. Quite amazingly, AAVs, if you modify the sequences, the sequence itself, and you test that in vivo, what you'll see is that different AAVs will go to different parts of the body. The critical implication of that is that these AAVs can enable a broad set of therapeutics to tackle many different types of diseases. AAV design can be cast as a 737-amino acid search problem. AAVs are complex proteins that are assembled through 60 protein subunits that come together. Structure-based and mechanistic insights have provided limited use cases for AAV design, particularly because we don't have a good understanding of the AAV biology or how AAVs deliver genetic material to target tissues, say, in either mice or non-human primates. So as a result, Dyno focuses on using machine learning to model that entire process end to end by coupling high-throughput biology, in vivo experimentation, and machine learning. Here's a schematic describing sort of on a high level what a black box design workflow might look like. The high level objective is to design novel sequences that have improved properties relative to what you might have observed previously. The ground truth function is a complex biological process. In our case, that is the delivery of genetic material to target tissues in vivo. So as a result, and because of resource constraints, this isn't an online learning problem. You don't get many interactions with the system that you're interested in. But instead it's an offline learning problem. You have a static dataset of XY pairs that we assume that you have access to to start with for this setting. And your goal is to do the best you can do to perform optimization, typically in one try. So you start with this training dataset of sequences and measurements of a target property. And there are many things you can do to produce candidate sequences. You can train a regression model, and you can couple that regression model with a discrete search algorithm, say using beam search or a genetic algorithm that can search in protein sequence space to find new sequences that have high predicted scores under the regression model. You can also use generative methods to produce sequences as an alternative. You can use generative models to produce sequences as well. Because hyperparameter selection for search algorithms, is a really hard problem, and there aren't really, frankly, any good methods to do that. In practice, what you do is you typically ensemble over, say, many different design methods and many hyperparameters for each of those design methods. All in all, this gives you a very large batch of candidate sequences and corresponding regression model scores for each of those sequences. While it's easy to generate, say, millions of these candidate sequences, The key challenge is how do you identify a small subset of those sequences that you want to submit for experimental testing? Now, in an ideal world, we would just trust our regression model and our job would be done. We would just sort the list of candidate sequences from highest to lowest model scores. We would take the top case sequences and we would submit those for testing. But we know that models are not going to be reliable in the entirety of the input space. Especially when the space is high-dimensional. So for proteins, that's 20 to the L, where L is the length of the sequence. So we need some way to identify inputs where predictions are reliable. And that's what this talk is about. So to study reliability, we deploy this black-box design approach in the real world. And here in Panel A, I'm showing you model scores from a discrete iterative optimization that goes for 15 steps. Each gray line corresponds to an independent trajectory, and the red line corresponds to the average of those trajectories at each step. The objective here is a regression model that maps from sequence to the target property. So in the in silico sense, if you look at this plot, this optimization has been successful. It correctly increased the score. It found sequences that resulted in an increased regression model score over time, over the course of these 15 steps. Now, one approach would be to simply take the sequences with the highest scores that would be at step 15. But we know that that might not work because our models aren't always reliable. So to test this, we experimentally validate sequences along the entire trajectory in an AV design in vitro experiment. And what we find is that the model error for the regression model increases as a function of the optimization step denoted as iteration here. Another way of looking at error is in panel C, which is the fraction of design variance from a particular step that is entirely nonfunctional. And we see by step 10, step 11, we're converging at a 100% nonfunctional rate. And I want to point out how surprising this might be to some people. In the inset in Panel B, I'm showing the performance of that regression model on a random holdout set. That model is quite good. It achieves a Spearman rank correlation of around 0.7 to 0.8 and MSC around 2. And that same model exhibits pretty substantial amounts of distribution shift as you go further and further from the data according to this iteration. And I want to note that this does not line up perfectly with edit distance. It has to do with some other embedding or model-based distance that results in this catastrophic failure of the design. I'd also like to emphasize that this effect is robust across search algorithms, reward functions, and regression models. So in the context of search algorithms, We tried two variants of discrete optimization. We tried a beam search method as well as AdaLead, which is a genetic algorithm. We also tried commonly deployed heuristics to constrain the reward function, such as geometric constraints, like looking at minimizing the— like constraining based on edit distance. We also looked at using ensemble scores instead of single regression scores. On the regression model side, we tried using CNNs that were trained from scratch on our XY pairs., but we also use pre-trained protein-based large language models coupled with linear probing to predict the target property, and we see basically the same trends. So, so far I've shown you that distribution shift is a big problem in design, and the intensity of this shift is not really seen in commonly used simulation settings. Mathematically, this distribution shift takes the form of covariate shift, Covariate shift mathematically is defined when the marginal distribution of the inputs for the training data does not equal the marginal distribution of the inputs for the test data. We're focusing just on covariate shift, which means that in these design applications, the conditional distribution of the property given the input stays the same because that's always going to be governed by some fixed underlying biological process. To illustrate what covariate shift might look like in a simple 2D example. In panel A, I'm showing this simple nonconvex function with 3 optima, and the ground truth function values are shown in the color map here. All of my training data is on one side. It's all centered around one of the 3 optima. If you train a regression model on these training data points and predict the function value, what you'll see is that model error increases as you get further away from the training data, which is what I'm showing in panel B. So if you don't have proper so-called guardrails on your search, optimization might veer off into regions where there's less data, resulting in high error. So the task we sought out was, can we define a score function that takes as input a particular sequence and outputs a scalar score that tells us whether a prediction for that input would be reliable? Fortunately, covariate shift is a well-studied problem in machine learning, and density ratios in particular have been used for outlier detection. Now, the way it's been used in the past always assumed, or typically assumes, that you have access to a fixed test set. And in the case where you have a fixed test set, you can learn this density ratio, the density of a point under the test distribution divided by the density of that point under the training distribution. But in the case of design, it's not obvious what a test set might be. So one of our key innovations is to use a design distribution as the test set. And that design distribution, if you recall from my earlier schematic, is the distribution of all of the candidate sequences that we designed from an ensemble of methods. We also provide evidence to show that this is a particularly good choice. And it's a critical choice to getting this type of method to work. We compare using the design distribution to two variants of a uniform distribution, uniform distribution and then edit distance constrained version of the uniform distribution. And we look at how well, what is like the discriminative capabilities of this particular score function identifying functional variance. And we see that using the design distribution is absolutely key. I also want to point out that we use the density ratio estimation trick to efficiently estimate this ratio. And the density ratio estimation trick is quite nice because it lets you avoid having to estimate these high-dimensional discrete distributions, which we know is very difficult, and we know the likelihoods under those models don't do a good job of detecting outliers. But instead, the minimizer of the binary cross-entropy equals exactly this ratio. So procedurally, this is so simple you can try it at home. The The idea is you take samples from the design distribution and you assign a positive label to those samples. Then you take samples from the training distribution and you assign a negative label to those. Then you just train a binary classifier to separate those two things, and the outputted probability scores is exactly equal to this density ratio. So now that we have an understanding of the mechanics of how the score function works, we've established that distribution shift is a problem. We developed, we designed an in vitro study to specifically test the validity of this method. At Dyno, we regularly do in vitro and in vivo studies. And one of the great things about Dyno is that we do in vitro studies at least partly to study our machine learning methods. So in this particular in vitro study, there are two properties we care about. The first is packaging. And packaging is a measure of whether the capsid assembles properly. You can sort of think of it as a proxy for stability or protein folding. This measure is a necessary but insufficient condition for optimization because at the end of the day, we're interested in function, and function will come from transduction. Transduction is a measure of how efficiently the virus can transfer the genetic material or the so-called payload to a target cell. So in this case, these are HEK293T cells in a cell culture. Finally, the last final point is that in this experiment, we're modifying a 63 amino acid region, of the AAV capsid. So we do this experiment, and we have these trajectory plots that we're all familiar with now. And in this top panel, what we're showing is we have this discrete iterative optimization. And for points along the entire trajectory where the reward function is a regression model, we measure all of these points in a real experiment. And the fraction of design— the y-axis is the fraction functional, which is the fraction of design variants that package at any given step., and we see a slow gradual decay from 100% to 0%. This is also the model precision for a packaging prediction model. What's kind of nice about this is if you were to have just taken points at iteration 15, nothing would have packaged. If you were to have done some sort of early stopping of your optimization, you would have seen that your, that your design worked really well in the sense that you have a high fraction functional, but you wouldn't have been able to observe this gradual change in the functionality rate based on the intensity. So now that we know that there is this gradual shift, a good score function should reflect that gradual shift. In the second plot, I'm showing our method, what we're calling the log OOD score, does qualitatively follow this pattern. The OOD score gradually increases in shift intensity from 0 to 15. In contrast, we use a baseline here which is considered the state of the art. It's a widely deployed method for uncertainty quantification that people use for selection, which is a deep ensemble standard deviation. You train, say, 5 deep neural network models with different random initializations, and you can look at the uncertainty of those predictions. And we see that those measurements are rather flat across the 15 steps. A more quantitative evaluation here would be to ask how well this OOD score can be used to identify variants that package and which ones that do not package. Packaging here is denoted as functional in red, and non-packaging variants are denoted as non-functional in gray. We see that our method produces an area under curve of 0.83, while the deep ensemble method produces an area under curve of 0.61. So it doesn't really do a good job compared to our method. And I also just want to clarify one more time that the packaging model says that all of these variants package. So the real information here is coming from this uncertainty score. So far, I've shown packaging results, which are a necessary but insufficient condition. Our main result focuses on transduction. So as I described in that schematic earlier, we have a setup where we have a large batch of candidate sequences. We have sequences, and we have the corresponding model scores. And we can only test a small subset of them. So in order to study this problem, we actually experimentally validate everything in the candidate set, which allows us to retroactively study something like regret. So if you were to only have been able to pick a small subset of those sequences, how well would you have done? So regret here represents the difference between the best possible variant in the entire candidate set and the best variant found in your top 100 selection. So the very simple thing to do would be to sort the list from highest to lowest model scores, take the top 100, compute regret. That's the gray dotted line here. Now, there are many ways that you might want to use the log OOD score, and it might depend on the practitioner's utility function. But one simple approach would be to use it as a pre-filter. So you can take this candidate set, you can filter that dataset based on your reliability, your OOD score, and then perform greedy selection as usual on the remaining subset. So let's take the example of the 10th percentile mark. Again, we know high scores correspond to more out-of-distribution, so we want things that are lower scores because those are more likely to be in-distribution. At the 10th percentile mark, we take the 10% of sequences that have the lowest OOD score, and we perform this greedy selection both for the OOD score and the deep ensemble score. And we see across the board that you're going to achieve lower regret with this method. In conclusion, I've shown that distribution shift is a major problem when deploying design in the real world. It's hard to see this in simulation. In simulation settings, but you see it when you do this in practice. We propose a very simple method for detecting the shift, just train a binary classifier. And we demonstrate that these scores are reliable, are a reliable scoring function for identifying varying shift intensities. I'd like to give a special thanks to my co-first author, David Brooks, on this work. This is a photo of our entire team based in Watertown, Massachusetts. We have a pretty large computational team and we do a bunch of really exciting machine learning. We have 4 people here at this conference, so please come find us if you're interested in learning more about our work. Oh, and I skipped this. We also have 2 other posters that we're presenting at MLCB, so please come check these out.

Thanks so much. Any questions?

Time, yeah. Thanks. Really, really interesting talk. When you're showing the uncertainty baseline that you're training an ensemble of models, are you still picking like when, and you're choosing which candidates to pick, Do you pick the 5, say the top N highest scoring by average, or you take into consideration the covariance of the scores? Because intuitively you want to take into consideration the covariance of the scores to hedge against any of the models being incorrect, right? You can't just do a greedy search and pick the top scoring according to the average, right?

[Speaker:JOSHUA_ELKINGTON] Yeah, it's interesting. Yeah, correlated residuals is like an interesting issue. Because of the bias in the particular selection procedure of this isn't a random sample of points. We thought greedy would be a fair baseline here because— so I'm taking the average of the ensemble as the score. I don't think that was your question, but I am taking the average of the ensemble for both methods. But no, we don't take into account the covariance. But I think that's an interesting question to think about.

I also had a question. So you mentioned that the— deviation, the ability to perform well in experiment is not correlated with edit distance. But you mentioned that maybe it's about the representation that's learned in the embedded space. So I'm wondering, how does— what if you use that as one of your baselines to figure out if the sparsity in the embedding space around a given design, is that informative of whether or not it will be successful?

Yeah, that's a really interesting question. The problem with a lot of embedding-based models is they assume some sort of smoothness over the sequence space. And what we find is that even a single, like, an edit distance, like 1 to 2 from a point that's viable, can be catastrophic. And those, a lot of those are not well predicted by embedding-based methods because they assume that sort of smoothness. Now, if you had a, some sort of model that was both a self-supervised model, but also was joined with the label or the measurement, that might be a really interesting approach to be able to identify, hey, for this particular area of training data, there's actually very little support around that data point. So maybe we need to have higher uncertainty for any designs around that point. But I guess the final point I just want to say is that the space is so high-dimensional that we've tried versions of that before. And even points that seem like there's a lot of data around it, I can basically guarantee that for any input within an edit distance 5, I can find these examples that will fool your model. And we found that the only way to really identify them is to use this design distribution, specifically using the unlabeled design data as information for your score function.

So one effectively tuning parameter you have is the number, is the size of your design dataset, right? Because you get to choose, you could always generate more of those, I guess, from just finding, find another algorithm that would generate some more. Yeah. So are you always better off having the design dataset being bigger, or does that actually hurt you at some point? Like you'd end up, I don't know whether it's overfitting exactly.

I'm not sure how to think about it.

You mean the candidate set? Yeah, the candidate set. Because when you're training your class, your binary classifier, if you have more data in that, candidate set, now you have more data to train that classifier on, which I guess is always good, but I was curious if that's how it actually works out.

Uh, my sense is, yeah, that's, that's sort of, that's essentially how it works out. It kind of ends up being quite interesting in practice when you implement it, because, um, if you ever have a situation where you want to evaluate a score for a point that is OOD with respect to the OOD classifier, So it's outside of the design distribution. What you should do is actually just retrain the classifier including that point and ends up being quite nice because everything just ends up being in distribution now with respect to this model in ways that for the deep ensemble, you're dealing with things that are out of distribution.

One more question here. If you assume, for example, that there is one really good oracle model that can extrapolate well beyond the training set and a really bad one, and the training set ranges from like -1 to 1, for example, but both of these models, good and bad one, make some designs that looks like a 2, for example, which is like beyond the training set. Wouldn't we expect both designs to be OOD from the training set and then flag them both as like kind of bad designs, although one is a really good extrapolation and one is really bad one. Or maybe in other words, is there any setting where we, um, we can extrapolate well beyond the training set while still staying within the training set?

Just to clarify. Yeah, that's interesting. I think the best you can hope for in design applications is some sort of compositionality where you can like combine particular elements of what you're seeing in the input to combine for design. And for those points, I wouldn't expect those points to have a high OOD score because they are well represented by subsets or subsequences that are in the input. Now regarding your point about like, well, the score function doesn't take into account the particular model, right? So you're right. If you build a better model, how, like, what is the interplay? I guess our view on that is that we've tried all the best models that we think exists for this type of problem. And for the particular types of issues we're focusing on, these so-called adversarial examples, we feel like there's basically like a limit to like how good a model can be given some finite number of XY pairs. And we don't, we don't really see that, that particular model improvements really make a difference.

Thank you so much.