WEBVTT

00:00:00.000 --> 00:00:03.000
And so I had many collaborators with many… from many different schools, and we were thankful to have a biologist, Stefan, from Harvard Medical School, who kept all of our experiments grounded in real biology. I think before I start, I guess I'm curious what everyone's backgrounds are.

00:00:03.000 --> 00:00:06.000
If you're, like, an interperson, could you raise your hand?

00:00:06.000 --> 00:00:14.000
Okay, and then biology? Okay, so I guess everyone else is more general like theory and Ml. Or.

00:00:14.000 --> 00:00:19.000
mathematics. Oh, they're all on Zoom? Okay, that sounds good. Okay.

00:00:19.000 --> 00:00:28.000
Um, so yeah, I guess I'll just start from the problem statement, which is really, like, why are we using AI protein folding? So I'll —.

00:00:28.000 --> 00:00:44.000
Okay, so I'm just gonna… so this is, like, the drug discovery problem, where we're really interested in finding early-stage binders. So let's say I have our target here, it's HER2, which is a protein receptor that drives tumor growth for breast cancer. So our goal is to design some binder.

00:00:44.000 --> 00:01:02.000
that will block the receptor signaling. So we'll… the two proteins will bind, and um… We'll get a… we'll get a successful drug. So how do we do this? Well, something that we can do is just generate a bunch of sequences, and then we'll have to go to the wet lab and express the proteins.

00:01:02.000 --> 00:01:23.000
And then we'll have to do some binding assays. And so this is a lot of physical work in order to check whether two proteins are actually going to bind. And in terms of cost, this is… Orders of magnitude more expensive than you'd normally think. So 1,000 proteins could be $50,000 to $100,000 in weeks of time. And this is only 1,000 proteins. The search space for sequences is huge.

00:01:23.000 --> 00:01:39.000
So here's a problem, and this is where AI protein folding comes into play. So the way I see it is it's reducing the search space. We want a way to have cheap and fast virtual screening of our proteins before we actually go through the whole process of putting them in the wet lab.

00:01:39.000 --> 00:01:54.000
So here we have the two amino acid sequences of our target and our potential protein binder. So just as a side note for non-biologists, so proteins are chains of molecules called amino acids. And so we can usually represent them with letters.

00:01:54.000 --> 00:02:09.000
Uh, here there are 20 common amino acid types, and when amino acids actually come together to form a protein, they lose atoms and are then called residues. So I'll interchange the two terms, but just for your information. So here, instead of needing to.

00:02:09.000 --> 00:02:24.000
Fully express them, and go through the whole wet lab process, and then what we can do is put the two sequences into AlphaFold, and it'll generate what it thinks is the protein complex. Here we have a huge speedup now, where we can virtually screen our proteins in order to.

00:02:24.000 --> 00:02:29.000
reduce the amount of proteins we actually need to put into the wet lab.

00:02:29.000 --> 00:02:45.000
But there's this question of, like, why do I want to interpret this protein folding model at all? Or why do this for scientific models? So on the way here, I was reading some papers on flat iron, and I saw this physics steering model that Alberto worked on, where they really were curious about, yeah, how does this?

00:02:45.000 --> 00:03:03.000
Other type of physical simulation model work. I think there are really 2 lessons that we can try and draw from interpretability. So the first question that we're trying to understand is, well, how does the model do these internal computations? And I think understanding this is helpful for building trust in the model, also for model debugging. And there's a lot of scientific curiosity too.

00:03:03.000 --> 00:03:09.000
Um, and I've seen some people do this type of work where they interpret the model, and then this helps inform maybe the inference setup.

00:03:09.000 --> 00:03:22.000
A second question is, can we, after we understand the model internals, begin to control its outputs and even extend capabilities? So these are the two questions that I'll be asking for protein models as a whole.

00:03:22.000 --> 00:03:32.000
So, when it comes to folding an actual protein, I'm just going to walk through the setup. And we're going to start with a protein sequence. So here we're just going to fold a single structure.

00:03:32.000 --> 00:03:44.000
So we're going to pass it into this protein language model, which produces a sequence representation where each token in our sequence receives some embedding.

00:03:44.000 --> 00:03:49.000
What happens next is we take the sequence representation and pass it into something called the folding trunk.

00:03:49.000 --> 00:04:06.000
Now, at this stage, we also initialize a pairwise representation. And so this is pretty unique to protein folding. It's not a normal representation you find in transformers, but all it is is it's meant to represent relationships between residues. And so it's just a 2D matrix. And here we just initialize it with positional embeddings.

00:04:06.000 --> 00:04:11.000
Um, all the folding trunk does is it just updates these representations.

00:04:11.000 --> 00:04:22.000
And then we pass it into a third module called the structure module. And it just renders our 3D protein structure. So it just outputs 3D coordinates of different atoms.

00:04:22.000 --> 00:04:31.000
So, from a whole, we're looking at ESM fold. It has 3 modules, and it maintains these 2 representations, the pairwise and the sequence.

00:04:31.000 --> 00:04:42.000
Now, the question that we want to ask is, how does the model actually fold proteins? And so we choose to look at a specific case where we're looking at secondary — oh, yeah.

00:04:42.000 --> 00:04:57.000
Alright, so, uh… Basically, you keep every intermediate step keeps having some sequence. Start with sequence, then you have a sequence again, and then you do something sequence again. Yeah? Yep.

00:04:57.000 --> 00:05:19.000
So… is just the intuition, because, like, at the end of the day, you just want to use LLMs. Like, why would you compress everything into segments? Why don't you keep, like… some embeddings, I don't know if that makes sense. Um, is the question about the…

00:05:19.000 --> 00:05:39.000
Yeah, I guess… Every step spits out a sequence. Oh, I see. So you keep compressing, you do something complicated, spit out a sequence, do something like… Oh, I see. So, instead, you could have just had some, uh, like, sequences compression, right? So you could, like, live in some embedding space, I don't know.

00:05:39.000 --> 00:05:46.000
Yeah. I guess in this case, so, um… The sequence representation here is always the embeddings, right?

