Scientist Stories: Matthew McPartlon, Protein Docking with Latent Diffusion

Yeah, okay. So this is going to be our last talk in the morning session. We are having Matthew McFarlane from VANT.AI. He gets his PhD from University of Chicago and is a senior research scientist at VANT.AI working on protein-protein and protein-small interactions. And he's going to talk about latent NOC. I think it's NIPS paper probably, or workshop. And yeah, let's move to the drug discovery part of the graph. Thanks. How about now? Okay, can hear you. Hear anyone? All right. Hi guys, I'm Matt. Um, I'm at Vance AI. I work Yeah, basically what Ali just said. I do protein-protein, protein-small molecule interaction prediction. This is a bit of a detour from the more standard graph neural network graph learning talks. So today we're going to talk about protein-protein docking and a method called latent dock. So this is an approach developed in collaboration with Celine Marquette, Tomas Gavner, Daniel Koda, Alex Zach Carpenter, Luke Benet, Michael Bronstein, and my PhD advisor, Timbo Zhu. So before we get into this, I'm going to first define the problem of protein-protein docking. So this protein-protein docking problem, the task is to predict the structure of a bound protein complex, so like the way that these two proteins might interact in your body or in a cell. Given the structures of the individual components. So you get to see each individual component protein, then you want to predict how they'll actually bind in 3D and interact. And kind of this problem, at least for now, is split between kind of two approaches. There's flexible and rigid docking. In flexible docking, you want to model conformational changes in these individual components. So you might see components, how they would look kind of individually apart in your cell, and then you want to model them in their bound state. So you want to model these changes that might occur only if one binds. So for example, like when an antibody binds to an antigen, like your body stimulates some kind of immune response, the actual loops on the antibody are going to change to like better match the 3D shape of the antigen when this binding occurs. So this is in contrast to rigid docking where you're really just predicting rotation and translation. Of each of the individual components. So you're just placing those as well as possible without changing the structure. Great. So why is this an important problem? So protein-protein interactions are kind of like— they essentially control all major biological processes. So they're at the core of like antibody-antigen interactions, foundational cellular processes like neuroprocesses, They also— this protein-protein interaction prediction forms the basis for most established like drug programs and therapeutics-based discovery programs. So you want to predict maybe ways of inhibiting or ways of inducing certain protein-protein interactions. Great, so now we'll talk about— yep. All right, cool. So now I'll talk about some motivation for this approach. So hopefully some of you guys, at least everyone in the bio world, is familiar with AlphaFold 2. This is like the DeepMind protein structure prediction program, which is set to solve protein folding. And AlphaFold 2 is this really, really interesting architecture. Like, you could probably get through a whole PhD just knowing this. It's really an incredible process. So it doesn't just work, it works really well. And there's some really interesting takeaways that you can get just from looking at this. So first, notice that Like, kind of the whole model is meant to predict a protein structure given some features. And what are those features? They're all derived from a protein sequence. A protein sequence is just some tensor of tokens. Each token has a value between 1 and 20. So that's a way to think about a protein sequence. And this protein sequence gets converted to a multiple sequence alignment, templates, all of these other things. But the important takeaway here is that all of these input features in AlphaFold 2 are invariant to rotations, translations. They don't say anything explicit about 3D geometry. So this gets passed into the encoder of AlphaFold 2. These are analogies that I'm making. They don't call this an encoder in the AlphaFold 2 paper, so just go along with this. These invariant features get passed to this invariant encoder. This invariant encoder is like this kind of large transformer network with some secondary updates for the pair features. I should say, like, the input is kind of your standard node features plus some edge features. These edge features are some kind of implicit representations of geometry. Great. So invariant features get passed to an invariant encoder. That encoder has some invariant latent representation of the nodes and edges, and that finally gets passed through this equivariant decoder, and that decoder produces 3D geometry. What's interesting is that along the way, like, so from the beginning to the end, there was no step where we had like these equivariant features. There's no like explicit notion of 3D coordinates. All of this is constructed by the network. So even though the input is fully invariant, we're able to predict this 3D geometry, which we'd usually think of as like only being possible with equivariance somewhere. And indeed, there is some equivariance in the decoder. However, the initialization of the decoder plus its features are fully invariant. So a question you might ask is like, why does this work? So shouldn't we need some notion of explicit notion of geometry in order to actually produce explicit geometry. And maybe we can think about— yeah, so sorry, this goes more with this slide. So we can think about like, why does this work? Why are we able to go from invariant representations of geometry to actual explicit 3D geometry? Great. And maybe the other thing is, if AlphaFold works with implicit representations of geometry, it should definitely still work with more explicit representations. So part of this motivation is to think about the multidimensional scaling problem. And in this problem, you want to go from some similarity matrix to a set of points which kind of recapitulate those similarities in the sense for, like, for— if the similarity matrices are the Euclidean distances between a set of points, then multi-dimensional scaling would be used to exactly recover a set of points that capture those distances. So there's a very simple formula to compute the distance matrix of a set of points X. So if you have a list of points in Rd, then you can write the Euclidean distance matrix of those points as a function of the Gram matrix plus some, you know, column vectors of weights. And then you can, after writing the Gram matrix— or sorry, the distance matrix is a function of the Gram matrix— you can just invert this and follow the classical multi-dimensional scaling algorithm, which just computes an eigendecomposition essentially of this Gram matrix. To recover the original X set of points. So that's one way to go from, you know, distance matrix to points. And along that way, we know kind of along the way that we haven't used any like explicit— or we've used explicit information of geometry, but we didn't have to worry about equivariance, rotations, translations, anything crazy. Kind of part of the reason for that is that this Euclidean distance matrix operator is invariant to rotations, reflections, and translations. So for any orthonormal matrix Q and like d by d matrix or translation B, the Euclidean distance matrix is invariant to Q and B. The other nice property is that this Euclidean distance matrix is low rank. So maybe you can see that in the previous slide here. So it's a function of the square root matrix. The square root matrix has rank d. Is the product of 2D-dimensional, uh, sorry, 2-rank D tensors, um, and this secondary term with these 1, 1 transpose, which adds rank at most 2. So it's nice that we have a very low-rank representation, um, which is also invariant to rotations and reflections, uh, and as a const— uh, so One consequence of this invariance to rotations and reflections is that we won't be able to recover the absolute point set. We lose this kind of rigid information, but we can recover a point set that is equivalent up to these rigid transformations. Great. So, kind of to summarize this, multidimensional scaling shows us the existence of an invariant mapping between Euclidean distance matrices and Euclidean coordinates., which realize those distances. And AlphaFold 2 has already shown us that implicit representations of geometry— so there exists a neural network architecture that can take implicit representations of geometry and use those to produce explicit representations of 3D coordinates. So this kind of motivates us to maybe learn in this implicit encoding space, so kind of the latent representation of some geometry encoder and use some decoder latent type decoding method to actually realize that geometry. So, right, so what does that actually give us? We're actually no longer constrained by these equivariance considerations that are typical in common kind of protein geometry-related learning tasks. We're able to leverage very deep standard transformer architectures with standard nonlinearities and so on. We don't have to care about, you know, whether or not any of our layers or any of our nonlinearities break equivariance. We're also able to leverage highly engineered, like, image diffusion frameworks, which have seen, like, a ton of engineering, and they're practically incredibly useful. So, yeah, a lot of progress has been made here. And rather than redefine our diffusion framework over some rotation translation space, we can just use the image diffusion framework directly because we're learning in an invariant representation. So I call this like geometric diffusion without the agonizing pain. So the other thing is we can actually naturally encode and decode other attributes apart from geometry in this invariant representation because there's nothing— there's really nothing stopping us from doing that. For proteins, we can encode the sequence, functional attributes, confidence estimates for every residue, so on. We're also able to handle much more natural representations of geometry. So it's very rare in the comp bio world that you have this exact explicit structure that you want to just like, you know, map around. Usually you just have some kind of idea of what the structure should look like. It might be noisy, it might be incomplete, you know, a certain section might be missing, but you still want to be able to model it. So in this type of approach, you can capture that in this distance matrix pairwise feature, you don't have to provide explicit distances. You can provide distances with some coarseness. You can also naturally mask out distances by just leaving that row, replacing this distance with a mask token or something like that. It's also very trivial to model conformational flexibility, which is something that most person-person docking methods can't handle. Especially in this— especially in a— sorry, yeah, especially like with equivariance concerns. So if you're trying to predict a single rigid rotation and translation, it's very hard to actually move the underlying like atoms in each chain. So one of the problems this addresses— so like why actually, why use diffusion here? Is it just because diffusion is popular and that's what gets into NeurIPS? Sometimes. But there's another issue. Yes, but no. So apart from this, like, so there is an issue with regression models for this problem. And this is kind of like, we saw this a lot like way back in the early protein folding days. What we try to do is like, predicts a distance matrix and then try to realize geometry with, like, some gradient descent algorithm over these predicted distance potentials. It's not what's important here, but kind of what you see from this type of analysis is that actually, like, regression models like AlphaFold 2, they tend to, like, yeah, the output is actually like the superposition of 2 or more predictions. So it's not often that AlphaFold 2 knows exactly where to place a protein. Typically, it's kind of like it's split between two sites, and depending on the confidence in each site, it'll kind of move the prediction towards one site or the other. So in the final output, you know, it has to make a single prediction, but if you look at like the intermediate state and what it thinks about like the protein structure via its like predicted histogram, predicted distance logits, it's usually kind of split between two or more sites. So you'll see like high likelihood of, for some pair of residues between like, say here, the green and pink chain, you'll see high likelihood that, you know, this upper site is the actual binding site, and you'll see a high likelihood that this lower site is the actual binding site. It's okay to actually like make both of these predictions when you're predicting just your distogram, but when you predict your final output, you have to decide on a site, and typically what happens is the protein gets placed between the two sites. So that's one thing, like, kind of one way in which regression models fail here. So diffusion models offer a nice alternative because you'll actually be able to sample from both sites, and then you can use some, like, secondary ranking algorithm or something like this to decide on which confirmation you want to select. Great. Yeah. So now let's talk about the method. So this method is LatentDoc, and it uses latent fusion to learn this protein-protein docking approach. So we kind of start— I'll get into the components here in a second. Great. So here our inputs to this algorithm are kind of invariant representations of 3D geometry along with other attributes to the protein. So we derive some explicit representations of geometry from the protein's 3D structure. So this is stuff like C-alpha distances. These are types of atoms at the protein residue level and some orientations. So I said before that multidimensional scaling is invariant for reflections. One way to break this is to provide dihedral angles. So if you know dihedral angles between points, kind of like torsion angles, then you can decide on the proper reflection. Apart from this, we have some standard sequence information which tells you kind of the relative positioning of the proteins. So in protein fold, from a sequence of amino acids, there is a start and end like natural language. So we have ways of kind of telling the network where that start and end is. And finally, we also provide the amino acid type. Great. So at a high level, our VAE architecture is just some encoder and decoder, and we train this before we actually train our diffusion model. Our diffusion model just recovers these latent representations, and those latent representations are converted to actual 3D geometry by the decoder. So our encoder consists of these pair-biased self-attention blocks. So it's almost like regular self-attention but with an additional term. We'll get into that in a second. The latent representation has like kind of the standard VAE KL divergence loss, and we also do some layer normalization there to kind of make it play better in standard diffusion frameworks. Finally, we keep the sequence dimension on our latent representation, which is not standard in VAE architectures. And technically, we don't have an autoencoder either because our input and output are different, but our output is just some realization of the input. So it's almost like a translation task. So yes, our encoder architecture, like I was saying, uses this pairwise self-attention. So if you look, the standard attention formula is given here. And all we're doing is just adding some additional term. That additional term is a linear projection of the corresponding edge feature. So when you compute the attention weight for feature i, node i and node j, you also have some edge associated with it. You use a linear projection of that edge feature to bias the attention weight. So essentially what this does is it gives us a way of transferring these pairwise geometric representations into the node representation. The decoder we actually borrow from AlphaFold. We use invariant point attention. So kind of, this is also quite similar to the pair-bias self-attention from the previous slide, but we have this additional term on the right here which captures some notion of geometry. So these TIs in the rightmost term are rigid frames. So they consist of some rotation and translation. And these get applied on kind of a per-node level. So let's see. All right, so start with our pairwise self-attention. And then we add on this other term. So apart from the standard dot product attention, we also have these point queries and point keys. And these are just linear projection of the hidden node feature reshaped to have a trailing dimension of 3. So you can think of these as kind of like local points, invariant local points that are associated with each node. And then we use this equivariant rigid translation to transform these points into the frame of every node.. And initially these frames are initialized by identity rotation and no translation. So kind of at the start of everything, or at the first layer, the points are all kind of lying at the origin, right? So you kind of look over here, so in the local frame you might have these points that are centered at the origin. Once you apply these rotations and translations you get the points of transform, and you have like the negative distance of this added to your attention weight. So the closer the points, the more they affect the attention weight. So the output of this encoder-decoder architecture is a full 3D protein structure. So we model all atom coordinates, And we train this with the BATE loss similar to AlphaFold 2. We also use a linear projection to recover the sequence, and we train that with cross-entropy loss. Along with that, we have some per-residue level confidence scores. This is like the LVDT score, which measures how well a residue's microenvironment kind of matches with the ground truth structure. Great. So now I'll talk about the diffusion model. First, does anyone have any questions so far? Now I'm going a little fast here. Yeah, I was just curious, did you try training this score network before convergence and it just like wasn't stable, or just a Yeah, so I just followed— like, the idea was to kind of follow the latent diffusion for images framework as closely as possible and kind of see what works there. Yeah, yeah, that was part of the reason. One of my favorite papers, and I like that idea a lot. And also just reading the literature for like equivariant diffusion, like a lot of these Diffusion, like the diffusion frameworks, are, they're very complex. They tend to be like unstable, very approximated. There's a lot of bugs. There's not a lot of open source implementations. They're just very under-engineered. It's mostly like grad students working as a one-off. And for like things like Stable Diffusion or like Image Diffusion, there's a million implementations. You can just grab one off the shelf and apply that to the problem. Kind of wanted to make this more of like an image diffusion-like problem rather than like this type of geometry kind of problem. Yeah, so part of the reason for applying this KL divergence loss and for layer normalizing this latent representation, so I talked about here, part of the reason we do this is to actually like make the representation look more image-like. So this is a bit of a hack. To make this work well with standard image diffusion frameworks. Okay, so yeah, now I'll talk about the diffusion model a bit. So we have a conditional diffusion model, and what we actually condition the diffusion model on— so our VAE learned to represent full protein structures, and they're given kind of full geometric information. Our diffusion model, on the other hand, only receives the structure of each, like, individual component So you have some protein complex, you break this into components, and you supply the conditional diffusion model with the structure of each component, but you don't tell it how the components should be arranged. So one way of doing that is kind of aligning all of the structural templates in a pairwise matrix and masking all the inter— like, off-diagonal features. Apart from this, we can also like easily incorporate things like residue-residue contacts or ideas of like where the chain should bind by saying like basically adding these one-hot encoding features to our pair representation or our node representation to say that, you know, the proteins, the individual components should be close in proximity between these two points or at this certain point on the surface. Right, and another way to think about this is like the score network is kind of like learning a type of matrix completion, and then the decoder is doing the multi-dimensional scaling. And then one other thing you can do for these template features, they don't have to be exact, they don't have to be true, you put in discrepancies. They can be some noisy version of a distance matrix. They can be coarse one-hot encodings of very wide intervals. And you can also optionally mask the distances to design parts of the sequence. Great. Great. So now we'll briefly talk about results. So when I talk about results here, there are different ways of evaluating protein complexes. So one is kind of all these basically have to do with this RMSD, root mean squared deviation. So you're computing some optimal alignment of 2 objects, the alignment that minimizes the L2 distance between the coordinates. One way to do this is over the entire complex for protein. Another way to do this is just by aligning one of the components and then kind of measuring how far off if you align on one component. So highlighted here in orange and gray, if you just align on this, how far off are the kind of red and blue components? And then finally, you can align just on the kind of surface where the two components interact, and that gives you another type of RMSD. So for all of these, you're just deciding what you want to find, either the entire structure, one of the chains, or just the portion where they're interacting. And then DocQ score, which is what I'll be reporting, is kind of a combination of those 3 RMSD types plus some term which says how well did you actually recover the biological interactions that occur between these 2 proteins. And typically it would say like a DocQ score of greater than 0.23 is acceptable. All you have to know is that you're mostly right if the score is above 0.23. Great. So our dataset for this task, it's currently an in-house dataset, but it's based on protein dimers split using a combination of what's called FOLD-Seq and MM-Seq. So what that is, is a way of saying that it gives us a way of saying that the proteins in our training set and test set kind of look their interactions look different, and the actual sequence tokens which make up these interactions are different as well. So it's a way of kind of reducing the amount of bias between your training test set by doing this kind of structured sequence-based split. On top of this, so we retrained 3 kind of state-of-the-art protein-protein docking methods on the same dataset. These are DiffDoc, Equidoc, and DockGPT. In order to have a fair comparison. Great. So here are results for 150 naturally occurring protein dimers. So you can kind of see our method does quite well in terms of recovering reasonable, so high DocQ protein complexes. So roughly 70% are predicted with acceptable or better scores. This is in contrast to other methods. So DiffDocPP is another diffusion-based protein docking method. It does fairly well, but we're able to outperform this using this latent diffusion approach. One thing to notice here is the percentage of very high-quality predictions, like very low RMSD, high-quality predictions, is not as good as DocGPT, another docking method, and that's kind of because our VAE, our encoder-decoder, had relatively high reconstruction error, so around 1.5 angstroms. Yes? So in the baseline, you didn't consider the equivalent which is using these graph geometric versions for the, let's say, these problems? Is it not useful in this setting, or there is another reason? So in general, let's ask this more general question. So you're using the autoencoder-decoder for learning the geometries, but why don't we use the geometry deep learning instead of the encoder-decoder? Yeah, I think like, so geometric deep learning can't readily be applied to this problem because there are— so it could be applied to a single structure, but as soon as you have two structures, well, how you decide to position those two structures relative to one another breaks equivariance. So if you want to just operate on a fixed static structure, that's fine. Geometric deep learning will let you do that. If you actually want to, like, fix a structure and now move some structure relative to it, that will break— like, a standard equivariant architecture can't really do that because you no longer have equivariance. Like, if you're randomly choosing, say, like, a starting pose, then you're just back to, like, standard augmented data augmentations. There are ways to get around this with geometric deep learning. So people have developed diffusion models and methods to overcome this. One is this DiffDoc method. So kind of the idea here is to predict this global rotation translation which maps these final ending positions. Yeah, but there are a lot of— there are reasons why this sometimes doesn't work well in practice. So there's like this huge truncated truncated, like, IGSO3 approximation in the actual diffusion formula, which is kind of unstable at very small time steps, very small noise scales. Equidock tries to do something similar, predicting some key points on the protein pockets. So you kind of think of this docking problem as predicting kind of points on each protein that will interact, and then you can think of docking as aligning these points. So that was the goal of Equidot, but what actually tends to happen in practice is it doesn't know where the points are. It's kind of a hard problem, and maybe it predicts— it doesn't have to predict like a fixed set of points. You can kind of predict points all around the protein with some level of confidence, and then when you actually do this alignment, the two proteins just kind of overlap one another. It's kind of like a low-variance way of making sure that the points are relatively well aligned. It's not actually like a biologically meaningful result, but it looks okay, say, in terms of like ligand RMSD or complex RMSD. Right. So here we just compare in terms of sampling efficiency. So I kind of like one way of thinking about this So for this DiffDocPP problem, you actually start like in their framing, you actually start with these two structures in a random rotation translation. And the goal is to kind of move one structure, move one of the structures over to its like full docked and ground truth position. And doing so, kind of the starting point, maybe like your initial starting point, your initial random rotation translation might be important if you're very far away from the site., it might not be able to get it there in time. So we actually show that in terms of sampling efficiency, our method tends to perform much better. So given the same number of samples, so if we only look at, you know, 5 generations for our method versus DiffDoc, we're able to get like significantly higher success rates. So DiffDoc, in practice, you actually have to sample from this a lot until you see a correct prediction. I'll quickly talk about some extensions. So we also extended our approach to other conditional protein generation tasks. So there's protein folding and protein inverse folding. So we use the same VAE. 'cause we can still recover the sequence and structure, but we kind of address each problem by doing some feature dropout in the denoising model. So for inverse folding, the task here is given a protein structure, predicted sequence. For protein folding, it's given a protein sequence, predicted structure. So I'll go quickly here. We drop out the sequence information and we use the same exact model. We just condition on these structural 3D templates, and we're able to get competitive performance with state-of-the-art inverse models. And finally, we apply this to protein folding, where we drop out the templates and leave this sequence information. We're able to significantly outperform existing kind of state-of-the-art folding-based diffusion models, and we're competitive with state-of-the-art monomer folding models. ESM, ESM4 developed by Meta. Yeah, thank you everyone for listening. Happy to take any questions.

