Scientist Stories: Samuel Stanton, Geometric Deep Learning for Protein Design

This will be a little bit of a, maybe a bit of a meta-analysis of what happens when you try to apply geometric deep learning to— a lot of my experience is informed by antibody design in particular. And yeah, maybe some unexpected obstacles that come, that come up, but also, you know, hopefully some interesting problems to think about to overcome them. I don't think I need to spend too much time on this slide for this audience, but broadly speaking, we're talking about exploiting the structure in data to achieve better generalization and understand why our models do what they do. As you would have noticed from the last talk, which was excellent, a lot of emphasis on invariance and equivariance. How much do you need it? When do you need it? These are all sorts of questions that we love to think about. Um, I think it's worth asking for a second about what are the overarching research aims of the geometric deep learning program. Oh, sorry about that. Yes, there we go. Yeah, thank you. Yeah, so I thought it was appropriate to take a step back and just ask, what are the research aims of the geometric deep learning program in general. And in, you know, kind of one of the seminal texts, which is this impressive monograph by Bronstein et al., Geometric Deep Learning, they lay it out right in the abstract. They say we want a common mathematical framework to understand the most successful neural network architectures and a constructive procedure to incorporate prior physical knowledge. And for the purposes of this talk, I'm going to be focusing on the second one. This is not going to be a talk about, uh, I guess, theoretical implications, but more to focus on, uh, how far has geometric, uh, deep learning actually gotten us, uh, so far when it comes to things like antibody design. Um, I do think that to really have a deep appreciation for where the barriers are in this problem, you really do have to know at least a little bit of biology. And so I'm going to go through kind of the introduction to the problem in a very slow pace. So the first question that you might want to ask if you're going to try to do protein design is, what is the protein? And at a high level, a protein is just a big molecule. Sometimes in freshmen, we like to say it's just all atoms in the end. However, it is really helpful to understand how proteins are made. And currently, our leading theory for how this happens is known as the central dogma of biology. You have genes in your DNA that code for proteins. If the gene is on, that gene will get transcribed into messenger RNA, which will get carried to the ribosome, and then the ribosome will translate the RNA into a literal sequence of amino acids, which is shown at the bottom, which is sometimes called the polypeptide chain. And it's called that because the bonds in between the individual amino acids are peptide bonds. So when you ask yourself what a protein is, this is one answer, that this is where they come from. However, when they operate in the world as molecules, they do not exist as these long serial chains of amino acids. After the ribosome translates the RNA sequence, it folds, right? And this is the protein folding problem. However, an important thing to note is that it doesn't fold into a particular static structure. It folds into a structure and then it moves., right? And this is what's known as the conformational ensemble. Again, this came up in the previous talk. And oh yeah, maybe one last thing that I'll mention is that again, because this is essentially the way that proteins exist in their environment, essentially the way that they function must be a consequence of the actual arrangement of the amino acids in 3D space. And so this is the reason for the general assertion that kind of order of information flow in understanding how proteins work is sequence determines the structural conformation ensemble, which in turn determines function. Now when you get to, so that's a protein. For protein design, what are we trying to do? We're trying to engineer molecules that accomplish specific objectives. There's a lot of different components to this. Actually, a talk this morning did a great job of laying this out. You need to have a theory of how, if you're trying to do drug discovery, how a disease works and how you're going to interact with it. Um, once you have a hypothesis for how you're going to treat a disease, now you need a library of candidate molecules that you hypothesize are going to interact with this specific biological mechanism. And once you have a library of candidates, invariably there will be more than you have time or resources to test. And so you have to do some kind of prioritization to decide what to try first. As a super concrete example that I spend a lot of my time thinking about, we're going to talk— spend a lot of the motivation for the examples that will come up come from antibody engineering. Why is antibody engineering such an important problem? Essentially antibodies are, they accomplish a surprising range of functions for how similar they are in terms of structure. So generally speaking, you know, during the panel I mentioned these exotic formats. What I showed here is kind of like the canonical format, which is known as its IgG form. You have these two chains, the heavy chain and the light chain. And what's interesting about it is that the constant region, which is shown in blue, actually basically doesn't change hardly ever. And different antibodies that bind to different antigens, or, you know, maybe serve to regulate things like allergic reactions, all this functionality is basically just incorporated in changes to the variable region. And in fact, even if you zoom in on the variable region, most of that also doesn't change that much. And in fact, most of the variation is observed in these hypervariable regions, which we call the CDRs. And these CDRs are typically what constitute the binding interface between an antibody and an antigen. And It's mislabeled here, but the— so on the antibody side, this should be the paratope, and then over here is the epitope. Yeah, that's correctly indicated there. And turns out that if you can engineer antibodies to, for example, bind to an arbitrary target, you can do all kinds of interesting things with them. You can bind to a critical protein in a cancerous cell and essentially sabotage the cancerous cell such that it ceases to reproduce. You can deliver a drug payload. There was actually an FDA approval this week for a drug called, well, the drug was previously approved, this was just an approval for a different use, but it was a drug called Zolair and it basically reduces severity of food allergies. They were— the big clinical trial success that they were celebrating was getting a patient from being able to tolerate less than a third of a peanut to 2.5 peanuts. So, but if, you know, this is a— if this is a life-threatening allergy, this can literally be a matter of life and death. So We're talking about geometric deep learning. We're talking about incorporating prior knowledge into our models. So a reasonable question to ask is, what prior knowledge do we have at our disposal? And a lot of it essentially comes down to your level of resolution. So we kind of move down the level of resolution here where at a higher level, biochemistry, you're thinking about Essentially cells is like systems of molecules and chemical reactions. You're not really looking things at like an atom-level resolution, but this is where a lot of the understanding that comes into play to do things like target discovery. Structural biology, now you're trying to reason about why specific molecules behave specific ways as a function of usually static 3D structure, and so there's a lot of discussion of things like loop conformations, for example. And then molecular dynamics, to me, this, you know, almost seems like a weirdly arbitrary distinction between that and structural biology because they're both looking at the atom-level resolution, but there's a difference in emphasis where there's much more focus on the dynamics of the protein over time. So how do people typically tackle protein design right now? Well, historically, basically the field was dominated by physics-based approaches, Rosetta being kind of the dominant model for a substantial period of time. And broadly speaking, the way that this workflow works is you, Specify your requirements in a structural fashion. So you say, perhaps I want a binding interface that's complementary to this particular epitope. Or you say, I want a structure that has this particular notion of symmetry or stability. The upshot of it is you've specified your objective in structural terms. Then you design a structure that meets these specifications. And then you map back to sequence space via inverse folding. This is starting to get a bit blurred with the advent of sequence structure co-generation. However, even in that case, if you want to do things like condition on a target epitope, a lot of the basic subroutines are shared. And a lot of the benefit of sequence structure cogeneration is really focused at trying to like eliminate this inverse folding problem and potentially eliminate some modeling error. A more recent approach of which there are several examples, the figure here is just a a particular example that I know about because I worked on it. This is what I'll call the black box sequence design approach. And here you say, okay, whatever it is I'm trying to engineer this protein to do, I'm going to assume that I can measure it somehow. In the case of binding, we have an assay called surface plasmon resonance, or SPR. And you can get a dataset of amino acid sequences and binding affinity labels. You can train a sequence model to do generation and discrimination, and it'll spit out some sequences that it thinks optimizes this function. And the reason that we call this black box is because even though we have this hypothesis that under the hood, inside the machine, there's something structural going on. We leave it inside the machine. We just focus on inputs and outputs. So you might be wondering, well, when— which tool should I be using? And what's really helpful here to understand is that even though certainly to me when I first started looking at this problem, I thought pretty much all protein structures looked like a bunch of ribbons, there's actually some very important differences between different types of proteins. And this emerges in essentially which types of proteins different protein design methods do best on. So structural methods, and I would say this goes for both traditional physics-based approaches and ML augmented approaches currently work best designing what I would call fairly generic proteins that have very rigid symmetrical structures. So what I mean by that is like, if your goal is to show that you can create a protein that expresses in the lab, you can do that very well with structural approaches. Papers like RFDiffusion and Chroma did exactly this. Not only did they create proteins that express, but they created very pretty proteins because Humans love symmetry, right? And this is exactly the sort of thing that structural methods are very good at. They don't tend to require as much data because they have very strong inductive biases. But as I mentioned in the panel, this sort of pulls you away from scaling. They tend to be extremely problem-specific. It typically takes quite a bit of time to kind of retool it for similar, but it's a substantially different problem. Black box sequence design currently has been working pretty well, at least for us, to design things like antibodies. And the reason why is because, as I mentioned before, the function of antibodies, at least when it comes to binding, is really driven by these CDRs. And these CDRs are like the worst thing for structural models. They are not rigid, they're pretty flexible, and they have, you know, very funky shapes. And that is exactly what makes them good at what they do. Um, and so even though you can achieve success with these black box models, it comes at the cost of data collection. And so you're always sort of presented with this question of like, could you have done more with less, basically? However, it does put you in a good scaling regime where it's not that hard to condition in a very general way in sequence space. So again, we kind of come back to notions of sequence structure cogeneration and essentially what if I supply both types of information to the model. The question is, is it, you know, what evidence is there that this could give us a boost in performance? And the answer is, right now, the evidence is pretty inconclusive. So when you look at papers that include structural information in the embeddings, the tasks that they evaluate on are often biased towards structure. What I mean by that is that they will say, well, we can predict how these different proteins fold better. And it's like, well, yes, that's a structural question. Or we can predict remote homologies better. Again, structural question. So this paper is super interesting. I want to kind of highlight this section here. Okay, I missed the C. What's interesting about this paper, it came out last year, is that they do not just look at kind of structural discriminative tasks, but they look at more general functional assays. And if you look specifically at the ones that are like not super well correlated with things like expression, which would be these up here, it's a mixed story. Sometimes it helps, sometimes it doesn't. Definitely not the type of impressive results that we love to show in journal papers. But this is exactly the sort of evaluation that needs to be done because, as kind of an aside, the discriminative generalization part, although it doesn't get talked about as much right now because of the popularity of generative models, really one of the big bottlenecks when it comes to like actually getting good results, that there's a lot of different methods you have to choose from that give you pretty reasonable libraries. The problem is picking out the stuff that you actually want to test. So you might ask, well, part of the reason that, part of the reason that so little of this evaluation has been done is because people will say, well, physics is a strong prior. I'm giving my model more information. Surely that must be better. It's not even a question. And I think the truth is a little bit more nuanced. And to illustrate this point, we're going to walk through— we're going to go all the way back to our Stats 101 class, and we're going to talk about estimation of the mean parameter of a Gaussian with known variance. So let's suppose I put, you know, a fairly generic prior, which is shown on the left. I draw one sample, I do my typical Bayesian update. This is all closed form, and I show my posterior belief on the right. Okay, so far so good. We've got some data, we've updated our model. Funky dory. Now you might ask, what if I had chosen a different prior? What if I had chosen this much more peaked prior, much closer to the true mean? Unsurprisingly, even though we see the exact same data, boom, we almost— we get almost the exact true mean with just one observation. And when you— this is often how I hear people talk about structural information, is that it's this unicorn. That gets us instant better generalization. And kind of as I just showed you, when you actually look at discriminative functional prediction, the story is much more blurry. And part of the reason for this is that you have to ask yourself, well, okay, yes, this is a strong prior, but was it the right prior? Because if I had, for example, gotten the sign wrong, I still see the same data, but my, my model is still badly wrong. And in fact, the strength of my prior has caused me to look less at the data when that's exactly what I should be doing. And so the takeaway here is that strong priors are extremely unforgiving. And so the strength of your prior ought to be proportional to your confidence. In the variety of your assumptions and the quality of your implementation. Something that I think actually, again, was alluded to in your talk, part of the reason why you were so happy to ditch equivariance in some portions of the model is because it's super hard to implement right. And another classic example of this is, You know, the classic story about equivariance of CNN is actually broken if you don't handle aliasing correctly. And how many people think to do this? I would say very few. And so this is the risk that you take when you bring stronger assumptions into the picture, is it's a high-risk, high-reward gamble that you're right and you did it right. And so you may say, well, I'm too grown up for Gaussians. I don't care about anything you just said, which, that's okay. However, this is actually a much more foundational understanding of the way that generalization works that comes from essentially a probabilistic information theoretic view that was really articulated well by, uh, Mackay, and, uh, that my previous advisor loved to talk about. We talked about almost every lab meeting. Um, and the, the key idea is that if the prior is too vague, you will not be sample efficient. And this is the criticism of black box methods, particularly purely sequence-based ones. But if it's too specific, yes, sure, you'll converge quickly, but you'll converge to the wrong model. And that's basically what's shown in this middle little blob here, is that because of, say, either an error in your assumptions or error in the implementation, you've actually excluded the true model from your hypothesis space, and no amount of data will ever get you to the truth. Which brings me to, I guess, an analogy. Which is to say that common wisdom in machine learning, one of the first things that I learned was that if I wanted to estimate a density ratio, I should estimate the density ratio and not model the densities individually and then take the ratio. Now, you may say, where did this come from? This— we were just talking about proteins. Why are we talking about Gaussians and density ratios? Point is that all of this is about error propagation. That's essentially the situation where you would want to model these densities individually is the uniform prior situation where you can just nail P and Q instantly. If that's not the situation, you're much better off modeling the density ratio directly, and there are ways to do this. And so this is, uh, This is essentially where people are coming from when they're talking about structure as a nuisance variable. I don't think that it is, but they are correct that the inputs to the system are not structured and the output to the system is not really structured. Inputs are sequence because that's what we order from TWiST, and the outputs are measurements because that's what we use to decide if a molecule goes forward. And that doesn't mean that structure is useless or doesn't have good inductive biases, but this is the picture that you should have in your mind of does it make sense to apply it to this particular problem. So I want to dive in a little bit deeper into what might be some of the error in assumptions or in the data that could be causing structure to not give us the benefit that we might expect. And I would say the first thing that, um, is good to know if you don't is that the protein structures in the PDB— and these are the protein structures that basically every folding model is trained on— are frozen crystallized structures. And the key thing to observe is that even if you believe that, you know, the primary determinant of function is structure, like let's say that essentially function is conditionally independent of sequence given structure, we are rarely looking at these proteins in frozen conditions, at least in terms of where we want to put them eventually. We want to put them in people, and people are not frozen. And so a reasonable question to ask is, does this distinction matter? And perhaps unsurprisingly, the answer is sometimes yes. There's a super interesting work on room temperature crystallography, which I didn't even know you could do, that essentially showed exactly some of the differences that you see between frozen structures and room temperature ones. And so the observation is essentially that when you use a— usually we don't have frozen crystal structure. Even if we did, we really care about room temperature. And that's not what you get from folding. You get the folding model is giving you its best guess at the frozen picture. And these are not necessarily the same thing. And so this is a classic example of domain shift, which every ML person spends a lot of time worrying about. But it's not just, um, it's not just that. Another aspect essentially boils back down to, let's say, the choice of metric space or your choice of loss, uh, which is that— and this is, you know, gets back to why structural methods struggle with antibodies so much. Is because even though the RMSD on the antibody structure really is not that bad on the whole, if you zoom in on the loops, it makes, you know, some pretty big mistakes. And these are mistakes that can completely change your prediction about whether an antibody binds or not. And I would say to me it's not 100%, but like the first observation here is that Here we're comparing to the crystal frozen crystal structure, so I don't think this is a domain shift thing. I think this is a metric thing that this is something extremely difficult for the, the model to predict. And so, and it's not particularly encouraged to predict it well as opposed to other things, and so it's just not. But this is essentially a no-go. If you are trying to make extremely fine-grained predictions based on the confirmations of these loops. The other issue, and I think this is like a super interesting question for geometric deep learning, is that essentially the way that we pose structural design problems is really not well posed. And this is again, like was alluded to in the talk earlier, that proteins are not static structures. I think people— this point is starting to be more broadly appreciated, but then if you actually look at how we go about conditioning, we're still conditioned in a very static way. And so the question is like, how do you condition on something that maybe looks more like almost like a video of here's like, you know, the starting unbound states, here's where I want to end up, and there has to be some sort of like reasonably low entropy path to get there. And this is like kind of the the criterion that I need to condition on. And I would say, as far as I know, like, this is at the moment essentially a wide-open problem. We haven't really quite figured out exactly how to do this kind of conditioning. The reason that I mentioned MD earlier is because I think MD comes up a lot. When you're starting to think about how to tackle some of these challenges. In particular, if you think about the domain shift problem, you might suppose that if you simulated, if you took the frozen structure and then simulated how it would change if you increase the temperature, maybe this gives you some useful data augmentations that helps eliminate some crystallographic effects. Not a guarantee, but possible. Similarly, if you want to condition on like a trajectory of conformational ensembles that are interacting, well, you at least need the conformational ensembles. And traditional MD is, uh, well, you'd wait a lifetime before you get a useful amount of training data. And so I think— I thought the AlphaFlow paper, which just came out a couple weeks ago, super interesting example of people that are starting to think more carefully about this problem, say, well, okay, how I at least get something that looks like MD but without actually running MD. Um, as I'm kind of drawing towards the end, I actually have not been paying attention at all the time, so hopefully it's, it's not too bad. Um, but, uh, I, I do want to say that, uh, you know, there was a, there was a question of what are the biggest challenges for geometric deep learning. And certainly there are technical challenges, and those are the ones that I just alluded to. But another, I would say, pretty big challenge is just one of perspective, which I would say tends to be extremely, let's say, a bit rationalist and a bit model-centric. And the truth is that Domain knowledge manifests in many different ways, sometimes in model architecture, but actually it's in, it's everywhere. It's in how you frame the problem. It's in what data you consider and what you don't. And in this sort of case study that was like extremely impactful for me, if you spent any time in like the, let's say, non-autoregressive sequence generation community, there was a huge focus on length change. Everyone was obsessed with length change. And so naturally, I, you know, had a solution for length change that was turned on. And I talked to an antibody engineer, and he's like, well, you know, like, once the germline's fixed, B cells really don't change the length of an antibody when they are doing somatic hypermutation to change to mature affinity. And this was a shock to me because I had assumed that length change was just always something that you wanted to do when you were doing protein design. So I decided to just turn it off, which was like a 1 or 2 line change in my sampling code, and my expression rate doubled. And the point is that the architecture didn't change. I didn't even, you know, like I didn't even retrain the model. But my results changed substantially because I talked to an antibody engineer. And so it's just, you know, it sounds cliché and it is cliché, but it's like you can't spend all of your time thinking about your model. You're missing out on a lot of the picture if that's all you do. The other thing I would say is Know your competition. And what I mean by that is in, let's say, protein design in particular, in the ML community, there's a lot of focus on what I would call hit identification, which is essentially identifying a binder for a new target. And I think that there's a perception that like if you find a computational method to do this, that this will somehow be a sea change for drug discovery. And you can imagine my chagrin to find out that this was essentially a solved problem in pharma companies, but their solutions were not computational. Their solutions were inject an animal with the antigen and then see what its immune system does. That's not to say that you can't work on this problem. Like, the fact that a problem has a solution now does not mean there aren't better ones in the future. But it is good to know what those solutions are because that's essentially your bar for impact. So just to kind of throw in a couple takeaways, if you get nothing else out of this talk, it's that geometric deep learning for proteins is very, very far from a solved problem. They are much more than static point clouds. There's a lot of issues to be resolved when it comes to trying to eliminate these crystal artifacts, how to appropriately condition on conformational ensembles, other points that I touched on. The other, and this is maybe a bit more of a philosophical point, but it's like at the end of the day, physics is just a model. Treat it like one, know its assumptions, know its limitations. And finally, this is the last one, but If you're working on an algorithm whose reason for existence is it works well in data-scarce regimes, try to pick a problem that's really data-scarce and maybe not binding affinity. A great example of this is a property called clearance, which is essentially like how fast your body flushes the antibodies out of your system. If it doesn't hang around long enough, it can't do what it's intended to do. This is something that our only currently reliable way to measure it is to see what happens in monkeys. And so, you know, you get maybe a handful of measurements a year. That's data scarce. And could really use some help on that one. But it is much, in some ways, a much harder problem because it's an extremely complicated system that you're looking at. And yeah, that's my talk. Thank you all for listening. There is a protein we could have one too. Advanced AI teams. Yeah. Sorry? Why do you hate structure prediction? Don't hate it. Right tool for the right job. Like I said, if you want to design rigid symmetric proteins, go for it. Actually, like capsids are like, you know, potentially a good example of that. Thank you. So just a quick question. At some point, it would be nice to get examples, like, like some expectations of problems, specific problems people are going to have. But I think it's not— I mean, just the dynamics, kinematics, you guys try to— Sorry, I didn't quite catch the question. Yeah, well, I mean, to me, like, I could be wrong, but this just seems like an MCMC problem to me. Yeah, it is. Um, like, I don't Not being an MD person, I guess I can't say for sure why this seems to present such a conceptual challenge, and I suspect maybe part of it is that the MCMC community has not really been made aware that this is something that could be brought to bear. Maybe a decent example of something that you might used to try to essentially explore the conformational landscape more efficiently actually would be exactly this sort of smooth discrete sampling, uh, formalism, where essentially like you, uh, you hop around in a smooth space and then drop the temperature, right, and look at the more local geometry. I don't know if that answered your question. Yeah, thank you.