00:05:46.000 --> 00:06:01.000
Um, so we actually… yeah. Yeah, yeah, sorry, sorry. So that's… that's what the, uh, yeah, that makes sense. That makes sense. So, yeah, here, when I say sequence representation, it is going to be an embedding per, uh, token that was originally in the amino acid sequence.

00:06:01.000 --> 00:06:07.000
Good. Yeah, that's a good question. Are there any other questions about this folding process?

00:06:07.000 --> 00:06:26.000
Okay, so yeah, we're going to ask now how the model actually folds a specific type of protein. Here it's called the beta hairpin. So I highlighted it in purple, but basically it's just one strand connected by a loop and then another strand. And so it's a very simple but very, very common structure in many, many proteins.

00:06:26.000 --> 00:06:38.000
And what we're interested in understanding is, well. How does the model use the sequence representation and the pairwise representation to come up with the 3D coordinates for that hairpin?

00:06:38.000 --> 00:06:58.000
So, something that we can do as a first intuition is… See, oh, are the hairprints actually present in the folding trunk representations? And this is actually… Okay, so what we can do is we take our beta hairpin, and we can throw it through the model. And we're mainly interested in the folding trunk because that is actually the component that.

00:06:58.000 --> 00:07:13.000
the folding really happens. Um, it's actually… it shows up in many different modules, but basically the folding chunk is where the pairwise representation is built up, and that is actually what's used for folding later on. So we're going to be looking at the representations there.

00:07:13.000 --> 00:07:19.000
Yeah, so the beta, so beta hairpin is this purple thing. Yep. I think I missed.

00:07:19.000 --> 00:07:34.000
So you're trying to see what causes these things to like hold in this particular shape. Yeah, I think right now we're just looking for, okay, are there linear representations of the hairpin that we can identify within the sequence or the pairwise representations?

00:07:34.000 --> 00:07:52.000
When you write these arrows and stuff, this is, like, you have a 3D geometry, and then some software figures out that this is a hairpin, or is it like there's a tag in the model output that tells you? So in this case, yeah, there's some software that just identifies patterns. Because these hairpins will form.

00:07:52.000 --> 00:08:10.000
structures that, like, are basically… have defined distances, so you can just look at that.

00:08:10.000 --> 00:08:25.000
The way that the model actually folds proteins. So, for instance, we can take a different protein that has no hairprints at all. So here it's just made up of alpha helices, and we can once again look at its sequence and the pairwise representations in the folding trunk. Now, there's this interesting question now of.

00:08:25.000 --> 00:08:30.000
Well, what happens if I just take the sequence representation of my hairpin.

00:08:30.000 --> 00:08:37.000
And I overwrite the sequence representation of this helix-loop-helix region.

00:08:37.000 --> 00:08:40.000
What's going to happen? I only need to do this at a single block.

00:08:40.000 --> 00:08:46.000
Similarly, what happens if I just replace the pairwise representation at a single block?

00:08:46.000 --> 00:09:01.000
Well, what we observe is actually that. The model produces a protein where instead of an alpha helix, it now produces hairpins. And so this is a really interesting observation, because we only need to patch one type of representation at a single block.

00:09:01.000 --> 00:09:14.000
And this… and so… What we can do now is try and understand, well, when does the patching work? And how — I guess reveals information about how the model is actually performing the folding computation.

00:09:14.000 --> 00:09:39.000
Okay, so can you explain again what is this pairwise patching? So it's just, uh… It's just a 3D matrix that relates items in the sequence. Yep. That's it? And then what I also observed there is some other dimension. What is this? Oh, if per sequence… Oh, yeah, here, um, the multiple blocks because.

00:09:39.000 --> 00:09:48.000
So the following trunk, it basically has 48 layers, and they call it 48 blocks. So that's what the stacking is meant to represent.

00:09:48.000 --> 00:10:00.000
Uh, no, it's not general dimension. Oh, it's… that's just meant to represent how many blocks there are, or how many layers in the folding trunk there there are.

00:10:00.000 --> 00:10:13.000
Yeah, yeah, it's all sequential. I understand. Is it like a channel, like internal, kind of like multiple filters, and yeah, for each filter? Uh, it just runs in sequence instead of having multiple channels.

00:10:13.000 --> 00:10:26.000
I'm talking about the one in depth. Yeah, it's like multiple layers, like, imagine… Yeah, yeah, it's like multiple layers. Your network has 10 layers, then each layer, he's extracting the pairwise.

00:10:26.000 --> 00:10:41.000
represent the pattern representation. And then the structure model looks depends on the final one. Yeah, sorry. Oh, on the final layer. Yeah, yeah, sorry. So after it'll it'll do all… it'll go through all the layers and then just produce a final output.

00:10:41.000 --> 00:10:57.000
And then the structure module. You only, in fact, need just the final layer. Yeah, final sequence and final… In terms of… in terms of the folding process, it only looks at the final one. But we're interested in understanding, well, how does the intermediate computation happen?

00:10:57.000 --> 00:10:59.000
So, can I?

00:10:59.000 --> 00:11:00.000
I see. Yeah.

00:11:00.000 --> 00:11:11.000
Can I ask a quick question? I'm wondering if this task is too simple, because there are models that can figure out, like, the secondary structure, that, like, you can just spit out the DSSP, and then you can just

00:11:11.000 --> 00:11:14.000
look at, you know, try to parse…

00:11:14.000 --> 00:11:16.000
I guess, beta strand…

00:11:16.000 --> 00:11:24.000
loop, beta strand, and just try to figure it out. You will miss a lot of things, but it's fairly good as a first approximation, so I wonder…

00:11:24.000 --> 00:11:28.000
Um, what's your thoughts on that?

00:11:28.000 --> 00:11:33.000
Uh, could you, I guess… I think I understand where you're going at. So I guess, what is the problem that?

00:11:33.000 --> 00:11:35.000
What I'm saying is that this is a fairly simple problem in terms of, like, figuring out, like,

00:11:35.000 --> 00:11:39.000
that that yeah.

00:11:39.000 --> 00:11:43.000
where beta herpians are for a majority of them.

00:11:43.000 --> 00:11:55.000
Yeah, I see what you're saying. So, and I guess, like, our… the problem is less so, like, oh, I want to identify where beta hairpins are in the protein, and more so about… Well, when I patch.