Hi, I have tried to talk One problem that I have frequently encountered is that although I have to get all this code execution and everything in this final job on the mining side, it ends up producing a lot of clashes in the final structure even when they are wrapped in there. What if I was for the latest job that we are using for interactions, do you suffer the same problem?

Yeah, so, so we, we compare the number of clashes with DiffDoc, Equidoc, and our method in the paper. On average, we have like 0.1 clashes per structure. Equidoc has, I think, like several thousand. And then DiffDoc, yeah, I forget the actual number, but it's like 10 or so. So the nice thing here is like, so with this approach, I can find it at some point. Yeah, but so when you learn this encoder-decoder architecture, you can just apply a very high clash loss on the decoder. So you can train the decoder to not produce any clashes. You can make that heavily penalize this. So anytime it sees an encoding, it knows, via adding this noise, whatever, it knows that protein structures do not have clashes. So you're enforcing this in a model that's really easy to train. With something like some method like TITTOCK. So you're starting with some random rotation translation of one of the protein chains. In order to get back to the correct transformation, you might have to pass through the structure. So you can't apply clash loss in this model because actually at intermediate steps you need clashes in order to fulfill the diffusion objective, like finding this rotation translation which moves you from point A to point B. On a direct path, that direct path might contain a lot of clashes along the way. So you can't actually enforce that loss in that.

So could that end up introducing bias in your models if you end up putting the hard class code?

If there are a lot of clashes in the training set? Yeah. And you're avoiding them. So yeah, if we have an example— so we use a pretty relaxed definition of clash. So we say 2 C-alpha atoms that are within distance of like 2.5 angstroms, that's a clash. We don't actually apply this over side chains. You could, but just the C-alpha atom clash tends to do pretty well. So typically you don't see— almost, it's almost impossible for 2 C-alpha atoms to be that close. Thank you. Yep, you're welcome. Yeah, I, I checked the workshop paper. It's still in code. You guys have time on it? Hopefully within the next couple weeks. We'd like to release it at the same time as this dataset that we use. Any other questions? Might be some in the chat.