00:11:55.000 --> 00:12:06.000
When I patch a single representation. How do I… what does… how does that change the downstream computation? And so I guess we're using… it's a very simple test setup.

00:12:06.000 --> 00:12:12.000
More for the interpretability side than for, like, trying to do some application to a protein design.

00:12:12.000 --> 00:12:13.000
Does that answer the question?

00:12:13.000 --> 00:12:22.000
Um, sure, yeah, I think that's… that's good, but I think you can also try to do that on the simpler models that just spit out, like, where these things are.

00:12:22.000 --> 00:12:23.000
And, um… but okay, no, I think that's good for now.

00:12:23.000 --> 00:12:29.000
Oh… I see what you're saying, yeah, that makes sense. Okay.

00:12:29.000 --> 00:12:42.000
Yes, but so here you have a bunch of data, a bunch of sequences with the hairpin and a bunch of with the alpha helix? Yeah. And these are potentially very different. Yes. And then you take an average vector.

00:12:42.000 --> 00:12:50.000
Like, how do you do the patching? I see, yeah, so it'll just be… I'll have a single alpha helix, and I'll just sample.

00:12:50.000 --> 00:13:05.000
the sequence representation from a single hairpin. Oh, a single from a single one. Yeah, I test over many of them. So we collect a bunch of these alpha helices, and then for a given alpha helix, I just find some random hairpin, and I just patch that representation in.

00:13:05.000 --> 00:13:34.000
Which tokens? Yes, yes. Um, yeah, any other questions? Just a question, please. So, when you do the sequence patching, it's only on the last layer of the folding trunk? Uh, so… Yeah, that's a good question. Could you, like, do it, like, in one layer in the middle, and then it starts, like, it understands that you changed something? Right, right. That's the, um… that's actually the whole reason why we're doing this type of patching, is because we're interested in understanding, like.

00:13:34.000 --> 00:13:41.000
Uh, which layers are these representations useful? Yeah.

00:13:41.000 --> 00:13:57.000
If the hotel as metrics, actually, as encoding some kind of, like, chemical or physical interaction between different components of the sequence? What was the first part? You asking about a matcher? Why is the matrix, that is the output of this, like, folding plan. Uh-huh.

00:13:57.000 --> 00:14:08.000
Whether it encodes, like, some kind of, like, physical or chemical interaction between, like, different components. Yeah, I think… that's a great question. I think it does, and, like, that's something that we look at here as well.

00:14:08.000 --> 00:14:23.000
So maybe a follow-up question. So, like, is, like, the intervention that causes the model to suddenly output the new structure? Is it kind of like chemically like faithful in the sense that if you see this kind of interaction as you expect to see this kind of structure?

00:14:23.000 --> 00:14:39.000
Oh, I see. Yeah, and um… that's something that we end up looking at downstream, but we see some… we see evidence of, like, yeah, if we do some chemical intervention, then we'll see faithful downstream outputs, even though… although the outputs sometimes are not, like.

00:14:39.000 --> 00:14:58.000
Perfectly realistic. Yeah, let's go ahead. Okay. Yeah, so what we end up doing is we get a bunch of these alpha helices, and then we just… For a given alpha helix, I'll just find some random hairpin, patch the representation, add a single block, and then we're… and then we'll just see if the output has a hairpin.

00:14:58.000 --> 00:15:13.000
So, this was actually… what happens is… We noticed that sequence patching only works in the first 7 blocks. And then we see a second success regime where pairwise patching begins to work only in the second half of the model.

00:15:13.000 --> 00:15:29.000
And this was actually, like, the critical observation that informed, like, the rest of our experiments. Because we see such a clear two-stage setup, it kind of helped us develop some hypotheses about maybe what's happening about the model within the model. So when we saw this.

00:15:29.000 --> 00:15:45.000
We thought, okay, well, maybe what's happening in this first stage is the model is taking information from the sequence, and then it's using that to update the pairwise representation so that it has a faithful chemical representation of.

00:15:45.000 --> 00:16:03.000
Everything that's happening in the protein. And then maybe what's happening in the second stage is, okay, this is where the model is really beginning to decide the spatial features of my protein and really just understand what the structure should look like. And so this kind of makes sense because, OK, if stage two is where the spatial features are determined.

00:16:03.000 --> 00:16:07.000
Then, if I just patch the spatial representation of my hairpin.

00:16:07.000 --> 00:16:13.000
At this part, at this stage, then I'm just injecting the structure into my protein.

00:16:13.000 --> 00:16:34.000
So, once again, I'm just going to look at the model as a whole, and… Specifically, we're looking at this folding trunk, and I think something that's important to understand is how the folding chunk operates. So the folding trunk actually is composed of 48 of these blocks. And so they happen in sequence. And what's most important to understand is, well, one.

00:16:34.000 --> 00:16:51.000
All that's happening in the folding trunk is we perform self-attention over the sequence representation, and then something called the triangular update over the pairwise representation. And so that's just, like, we're doing attention over the rows and the columns, and there's some other operations. But the most important piece is to understand that.

00:16:51.000 --> 00:16:57.000
The pairwise and the sequence representations communicate with each other during the folding process. And something that.

00:16:57.000 --> 00:17:10.000
And so that kind of informs why we think, okay, well, the sequence representation is probably talking to the pairwise in the early blocks, and then something else could be happening later.

00:17:10.000 --> 00:17:18.000
Um, oh, no, it's a metric. Yeah. Why is it asymmetric?

00:17:18.000 --> 00:17:42.000
Yeah, I mean, it just tells you about. Oh, I see. Um… Well, I guess, yeah, similarity is, like, maybe… Yeah. So there's no like enforcing of symmetric of having a symmetric 2D matrix here. I think that's probably just a design choice from the people who made the model.

00:17:42.000 --> 00:17:48.000
Well, you can just… are you asking about symmetriness of the pair representation?

00:17:48.000 --> 00:17:59.000
then I can just imply some order, right? So ij does not need to be equal to Ji.

00:17:59.000 --> 00:18:08.000
Yeah, and I guess maybe… yeah. The order matters, right? It's a sequence. Yeah.

00:18:08.000 --> 00:18:13.000
I mean, in terms of encoding the sequence of the.

00:18:13.000 --> 00:18:31.000
the input. Yeah, that's a really interesting point. We also looked at that, and there's a lot of choices in the model that are somewhat ad hoc, but that's that's like something that shows up. Quick question. So how is the triangular update defined? It's…

00:18:31.000 --> 00:18:37.000
So the the the triangular update, like the main motivation there is.

00:18:37.000 --> 00:18:56.000
They perform attention operations. Yeah, overpairs. So I see. So I will compute the attention score between, say, iJ and LK. Yeah, I think it's they'll do — yeah. Yeah, it's triangular because like.

00:18:56.000 --> 00:19:04.000
Any update that's performed to like ij and then jk, they'll take into account for the update to Ki.

00:19:04.000 --> 00:19:09.000
or IK. It's like, yeah, they consider triplets. Yeah.

00:19:09.000 --> 00:19:23.000
So they will produce a yeah 3D. Yeah. Yeah, but for this case, we don't look at the triangular updates or self-attention that much, so… but that's just, like, an interesting thing to understand.

00:19:23.000 --> 00:19:34.000
Okay, so now we're going to start looking at the mechanisms. So here we suggested that in the early blocks, there's likely information biochemical signal flow from the sequence into the pairwise representation.

00:19:34.000 --> 00:19:51.000
And when we look at the folding block, well, this is definitely happening through this seek the pair update. And all that is is really we take our sequence representation, we basically perform an outer product-ish operation so that it becomes a 2D matrix, and then we just add that to the pairwise representation.

00:19:51.000 --> 00:20:08.000
So, there is this question of, like, well, what type of biochemical information do we expect is flowing? And when it comes to protein folding, something that we'd expect is like the electrostatic charge, because if I have one residue that's positively charged, another residue that's negatively charged, and they're close together, they'll probably come and reduce the distance between them.

00:20:08.000 --> 00:20:26.000
Or we'll have repulsion between two residues that are of the same charge. So here, when it comes to amino acids, there's three canonically positively charged amino acids and two negatively charged ones. So just as, like, a first observation, well, we can probably just look for this in terms of using linear probing.

00:20:26.000 --> 00:20:38.000
So what that really means is we just train a linear classifier on the seek to pair representation and the pairwise representation, just to see, well, can we actually extract charge information?

00:20:38.000 --> 00:20:46.000
Uh, for pairs of residues. Uh, so all we find, really, is that, yeah, in the basically the seek-to-pair update.

00:20:46.000 --> 00:20:52.000
Includes information about the charge. And what's most interesting is that plot on the right, where we have.

00:20:52.000 --> 00:21:08.000
The test accuracy jumps the most after the first 7 blocks, which is also when we found the sequence patching to work. And so this is just giving us some more motivation to think, OK, charge would probably be something that is propagating from the sequence into the pairwise representation in this early stage.

00:21:08.000 --> 00:21:26.000
And this is to be expected. But what we're more interested in is, well, let's look for linear representations of charge in our latent space, and then can we use that to steer the model outputs? So here's something that we can do. Well, we have this sequence, and we can put it into our model. And we can just extract the sequence representations.

00:21:26.000 --> 00:21:42.000
of positively charged and negatively charged residues. And what we can end up doing is we can look at the mean representation of all of my positively charged and all the mean representation of all my negatively charged residues in latent space. And if we just draw a vector between them.

00:21:42.000 --> 00:22:06.000
Well, this is just a hyperplane that separates the two concepts. And what we would expect is, well, now if I take a new protein and I apply this vector to the sequence representations, then… If I add it, it'll like, it'll make the residue look more positively charged and if I subtract it, it'll make the residue look negatively charged. So now the question is, does the model actually use charge? So…

00:22:06.000 --> 00:22:18.000
Something that we want to do in mechanistic interpretability is always confirm that our representations are causally meaningful, so that the model actually uses — looks at this representation. So here, what we can do is we can look at.

00:22:18.000 --> 00:22:38.000
Helix loop helix motifs and our proteins. And we can just make all the residues in this first helix positive, and all the residues in the second helix negative. And so what we would expect is that if this is a charge representation, and if the model actually uses this, then the two helices should come close together.

00:22:38.000 --> 00:22:42.000
and form what looks more like a beta herpen.

00:22:42.000 --> 00:22:52.000
And so the way that we can measure this is just by looking to see if the two strands form what's called hydrogen bonds, which are those bars in purple.

00:22:52.000 --> 00:22:58.000
And so this is question. And now we see we apply the intervention, and we find that the success rate.

00:22:58.000 --> 00:23:07.000
Peaks within the early blocks, and we apply the intervention across… with a sliding window. But this kind of matches our expectation. So if we apply this.

00:23:07.000 --> 00:23:12.000
Vector to sequence representations. In the early blocks.

00:23:12.000 --> 00:23:27.000
It looks… it produces outputs that look like a hairpin. But something that we can also do is sort of an opposite experiment. So what happens if I take a beta hairpin, and then I make both of the strands look like they're positively charged? Well, we would expect that the two strands would begin to repel.

00:23:27.000 --> 00:23:38.000
And so what we find is that this is actually the case. So here's a plot where the red is making both of them negative. And we find that if we apply the same vector and make residues look.

00:23:38.000 --> 00:23:51.000
Both strands look negative. The two strands will repel, as we hope to see. And so I think the most… The important piece is that this effect is focused in the early blocks, which aligns with what we proposed the early block mechanism would work.

00:23:51.000 --> 00:24:06.000
And yeah. So I don't know about this hydrogen bonding. So is that something just geometrically you can measure if there's hydrogen? Yeah. So it's like there's a hydrogen atom in the middle. The way that it's commonly measured here is by taking.

00:24:06.000 --> 00:24:18.000
Uh, is by measuring the distance between the nitrogen and the oxygen atom, which is, like, which are always present in all of the… in all amino acids, and then that's where the hydrogen bond will form.

00:24:18.000 --> 00:24:37.000
Yeah. Can you compel the headplot with what you would expect to see based on some, like, simulation? Yeah, that's a good question. We thought about doing that, um… I guess the… We didn't end up going down that route, mainly because we weren't sure, like, how to…

00:24:37.000 --> 00:24:52.000
Calibrate that in terms of, like, oh, I don't know what… As I steer like what the charge like value will look like. But that was something that we thought about doing, because, well, let's say we find that this plot correlates very well with a standard.

00:24:52.000 --> 00:25:02.000
Like, what the physics simulation suggests, then that's suggesting that, well, the model not only has learned charge, but also the physical properties of charge. So yeah, that's a good question.

00:25:02.000 --> 00:25:03.000
Yeah.

00:25:03.000 --> 00:25:08.000
Um, I have a question here, like, I remember, like, all the folding modes actually does not really include any hygiene items, right?

00:25:08.000 --> 00:25:10.000
And, like, what do you mean by hydrogen bond here?

00:25:10.000 --> 00:25:11.000
You only have Javier that I'm seeing the models, right?

00:25:11.000 --> 00:25:14.000
Yeah. Yeah.

00:25:14.000 --> 00:25:19.000
Yeah, yeah, that's a good point. So, uh, the… I guess the hydrogen bonding is.

00:25:19.000 --> 00:25:26.000
And, like, the way that we measure it is not necessarily, like, oh, we look for hydrogens here. And more so just, like, that canonically.

00:25:26.000 --> 00:25:36.000
If… if our protein… okay, I'm just going to skip that, sorry. So if the protein ends up making a hydrogen bond, then.

00:25:36.000 --> 00:25:54.000
The distance between the two… between the nitrogen atom of one residue and the oxygen atom of the other residue is going to be like 3.5 angstroms or less. So that's just the way we measure it. But I'm also going to be honest and say like I'm —.

00:25:54.000 --> 00:25:55.000
notes.

00:25:55.000 --> 00:25:56.000
Not super like the the biology and like the chemistry behind it is not my strongest suit. But yeah.

00:25:56.000 --> 00:26:02.000
I see. But I guess, like, my major question is about, like, during the whole training of the models, like, the model has not really accessed any of the hydrogen bond information during its training, right?

00:26:02.000 --> 00:26:06.000
Yeah. Oh, I see. Yeah, that's true.

00:26:06.000 --> 00:26:09.000
And how do you imagine, like, to learn something about, uh…

00:26:09.000 --> 00:26:17.000
the hydrogen bonding stuff, because, like, hydrogen bonding is not only about distance, it's all about, like, the orientation of the bonds, right? It's actually very, like, uh…

00:26:17.000 --> 00:26:18.000
like, angle-related stuff.

00:26:18.000 --> 00:26:20.000
Right. Yeah.

00:26:20.000 --> 00:26:26.000
Do you think I get to learn this kind of information from only Abby Adams from the… from training datasets, or…?

00:26:26.000 --> 00:26:27.000
It's…

00:26:27.000 --> 00:26:33.000
Yeah, I mean, I guess my… my intuition was always, like, yeah, we're… even though we're only looking at.

00:26:33.000 --> 00:26:53.000
Like, these PDB structures, which are pretty coarse, um… I guess, like, the… there's… the pattern is, like, implicit within all the structures, and so the model just picks up on things like this. Yeah, it's pretty — I guess it's like a note to how powerful these deep learning approaches can be.

00:26:53.000 --> 00:27:02.000
Yeah. What's the final representation of the geometry? Is it like angles?

00:27:02.000 --> 00:27:18.000
It's like XYZ coordinates of all the atoms in the — and then maybe I missed this, but in each layer, the representation is kind of the same, right? It's the same format. So can you?

00:27:18.000 --> 00:27:33.000
Just, like, dump out the coordinates in the middle of the… Yeah, yeah. So we've… the intermediate decoding is something that we've looked at. Um, do you have a follow-up question for that, or is it… oh, yeah. Um, but something that we do notice is, like, sometimes it's out of distribution, um, even… even if we, like, do some.

00:27:33.000 --> 00:27:44.000
Clever, like, layer norming. The structure renderer usually only works in, like, the second half of the model, but yeah.

00:27:44.000 --> 00:27:46.000
Just plug it into the middle and it'll work.

00:27:46.000 --> 00:28:02.000
Yeah, it'll, like, produce… sometimes it'll produce, like, decent structures. And that's something that we also — we look at something very similar for the next experiment. But yeah, here's just some more verification of the mechanism. So this was something that we just like.

00:28:02.000 --> 00:28:11.000
We hope to see, or if we apply this intervention to make all the residues in strand 1 positive, then hopefully my linear probe agrees with this. And so that's the case that we see here.

00:28:11.000 --> 00:28:18.000
Um, and something else that we also checked was, well, what if it's… what if we just apply the intervention to the sequence, and then we don't let.

00:28:18.000 --> 00:28:35.000
the sequence talked to the pairwise representation. If the output still follows still forms hairpin-like structures, then this mechanism is not what we expect. But we actually see, OK, in the early blocks, if we just prevent the sequence from talking to the pairwise representation.

00:28:35.000 --> 00:28:50.000
Then we don't see the expected hairpin formation. But this is just going to be a visualization of this charge mechanism now. So here we have a toy protein, which is just a single hairpin. And then I'm going to apply the intervention.

00:28:50.000 --> 00:29:06.000
From block 0 to 15. And we're just going to slowly increase the magnitude of the vector that we're adding. And so you can see that, like, as I make both the strands more positive, they begin to repel, and the distance between them increases. So it's like a pretty nice visualization of what's happening.

00:29:06.000 --> 00:29:24.000
Within our proteins. Uh, sorry, so each arrow is composed of multiple atoms, like in this case 4. Yeah, that's 4 residues, and then each residue has, like, maybe 15 to 20 atoms in it. Yeah. Okay.

00:29:24.000 --> 00:29:37.000
And you just visualize them as arrows? Uh, yeah, that's like the cartoon representation. But yeah, it's helpful to see this as opposed to the actual atoms themselves.

00:29:37.000 --> 00:29:43.000
And I guess, like, technically, because the atoms are, like, fields, so it would be blobby and weird. But yeah.

00:29:43.000 --> 00:30:00.000
This is only for beta. When you have arrows, and you already have this beta thing, or it could also be alpha. Yeah, I guess. And yeah, you would just assume like this is beta content, and that's when the arrows show up. Yeah. So it's like you're already taking something that would fold.

00:30:00.000 --> 00:30:19.000
Together, and then you're trying to… Break it apart, but it will not change the nature. Uh, yeah, I think technically this visualization is, like, um… Not entirely accurate, but… Yeah, it's like a like it actually shouldn't be arrows anymore, but it's just there so you can see it.

00:30:19.000 --> 00:30:23.000
Okay. So yeah, we we have this mechanism now.

00:30:23.000 --> 00:30:36.000
Now, something else that we looked at. Well, if we looked at this first stage, then what's really going on in the second stage? And so here we're assuming we're hypothesizing that this is where the structural decisions are being made. So.

00:30:36.000 --> 00:30:44.000
Something that we can do — I suggested like the intermediate decoding is, well, then we can just basically say, well, what if we just look for?

00:30:44.000 --> 00:30:53.000
This is… pairwise distances in the pairwise representation. And so we just train a linear probe where it's just a simple linear regression. And what we find is.

00:30:53.000 --> 00:30:57.000
Sort of what we expect is that, well, the accuracy of this probe increases.

00:30:57.000 --> 00:31:04.000
As we go down the layers and the blocks of the folding trunk.

00:31:04.000 --> 00:31:05.000
Which is definitely what we would expect. But the question that I'm more interested in is, well, if I perform an intervention with this, does this help us understand when the structure is decided? Yeah. I'm sorry. What's d hat? Oh, sorry. So this is the distance between.

00:31:05.000 --> 00:31:19.000
Okay, thanks.

00:31:19.000 --> 00:31:26.000
residue I and J. And then we all we take is the pairwise representation between i and j.

00:31:26.000 --> 00:31:40.000
And just to further clarify, so for each block capturing a different code. And so it's assuming you're training the best one that you're comparing which one gives you the best tests.

00:31:40.000 --> 00:31:48.000
Uh, I have a question here, like, because basically, like, the pair representation is not symmetric, as you just mentioned, right? But D matrix is actually symmetric.

00:31:48.000 --> 00:32:00.000
So how do you, like, do that here?

00:32:00.000 --> 00:32:01.000
Yeah, I would say…

00:32:01.000 --> 00:32:02.000
Yeah. Yeah, that's a pretty good question. I mean, here, we just take… Well, I guess, um, is the question, like, how do we account to make sure, like, there's no…

00:32:02.000 --> 00:32:08.000
Because, like, in your Pell representation matrix, like, ZIG does not equal to ZJI, right?

00:32:08.000 --> 00:32:09.000
like, in a real distance matrix, DIG actually have to equal to DJI.

00:32:09.000 --> 00:32:12.000
Right. Right.

00:32:12.000 --> 00:32:17.000
Because, like, business metrics are kind of symmetric. So how do you do that in practice?

00:32:17.000 --> 00:32:26.000
I guess we didn't take anything special into account. Like, I'm not sure what the.

00:32:26.000 --> 00:32:27.000
Yeah, I guess.

00:32:27.000 --> 00:32:33.000
Um, I guess, like, my general question would be, like, because, like, to compute the distance DIJ, you have to

00:32:33.000 --> 00:32:37.000
Uh, the… you can use… you can use DIJ or use DJI.

00:32:37.000 --> 00:32:40.000
Because, like, as the end, it should cost money the same.

00:32:40.000 --> 00:32:42.000
Due to average, or…?

00:32:42.000 --> 00:32:47.000
you choose one of the two.

00:32:47.000 --> 00:32:48.000
Yeah, yeah, yeah.

00:32:48.000 --> 00:32:51.000
for the average. Yeah, yeah, I think we… uh, oh, oh, I see, in terms of, like, the performance, like, we… yeah, I think we just take the average.

00:32:51.000 --> 00:33:14.000
Oh, you take the average of… Oh, so we… the probe will, like, we'll have one case where the probe will take zij, and then… Oh, I actually remember now. Sorry. Both directions encode the same… Yeah, both. But they might have other asymmetric stuff, but… Yeah, I think in this case, so what we actually end up doing is we only… we basically just look at, like, I think the lower triangle here, where it's… we'll only pass in, like, zij.

00:33:14.000 --> 00:33:20.000
And then we just ignore ZJI. Yeah.

00:33:20.000 --> 00:33:29.000
Okay. Uh, I mean, we could just, like, redo it and check. I guess I'm pretty sure, like, it would just follow the same pattern, yeah.

00:33:29.000 --> 00:33:44.000
Um, okay. So, yeah, now we're curious, like, when does the model decide structure? And the way that we could do this is… Well, we have this probe, which predicts pairwise distances, and so now we can look at.

00:33:44.000 --> 00:33:59.000
Once again, the alpha helix… the helix loop helix motif that shows up in a lot of our target proteins. And we can just say, like, well, what if we just use the probe to decrease the distance between the two proteins? So here, on the left, we're visualizing the cross-strand distances.

00:33:59.000 --> 00:34:07.000
Basically, like, EG is the distance between atom E and atom G. And so here we can just see, like, OK, they're all pretty far apart.

00:34:07.000 --> 00:34:23.000
Uh, here's, like, 16 angstroms. So, the goal… the question is, well, can we just, like, force the two strands to come together? Uh, and here we can just use essentially its gradient steering. But because our model is a linear arrestor, then the gradient is just the weight space. So we're just going to just subtract the weights.

00:34:23.000 --> 00:34:35.000
from my parents representation. And the goal is to sort of force them closer together to form sort of like a hairpin-like structure again. So… and we can see, like, now the… now EG.

00:34:35.000 --> 00:34:47.000
It's supposed to become a lot closer together, and so that's why it's dark here. So that's what we're expecting to see. And once again, we measure the hairpinness using the hydrogen bonds again. Yeah.

00:34:47.000 --> 00:35:02.000
I'm not sure, like, I understood what you just did, like, could you re-explain, please? Yeah, um, so we're going to do an intervention on the pairwise representation. And we just trained linear probes to predict distance from that pairway space again.

00:35:02.000 --> 00:35:22.000
So… Now what we want to do is — our goal is to decrease the distance. And then eventually we want to see, OK, we're going to apply this with a sliding window. So we want to see when does this intervention occur most effectively? That's the end goal. So in order to apply the actual intervention, what we're going to do is.

00:35:22.000 --> 00:35:29.000
During the forward pass, all I'm going to do is, for any given Zij, for, like, these.

00:35:29.000 --> 00:35:39.000
Like, any given representation where it's, like, relating residues in the first helix to residues in the second helix, I'm just going to subtract the weights.

00:35:39.000 --> 00:35:51.000
of the linear probe from that representation. And, like, in theory, that should just decrease the distances, uh, in the same way that the gradient steering would work. Um, and so does that make sense?

00:35:51.000 --> 00:35:56.000
Okay, and so we just… we…

00:35:56.000 --> 00:36:05.000
Right, and so I guess, like, we've been thinking about how to… Like, we've been thinking about how does the model form hairpins?

00:36:05.000 --> 00:36:21.000
So I guess here, like, the previous experiment with the charge, we were trying to force the model to make a hairpin by just editing the chemical information. And so here we're just going to physically just push the distances together to see if we can make it happen again.

00:36:21.000 --> 00:36:28.000
Do you know anything about how the charge stuff affects this W direction?

00:36:28.000 --> 00:36:43.000
Maybe that's Marcus. Oh, I see. Yeah, we didn't look at the charge. But yeah, that's… I think that's an interesting circuit to look for, which is, like, what are heads that are attending to information that's related to charge? Or if I edit the charge of one residue, will I begin to attend to it more now?

00:36:43.000 --> 00:36:55.000
Or does it change the activity? Like, yeah, the embedding's in this W direction. Yeah, oh, that's… I see what you're saying, yeah. That would be pretty interesting to look for. But yeah, so we apply this intervention.

00:36:55.000 --> 00:37:12.000
And following the pattern, we noticed that, like, this intervention, when we apply it with a window of 15, is most effective, once again, in this second stage. But actually, right where the pairwise patching was occurring. And so this is yet again, another case where, OK, it's very likely then that this is where.

00:37:12.000 --> 00:37:23.000
The structure is being decided, at least in the pairwise space. And so if we edit the structure in this window, then the output structure is going to follow what we hope for.

00:37:23.000 --> 00:37:43.000
So, yeah. And so this is just like a visualization of what that, um… What that process sort of looks like. So here, we're going to… Take the alpha helix, and then I'm just going to force the residues to be closer together. And it doesn't actually form, like, the arrows, but it just, like, looks sort of hairpin-y. So that's, like, the…

00:37:43.000 --> 00:37:57.000
That was the goal. So, and you can look at the… those are the cross-strand distances. So we're hoping to see dark diagonals, which will represent the cross trend distances decreasing. So that's what we expect, and that's what we end up seeing.

00:37:57.000 --> 00:38:22.000
Yeah. Okay. So… Yeah, just like us in sort of like a… High-level overview of what we just looked at. So, we started out with trying to ask, like, when and where is hairpin information computed? And so what we… what we saw is, yeah, in the early blocks, there's some sort of hairpin information in the sequence, and in the linear box.

00:38:22.000 --> 00:38:32.000
We see that there's spatial features. And then we took this step of, well, can we actually confirm that this is the mechanism? And so we looked at mechanism 1, which is a biochemical signal flow, and we find that it aligns with.

00:38:32.000 --> 00:38:46.000
Our sequence patching results. And then we look at mechanism two, which was structural spatial feature or structural decisions, and that once again lines up with our pairwise patching results. So I think it's pretty pretty, yeah.

00:38:46.000 --> 00:38:57.000
Um, so can you explain again the first picture? So, basically, when you intervene, um.

00:38:57.000 --> 00:39:08.000
at this sequence representation, um… Everywhere, like, after level 10, then, uh, you don't get the hairpins.

00:39:08.000 --> 00:39:12.000
And then how did you jump into a biochemical signals?

00:39:12.000 --> 00:39:37.000
Oh, I see. You're saying, like, when we saw this, like, why did we just assume that this is, like, that was the case? Yeah, and I think that was… the intuition for that came from looking at the… Basically, the structure of the folding block, along with the fact that the pairwise representation is initialized basically as empty. So, like, the only way that the pairwise information, or the pairwise representation is populated at the beginning is most likely due to the fact that.

00:39:37.000 --> 00:39:43.000
The sequence will add some updates. And so… because otherwise, like, the… how will the pairwise.

00:39:43.000 --> 00:39:59.000
operations understand anything about the protein? Uh, and so that was… I guess that was, like, our intuition. Uh, and then the second half was like, okay, it's very likely just going to be biochemical signals along with some other stuff that's going to be flowing into the pairwise representation at the beginning.

00:39:59.000 --> 00:40:15.000
Yeah. So this does the ESM embedding have a lot… probably has a lot more than just charge, right? Yeah, yeah, yeah. So it's just one direction that you could look for. Do you feel like that other stuff might not be as powerful as this charge stuff, or… Um…

00:40:15.000 --> 00:40:31.000
I think… so, part of the reason why we went for Charge was because it was a lot more obvious to, like, how to actually isolate that signal. Something else we were looking… potentially looking at was like, oh, the hydrophobicity is very important to protein folding. So that's something that we could have looked at. But we just.

00:40:31.000 --> 00:40:43.000
Yeah, we just chose to look at charge, because it's very intuitive, too. Like, if I make one residue, like, hydrophobic, then, like, it is not quite clear, like, how the protein will end up changing.

00:40:43.000 --> 00:40:49.000
So, yeah, this is just, like, an overview of the mechanisms that we just looked at.

00:40:49.000 --> 00:40:58.000
And so, yeah, I guess thinking about what's next. So this is basically some more like BioE experiments that we're thinking about doing.

00:40:58.000 --> 00:41:12.000
Something that happens with these proteins is, like, they're not actually static. So they take on many conformations, and so there's this potential, like, oh, when we take a… when we do a forward pass of a protein that has many conformations, or even disordered regions, like, can we somehow predict that?

00:41:12.000 --> 00:41:28.000
from the representations, like the pairwise representations, and can we even, like, do some interventions to explore other possibilities for protein structures? Which is already being done using MD trajectories, but this is just one potential example. Another idea is, well.

00:41:28.000 --> 00:41:48.000
We have this interpretable human charge representation. Well, what if we just use this to simulate environmental conditions now for protein folding? So proteins will change depending on maybe if they're in solution and it has a different pH. So we could apply the charge vector to all these different residues to now extend the model capabilities. Because the model was never trained.

00:41:48.000 --> 00:41:56.000
to account for environmental conditions. But we can extend that capability by using this interpretable representation.

00:41:56.000 --> 00:42:13.000
Um, but I think something else that's important is we talked about these two questions of, like, what will interpretability give us here? Which was, will, one, understanding the model internals, and two, being able to steer. But I think the third piece that is really important when it comes to AI for science is.

00:42:13.000 --> 00:42:24.000
Can we discover new science, or even discover new mechanisms of biology? I think here in terms of the protein space, it's less clear. But something else that I'd love to explore.

00:42:24.000 --> 00:42:39.000
We'll probably be looking at, like, these virtual cell components. So, something like looking for… looking at single-cell perturbation models, uh, and, like, finding new gene perturbation networks, stuff like that. And there's always a question of, can we do this in an unsupervised fashion because.

00:42:39.000 --> 00:43:04.000
This project relied on a lot of human intuition, and there's this question of, well, if you keep applying human intuition to a model, how will you ever find something new? But yeah, that's just where… This project is out right now. So check out our website at the Bao Lab, and yeah, thank you guys for your time.

00:43:04.000 --> 00:43:20.000
Any questions? Yeah. Hi. Thanks for a great talk. Can you please explain? So for that matrix, so it would make sense to make it symmetric, right? Because it's the same sequence. It's not like.

00:43:20.000 --> 00:43:39.000
keys and queries, right? Yep. Okay. So, why don't you just take an average of itself and it transpose, and, like, work with AI? Yeah, and I think… so there is one… If I was to, like, make ESM fold, I do think…

00:43:39.000 --> 00:44:02.000
They had a statement about why they actually wanted the asymmetry. And I guess their idea was like, oh, maybe this asymmetry will allow the model to capture a richer representation than if you force symmetry. I think symmetry would be a nice to have.

00:44:02.000 --> 00:44:03.000
But the thing is, is that…

00:44:03.000 --> 00:44:07.000
But if we're hoping for the model to capture new relationships that maybe we don't understand, like maybe having more representation space is helpful. But honestly, from a geometric deep learning perspective, they probably should have just made a symmetric.

00:44:07.000 --> 00:44:11.000
But the whole thing is not symmetric, so if you have, like, an alpha…

00:44:11.000 --> 00:44:13.000
Helix. This is a right-handed helix that has a handedness. There's a directionality to it.

00:44:13.000 --> 00:44:17.000
Yeah.

00:44:17.000 --> 00:44:18.000
So, from residue to residue, if you invert the whole thing, it's not the same thing.

00:44:18.000 --> 00:44:23.000
That's true, yeah.

00:44:23.000 --> 00:44:45.000
Yeah, I see. Yeah. Yeah, and I guess, like, even the way that the proteins fold, like, in… I don't understand what is that matrix doing, like… What is this direction ij to Ji? Yeah. I thought it's, like, some kind of notion of distance or similarity between.

00:44:45.000 --> 00:45:05.000
those positions. Yeah. I think… that direction. So maybe I think you can think about it as just, like, a feature encoding, right? Rather than encoding just individual I'm not a bio person, and then now you're just encoding pairs of them.

00:45:05.000 --> 00:45:06.000
And the pair does not have to basically. So it's like, when we model language, right, we can encode each word individually as a token, but we can encode phrases, pairs of words, that's called, like, bigrams and trigrams and whatnot. Yeah, and there is…

00:45:06.000 --> 00:45:26.000
So I think, like, it's good to be able to capture that.

00:45:26.000 --> 00:45:29.000
There is meaning, right? Because they just point out it's different physically.

00:45:29.000 --> 00:45:34.000
I guess you can think of this as a directed graph.

00:45:34.000 --> 00:45:37.000
And, like, chemically, like, you have, like, um…

00:45:37.000 --> 00:45:40.000
I guess elements, I guess items that…

00:45:40.000 --> 00:45:42.000
blend electrons and atoms that…

00:45:42.000 --> 00:45:50.000
take them. So that's… there's definitely directionality here. And to have… thinking about this as a directed graph is a… is a good way of thinking about it.

00:45:50.000 --> 00:45:53.000
And, uh, this is the representation of…

00:45:53.000 --> 00:45:59.000
The edges of that directed graph.

00:45:59.000 --> 00:46:05.000
There's only more than distance, probably going… Oh, yeah, there's definitely more than distance.

00:46:05.000 --> 00:46:23.000
Now, did you try cyclic proteins, and then you should see in the… That's… we didn't look at… yeah, we only looked at, like, secondary structure less, like, but I think looking at cyclic proteins would make a lot of sense.

00:46:23.000 --> 00:46:40.000
Also, how are you initializing the pairwise matrix? Is it just, like, identity or random, or… It's just initialized with some positional embeddings. So it's like, basically, like, if I have, like, um, residue ij, then, like, I have an embedding that.

00:46:40.000 --> 00:46:44.000
Tells you how far apart they are, actually, in, like, in the sequence.

00:46:44.000 --> 00:46:51.000
Yeah.

00:46:51.000 --> 00:47:07.000
Do you know how well PNNs perform on that type of task? Or if it's used at all? Yeah, I know for like DNNs are used a lot for sequence design. I think for right now, the main paradigm is still just like.

00:47:07.000 --> 00:47:14.000
We have this… this folding module, and then, like, diffusion afterwards. But yeah, there was a phase.

00:47:14.000 --> 00:47:29.000
15… Yeah, protein MPNN is graph-based. Yeah, a lot of, like, the protein representation is often… lends itself to graph representations, so… For sequence design, people usually use graphs.

00:47:29.000 --> 00:47:36.000
Alright, okay, thanks again, Kevin. And we'll go upstairs for lunch if people want to join, feel free to join. Thank you.

00:47:36.000 --> 00:47:50.000
Mm-hmm