BERTology Meets Biology: Interpreting Attention in Protein Language Models (Paper Explained)
deep learning machine learning arxiv explained neural networks ai artificial intelligence paper bert transformer mlm language model masked language modeling proteins protein amino acid primary secondary tertiary structure helix strand band sheet turn binding site contact map dna rna amino acids proline phenylalanine
Proteins are the workhorses of almost all cellular functions and a core component of life. But despite their versatility, all proteins are built as sequences of the same 20 amino acids. These sequences can be analyzed with tools from NLP. This paper investigates the attention mechanism of a BERT model that has been trained on protein sequence data and discovers that the language model has implicitly learned non-trivial higher-order biological properties of proteins.

OUTLINE:
0:00 - Intro & Overview
1:40 - From DNA to Proteins
5:20 - BERT for Amino Acid Sequences
8:50 - The Structure of Proteins
12:40 - Investigating Biological Properties by Inspecting BERT
17:45 - Amino Acid Substitution
24:55 - Contact Maps
30:15 - Binding Sites
33:45 - Linear Probes
35:25 - Conclusion & Comments

Paper: https://arxiv.org/abs/2006.15222
Code: https://github.com/salesforce/provis

My Video on BERT: https://youtu.be/-9evrZnBorM
My Video on Attention: https://youtu.be/iDulhoQ2pro

Abstract:
Transformer architectures have proven to learn useful representations for protein classification and generation tasks. However, these representations present challenges in interpretability. Through the lens of attention, we analyze the inner workings of the Transformer and explore how the model discerns structural and functional properties of proteins. We show that attention (1) captures the folding structure of proteins, connecting amino acids that are far apart in the underlying sequence, but spatially close in the three-dimensional structure, (2) targets binding sites, a key functional component of proteins, and (3) focuses on progressively more complex biophysical properties with increasing layer depth. We also present a three-dimensional visualization of the interaction between attention and protein structure. Our findings align with known biological processes and provide a tool to aid discovery in protein engineering and synthetic biology. The code for visualization and analysis is available at this https URL.

Authors: Jesse Vig, Ali Madani, Lav R. Varshney, Caiming Xiong, Richard Socher, Nazneen Fatema Rajani

Links:
YouTube: https://www.youtube.com/c/yannickilcher
Twitter: https://twitter.com/ykilcher
Discord: https://discord.gg/4H8xxDF
BitChute: https://www.bitchute.com/channel/yannic-kilcher
Minds: https://www.minds.com/ykilcher
 Hi there! Today we'll look at Bertology meets Biology interpreting attention in protein language models by Jesse Vig, Ali Madani, Lav R Varshini, Kaiming Xiong, Richard Sokker and Nazneen Fatima Rajani. This paper is a investigative paper into models that are trained on biological data, specifically into BERT models. Actually into one specific BERT model that is trained on protein sequences. Now it is trained to simply perform language modeling on these protein sequences, but out of this language model you can then inspect this BERT model and read important biological data of these proteins, higher order data from the attention heads of the BERT model, which is pretty interesting. Basically means that the information of these higher order functions is at some point encoded in the structure of the language of the protein sequence. So we're going to go through what this means and how this comes about and what they did in order to investigate. I think this is a pretty cool investigative work and probably very promising for future research. Yeah, as always if you like content like this, consider sharing it out and leaving a like. Also tell me what you think in the comments. So biology. Really quick for people who maybe have never heard this. In your every cell you have this thing called DNA, which basically is an encoding of all of your biological functions. Now usually biological functions are realized through proteins. So DNA is basically a building plan for all of your proteins. This happens in the following two steps. First there is this transcription step where RNA is built. This is basically a copy of your DNA, but it's only single strand as you can see right here. And then there is a translation step that finally translates the RNA into the protein. What will end up is just a sequence of these beads right here. Now these beads are what are called amino acids. So a protein is simply a chain of these amino acids. There are 20 different amino acids and the order of these amino acids in the chain makes the function of the protein. Now specifically we know about these proteins that it seems to be very important how their three-dimensional shape is. So a lot of these different amino acids have different chemical properties. Some are sort of I think negatively charged, some are neutral, some are acids and so on. So they have very different chemical properties. So once you build this protein and you kind of release it into the cell it will curl up into a three-dimensional structure. So this one might be doing something like this and sort of form a circle or something like this. Just because these proteins here they kind of attract each other maybe electrically and thus the protein forms a circle and the function of the protein is very much related to its shape. So if it is a circle it could maybe trap something else in here. So you really have to think of these things like kind of tools. There are proteins that cut other proteins and they are really shaped sort of like a scissor that exactly fits these other proteins such that you can effectively cut them. So sometimes you can substitute an amino acid for a different amino acid like this here. If it doesn't change the shape very often you're fine. The protein function isn't changed. But if you change a different amino acid that is sort of vital to the shape and the shape changes then your protein very often loses function. So mutations in DNA sometimes lead to mutations in protein. Not always because there is some redundancy in this translation step from RNA. But if they do lead to a different amino acid it doesn't actually mean that the function changes. So there is sort of value in analyzing the sequence of the structure of proteins rather than the structure of DNA. Of course it's also important to analyze the structure of DNA but it is equally important to analyze the structure of proteins because not all the information is in the sequence. Not all the obvious information is in the sequence. So what does this paper do? This paper goes and takes a model that has been trained on protein data. So if you look at this protein it is simply a sequence of amino acids and these amino acids they all have names. I think I have a table somewhere here. Yes so these are the different amino acids that exist and you can see a protein is simply a sequence of these names. So usually they're abbreviated by like a three-letter abbreviation or just a one-letter abbreviation. So a protein might be AVMMVAG and so on. And this is just a string of text. So what I can do is I can train a language model on this. A language model is simply a model that takes a piece of text and tells you what's the next piece of text. So what's the next letter, what's the next word, in this case what's the next amino acid. And we can use tools from NLP for that. Specifically we can train a BERT model. Now BERT works a bit differently than a standard language model. BERT does what is called masked language modeling. So you would take this string, you would feed it into a BERT model right here. And I've made an entire video on BERT if you want to check that out. And what you'll do by inputting that you'll mask out some of the tokens. So you'll maybe mask out this one, mask out this one, and then you ask the model to reconstruct those. We say that here is an M and here is an A without seeing them. So the model somehow has to learn from the surrounding amino acids what this amino acid could be. So it has to reconstruct this sequence. So the hope here is, in natural language, is that BERT somehow learns something about language itself. By being able to reconstruct these things it has learned something about language, about which words appear together and when. It might even learn very long distance relationships between words just because it has to predict those. And the idea carries over to biology. So we might hope that a BERT trained on an amino acid sequence will learn something about the language of proteins, about the amino acid sequence. And our goal here is to ask, can we somehow infer the 3D shape of a protein, which is the important part, from its sequence right here? So given its sequence, can we infer the 3D shape? Now as I understand it, usually this has to be done in like a simulation. So you would build this in a simulator and then you do like some sort of a molecule simulation to see how this ends up in a 3D shape. You could train a model to just predict the 3D shape, but in this case we're just interested in what does the BERT model learn about the 3D shape while only ever having been trained on predicting the sequence of amino acids. So it's never been trained to look at the 3D shape. And that's our goal here. So specifically we'll look at two different things. So here you can see examples of proteins and their high-level structures. So in these proteins what you call the primary structure is this sequence of amino acid. This is simply which amino acids are in which order. There is a thing called the secondary structures and we often observe that spans of these amino acids like substrings form these what are called these helixes as you can see here or these sheets. I don't know how they're strands in English. We call them sheets or I think these are the alpha helixes and these are the beta sheets and there is also a turn. I think this here might be a turn. So there are these kind of secondary structures and then the tertiary structure is how these this is still one protein. This is one unbroken chain of amino acid. And you can see this here kind of forms this double ring which would be its tertiary structure. Very important for predicting the tertiary structure is to predict when two amino acids are close to each other. So if we have a chain right here and the chain as we saw before kind of turns and bends on itself then these two amino acids here are very close in close contact. And to predict which amino acids are in close contact to each other helps you determine the the tertiary structure. So that's a consequence of it. So we wonder does BERT know intrinsically which of these amino acids are going to end up being in contact with each other without ever having been trained to do it? The second thing we're interested in are binding sites. So here what you might not be able to see but we made this example before where this sort of forms a loop and then I say can trap something here right like another molecule. And this is what we would call a binding site. A binding site is a one amino acid that maybe through the structure of the surrounding amino acid as well but also through its properties and how it is exposed in 3d shape acts as sort of a receptor for other molecules. It binds to other things. So think of your hemoglobin that traps the oxygen in your blood or something like this. It is where a chemical reaction or a reaction with something else will happen. That's a binding site. And we are interested does BERT, the BERT that is only trained on a language modeling objective, know which ones are the binding sites? Because you know that would be very interesting and not something BERT was trained on. By the way I particularly liked Richard Sockers tweet on this. I think he tweeted out, oh BERT trained only on language model can predict binding sites and biological properties and formulated it like it was you know like GPT-3 was formulated like if we train on Wikipedia our model can do math. I thought it was kind of a satire headline. If we train on Wikipedia our model can predict biology and also it can tie your shoes and cook your dinner. So it's trained on language modeling on biological data and now that makes sense. So they're going to look at two different things or actually more than two different things but they formulate this in an abstract way right here. So what they'll look at is the so-called properties. A property F and this property F can be for example that a amino acid is a binding site. The property F can also be that two amino acids are in contact with each other. So F always takes I and J. If in the case for example where this is the contact property then it simply is the indicator function for when I and J are in contact. And if it is a just a binding site then I think we are looking at J. At the token level property we define to be an indicator that returns one if the property is present in token J. So whenever J is a binding site then that holds. So what we're looking at are these attention heads in BERT. If you don't know BERT has an attention mechanism which basically means from layer to layer each token can attend to all other tokens. So here the amino acid sequence I've drawn it twice and the next layer representation of this amino acid will be able to gather information from all of the other amino acid through an attention mechanism through a dynamic routing algorithm. I've made a video on attention is all you need if you want to find out more how this works. Now what we're interested in is the strength of these connections. So the hypothesis is if molecule 1 and 3 are contact sites then maybe we will find a layer where this connection between 1 and 3 is very strong. That would indicate that there is a connection site or that would indicate that BERT has learned something about the connection sites. If we find this repeatedly so if we look at many many proteins and whenever we know that there is a contact between two things and then we observe that the corresponding attention is very high then we can be pretty sure that BERT has learned something about contact between amino acids. The same goes for binding sites so if 4 here is a binding site and then all the connections all the attention that the higher layer gets from 4 so all the information routed away from 4 is very strong that means all these other tokens are paying special attention to the token number 4 to this amino acid and if we find that there's a big correlation with this being a binding site then we can reasonably conclude that BERT has learned something about binding sites. So we're going to do a correlative analysis for proteins where we know the binding sites where we know the contacts. We can analyze them we can run simulations therefore we can know them. So we're going to look at this quantity right here which is simply a normalized quantity. So we're going to look at the attention in a given attention head so as you know BERT has many layers with many attention heads and we're going to look at whether or not this property is active and just normalize it by the total attention in that head so that we get some kind of a percentage number. That's the first task we're basically going to look at how does the attention correlate with these properties and the second task we're going to do is this probing task. So a probing task is like a linear probe in like a classifier so what we're going to do is we're going to take a layer right here and even though it's an intermediate layer we're simply going to run it through a linear classifier and then decide is this a binding site for example or not. Is a given amino acid a binding site or not? Is a given pair a contact or not? So this is kind of a linear probe but this sort of takes a backseat in this paper the analysis is really on the attention heads and what the attention heads learn. And that's already it they don't they take a pre-trained BERT model so there are these BERT models that are already trained on these protein databases and first they look simply can we find attention heads that correlate with a given amino acid. So here you see the attention to the amino acid this is proline I believe and this is phenylalanine is that the same in English? Yes phenylalanine and proline right here. So you can see that the the plots here are there's almost no attention pretty much throughout the network that pays special attention to the amino acid proline except this head right here seems to have if you look at the scale over like a 70% of attention always goes to proline in this particular head so this is layer 1 head number 11 focuses 78% of its attention on proline. Now this is not that special if you think about it because in language models as well in natural language models you might want to think that you have some mechanism in your neural network that's especially specialized on like a very particular word in the language because that might just be a often occurring very particular word for example in English maybe the is very important or the word what these these are like very indicative very often occurring words so it is reasonable to expect to find a an attention head that pays a lot of attention to these things especially here where our vocabulary size is 20 instead of like 30,000 in natural language. And the same goes for this phenylalanine where you can see that in the layer in the last layer and in the first layer you have attention and also in the proline you have in the last layer so why does this make sense because what we would expect from like single tokens these are not interactions yet these are not biological functions yet so we know that in the lower layers of a neural network we have these kind of basic features basic feature extractors and here these basic feature extractors appear to be simply paying attention to one specific token in the vocabulary a lot okay so they kind of these heads sort of specialize for single for single amino acids and the same in the last layer so in the very last layer the the task of the very last layer is to prepare for the classification tasks so if you remember the BERT model you have layer layer layer layer and at the end you'll have to predict which ones are masked down here so at the end you'll have to predict single amino acids again so if there's a proline masked here you'll have to predict the proline so it also makes sense that the last layers would very much specialize to single tokens so this does make sense now our question is going to be do we find the biological function where where would you expect them we would expect the let's say the tertiary sorry the secondary structures which are sort of one level higher than the primary structure we would expect to find them maybe here and then we would expect to find the tertiary structures maybe somewhere here okay because these are most highest level and then it goes it goes back again or maybe it's like we find the tertiary structures rather here and here again and then in the middle we'll find the the most high-level the tertiary structure sorry yeah blue secondary this drawing is getting too too too weird but there there could be multiple scenarios but that could fit here but until now it sort of makes sense so they do it as an additional investigation where as I told you sometimes you can substitute an amino acid and nothing really happens right and in fact that this probably happens in you right now you probably might have some mutation that changed some amino acid and you don't even realize because it's just it's fine no notice so the biologists can build these matrices of how much you can substitute proteins with each other so here you see this blossom 62 substitution scores which are very I guess very high if you can substitute two protein two amino acids with each other and the effect is negligible and it's very low if it's the other way around now this is interesting so far but you compare this to this matrix right here this is the attention similarity so what we'll do is for each two amino acids we take those two attention things those two attention matrices and we'll calculate the correlation between the attention matrices and our hypothesis is that the more correlated the attention patterns are between the two amino acids the more likely we are to substitute them because as a direct result of our language model our language model is it's reconstructing right these things so our language model is going to treat if in natural language is like a synonym right is our language model is going to treat synonyms very similar to each other because they're synonyms they can be exchanged so a good language model should learn that they are almost the same and therefore the attention pattern is going to be almost the same so a high correlation we hypothesize is a means that the function of the amino acid is similar and therefore we can substitute it easily so this here is the matrix of the correlations between each two attention patterns of these amino acid and if you compare the two right here they look extremely similar just have a have a look for a little while and you'll see that the patterns they do not match perfectly but they are very very similar the dark spots are in the same places the light spots are in the same places so this already makes a good case that the language model here has learned something about biology now what we want to do is we want to investigate higher order higher order functions so here we are interested in these contact maps right so how likely is it that two amino acids are in contact and we'll look at it through the lens of attention as we did before so here you'll see percentage of each each head of each head's attention that is aligned with contact maps averaged over data set suggesting that had 12 for is uniquely specialized for contact prediction so look at this this this head here is just spiking so remember before we said our analysis is whenever whenever we're basically measuring the correlation of two things being in contact because we know it from our simulator or from our data set the correlation of that with an attention connection being particularly strong and will we find it in this attention head right here so this layer 12 head number four will always peek out whenever two things are in contact now you can see that it's it's not like always it's like 25% of its attention but significantly more than anything else right here in fact if you group the things by this attention you can build the following plot so you can see right here probability two amino acids are in contact as a function of attention between the amino acids in head 12 for showing attention approximates perfectly calibrated estimator which would be the green line so here we simply for each for each pairs to amino acids for each pair of amino acids we plot we make a histogram right here of what they're sorry not a histogram we plot the probability if they have the attention weight point nine we plot how likely is it that they are in contact so this is this if we just look at the data and we simply take this attention weight as a measure as a predictor of being in contact we get the blue curve and the green curve would be if we could perfectly predict from this attention head what the probability of contact would be and you can see that the fit is fairly good you can't predict with super high accuracy but the fit is fairly good and you can see that general trend that as the attention in this head rises the probability of the two amino acids being in contact with each other also rises so we can sort of confidently say that BERT has learned something about a higher level funk a higher level biological structure just from the language modeling objective how can we interpret this this must somehow mean that it is it is possible it is vital to it is vital for reconstructing the sequence from its surroundings so if we delete this right here if if this if these two are in contact in the 3d structure that makes probably probably means that this thing right here is a very good predictor of what was here right if we mask this out and we're asked to reconstruct which amino acid was there then it probably helps to look at its neighbors right it probably always helps to look at one's neighbors especially also in natural language but if these two are in contact they they have very special they have very special connection to each other it's very you can basically read out from this one which one this was this is sort of like if you have a sentence and you say does I don't know I can't come up with with one right now but if it's like da da da da da and then there is a name like mark and then da da da da da da and then there is him right and you would expect if I drop out okay let's do it the other way around if I drop out him then from the text right here you can probably determine that it is some sort of pronoun but then you go back and you see ah it's mark okay so it's not like it's not like it or some or or or she it's probably he or him this is sort of the analogous structure right here in biology the second thing we're looking at is these these binding sites now these are single properties of of of different amino acids and we're simply looking at all the incoming or sorry all the F all the other tokens that focuses their attention why is this important because these binding sites are central to the structure of the or to the function of the protein right if this here is a binding site then that's a very central important point of the protein so a lot of these other things are going to be determined by what this binding site is this binding site needs to have a very spurted particular function and therefore probably needs to be a very particular amino acid and the other things here are sort of supporting this binding site because they form the 3d structure around it and so on so you would expect a lot of attention to be put on this binding site and what do we find the we find that it's a bit more murky than before so you can see that the attention is kind of spread out percentage of each head's attention that focuses on binding sites especially in the deeper layers binding sites are targeted at much higher frequency than would occur by chance head 7 1 has the highest percentage with 34% so also here you can see that it is spread out but this is because multiple heads are now focusing on these binding sites because probably binding sites come in different variations so you'll have lots of heads specializing on attending to binding sites and they say it is much higher frequency than would occur by chance and you can see here this head is the highest with 34% of its attention focused on binding sites you can also see the general trend of the attention being rather in the later layers which we would expect from a tertiary structure now yeah it would be interesting here here you also see that actually most of the things are in the in the last layer which points to points to rather maybe lower lower level information because we reasoned before about the last layer or I was just wrong but also in a general trend you can see that the attention is rather shifted towards the later layers because this is sort of a higher order function okay if you look at the same calibration experiment you can see that the picture is not as clear there is the general trend at the beginning but then it sort of flattens out so you can sort of differentiate the very probably not a binding site from the somewhat probably a binding site but it's not a perfectly calibrated classifier and that might just be because there are many things specializing in different types of binding sites so you can't just go to this one head so this is just for this one head you can't just go to that one and expect that to classify all the binding sites because you might want to be you might want to combine all of the high ranking ones here to form a classifier the last experiment they do is these linear probes which where they just go and they just build classifiers from different parts of the network and you can see right here that what is predicted and how well they work so each bar here is going to be the difference of performance so this is differential performance of diagnostic classifier by layer sorted by task order in figure 8 each plot shows the change in performance between the given layer and the previous layer okay so a bar up shows it's performing better than the previous layer bar down shows it's performing worse than the previous layer so you see right here that the these are the secondary structures right here and you can see that there is a lot of performance in the earlier layers right here and sort of not that high performance in the later layers whereas for the tertiary structures the binding site and the contact you can see that there is a bit of performance on in places but it sort of tends to be more towards the middle certainly more at towards the middle of the end of the network than the the secondary structures which sort of makes sense with our hypothesis you can also see this here where they show the percent of attention focused as a function of layer and the red is the center of mass and you can see that as the the secondary structures this their center of mass is at a lower layer in general than the tertiary functions all of this is not perfect of course but it's still an open question I guess whether or not it's not perfect because we haven't built a strong enough language model yet do I want to say GPT-4 is now for biology and not for language or is it because there is really you need you really can't very well predict the these things just from a language model I mean you should technically all the information is there but maybe the language model objective as such isn't able to capture that information so yeah this this was the paper it's pretty simple they have the in the appendix they have a lot of a lot of these additional experiments or full experiments I believe for all the amino acids and so on and I invite you to check that out in general I like this kind of work because it's very applied it's and it can you know tell us something about the nature of both these language models and the biological things that we that we care about in biology okay I'm just talking crap right now thanks for being here I hope you enjoyed it and bye bye
 Hi there! Today we'll look at Bertology meets Biology interpreting attention in
 protein language models by Jesse Vig, Ali Madani, Lav R Varshini, Kaiming Xiong,
 Richard Sokker and Nazneen Fatima Rajani. This paper is a investigative paper into
 models that are trained on biological data, specifically into BERT models.
 Actually into one specific BERT model that is trained on protein sequences.
 Now it is trained to simply perform language modeling on these protein
 sequences, but out of this language model you can then inspect this BERT model and
 read important biological data of these proteins, higher order data from the
 attention heads of the BERT model, which is pretty interesting. Basically means
 that the information of these higher order functions is at some point encoded
 in the structure of the language of the protein sequence. So we're going to go
 through what this means and how this comes about and what they did in order
 to investigate. I think this is a pretty cool investigative work and probably
 very promising for future research. Yeah, as always if you like content
 like this, consider sharing it out and leaving a like. Also tell me what you
 think in the comments. So biology. Really quick for people who maybe have never
 heard this. In your every cell you have this thing called DNA, which basically is
 an encoding of all of your biological functions. Now usually biological
 functions are realized through proteins. So DNA is basically a building plan for
 all of your proteins. This happens in the following two steps. First there is this
 transcription step where RNA is built. This is basically a copy of your DNA, but
 it's only single strand as you can see right here. And then there is a
 translation step that finally translates the RNA into the protein. What will end
 up is just a sequence of these beads right here. Now these beads are what are
 called amino acids. So a protein is simply a chain of these amino acids.
 There are 20 different amino acids and the order of these amino acids in the
 chain makes the function of the protein. Now specifically we know about these
 proteins that it seems to be very important how their three-dimensional
 shape is. So a lot of these different amino acids have different chemical
 properties. Some are sort of I think negatively charged, some are neutral, some
 are acids and so on. So they have very different chemical properties. So once
 you build this protein and you kind of release it into the cell it will curl up
 into a three-dimensional structure. So this one might be doing
 something like this and sort of form a circle or something like this.
 Just because these proteins here they kind of attract each other maybe
 electrically and thus the protein forms a circle and the function of the protein
 is very much related to its shape. So if it is a circle it could maybe trap
 something else in here. So you really have to think of these things like kind
 of tools. There are proteins that cut other proteins and they are really
 shaped sort of like a scissor that exactly fits these other proteins such
 that you can effectively cut them. So sometimes you can substitute an
 amino acid for a different amino acid like this here. If it doesn't change the
 shape very often you're fine. The protein function isn't changed.
 But if you change a different amino acid that is sort of vital to the shape and
 the shape changes then your protein very often loses function. So mutations in
 DNA sometimes lead to mutations in protein. Not always because there is some
 redundancy in this translation step from RNA. But if they do lead to a
 different amino acid it doesn't actually mean that the function changes. So there
 is sort of value in analyzing the sequence of the structure of proteins
 rather than the structure of DNA. Of course it's also important to analyze
 the structure of DNA but it is equally important to analyze the
 structure of proteins because not all the information is in the
 sequence. Not all the obvious information is in the sequence. So what does this
 paper do? This paper goes and takes a model that has been trained on protein
 data. So if you look at this protein it is simply a sequence of amino acids and
 these amino acids they all have names. I think I have a table somewhere here.
 Yes so these are the different amino acids that exist and you can see a
 protein is simply a sequence of these names. So usually they're abbreviated by
 like a three-letter abbreviation or just a one-letter abbreviation. So a protein
 might be AVMMVAG and so on. And this is just a string of text. So what I
 can do is I can train a language model on this. A language model is simply a
 model that takes a piece of text and tells you what's the next piece of text.
 So what's the next letter, what's the next word, in this case what's the next
 amino acid. And we can use tools from NLP for that. Specifically we can train a
 BERT model. Now BERT works a bit differently than a standard language
 model. BERT does what is called masked language modeling. So you would take this
 string, you would feed it into a BERT model right here. And I've made an entire
 video on BERT if you want to check that out. And what you'll do by inputting
 that you'll mask out some of the tokens. So you'll maybe mask out this one, mask
 out this one, and then you ask the model to reconstruct those. We say that here is
 an M and here is an A without seeing them. So the model somehow has to learn
 from the surrounding amino acids what this amino acid could be. So it has
 to reconstruct this sequence. So the hope here is, in natural language, is that
 BERT somehow learns something about language itself. By being able to
 reconstruct these things it has learned something about language, about which
 words appear together and when. It might even learn very long distance
 relationships between words just because it has to predict those. And the idea
 carries over to biology. So we might hope that a BERT trained on an amino
 acid sequence will learn something about the language of proteins,
 about the amino acid sequence. And our goal here is to ask, can we somehow
 infer the 3D shape of a protein, which is the important part, from its sequence
 right here? So given its sequence, can we infer the 3D shape? Now as I understand
 it, usually this has to be done in like a simulation. So you would build
 this in a simulator and then you do like some sort of a molecule simulation to
 see how this ends up in a 3D shape. You could train a model to just predict the
 3D shape, but in this case we're just interested in what does the BERT model
 learn about the 3D shape while only ever having been trained on predicting the
 sequence of amino acids. So it's never been trained to look at the
 3D shape. And that's our goal here. So specifically we'll look at two different
 things. So here you can see examples of proteins and their high-level structures.
 So in these proteins what you call the primary structure is this sequence of
 amino acid. This is simply which amino acids are in which order. There is a
 thing called the secondary structures and we often observe that spans of these
 amino acids like substrings form these what are called these helixes as you can
 see here or these sheets. I don't know how they're strands in English. We call
 them sheets or I think these are the alpha helixes and these are the beta
 sheets and there is also a turn. I think this here might be a turn. So there are
 these kind of secondary structures and then the tertiary structure is how these
 this is still one protein. This is one unbroken chain of amino acid.
 And you can see this here kind of forms this double ring which would be its
 tertiary structure. Very important for predicting the tertiary structure is to
 predict when two amino acids are close to each other. So if we have a chain
 right here and the chain as we saw before kind of turns and bends on itself
 then these two amino acids here are very close in close contact. And to predict
 which amino acids are in close contact to each other helps you determine the
 the tertiary structure. So that's a consequence of it. So we wonder does
 BERT know intrinsically which of these amino acids are going to end up being in
 contact with each other without ever having been trained to do it? The second
 thing we're interested in are binding sites. So here what you might not be able
 to see but we made this example before where this sort of forms a loop and then
 I say can trap something here right like another molecule. And this is what
 we would call a binding site. A binding site is a one amino acid that maybe
 through the structure of the surrounding amino acid as well but also through its
 properties and how it is exposed in 3d shape acts as sort of a receptor for
 other molecules. It binds to other things. So think of your hemoglobin
 that traps the oxygen in your blood or something like this. It is where a
 chemical reaction or a reaction with something else will happen. That's a
 binding site. And we are interested does BERT, the BERT that is only trained on a
 language modeling objective, know which ones are the binding sites? Because you
 know that would be very interesting and not something BERT was trained on. By the
 way I particularly liked Richard Sockers tweet on this. I think he tweeted
 out, oh BERT trained only on language model can predict binding sites and
 biological properties and formulated it like it was you know like GPT-3 was
 formulated like if we train on Wikipedia our model can do math. I thought it was
 kind of a satire headline. If we train on Wikipedia our model can predict biology
 and also it can tie your shoes and cook your dinner. So it's trained on
 language modeling on biological data and now that makes sense. So they're going to
 look at two different things or actually more than two different things but
 they formulate this in an abstract way right here. So what they'll look at is
 the so-called properties. A property F and this property F can be for example
 that a amino acid is a binding site. The property F can also be that two amino
 acids are in contact with each other. So F always takes I and J. If in the case
 for example where this is the contact property then it simply is the indicator
 function for when I and J are in contact. And if it is a just a binding site then
 I think we are looking at J. At the token level property we define to be an
 indicator that returns one if the property is present in token J. So whenever
 J is a binding site then that holds. So what we're looking at are these
 attention heads in BERT. If you don't know BERT has an attention mechanism
 which basically means from layer to layer each token can attend to all other
 tokens. So here the amino acid sequence I've drawn it twice and the next layer
 representation of this amino acid will be able to gather information from all
 of the other amino acid through an attention mechanism through a dynamic
 routing algorithm. I've made a video on attention is all you need if you want to
 find out more how this works. Now what we're interested in is the strength of
 these connections. So the hypothesis is if molecule 1 and 3 are
 contact sites then maybe we will find a layer where this connection between 1
 and 3 is very strong. That would indicate that there is a connection
 site or that would indicate that BERT has learned something about the
 connection sites. If we find this repeatedly so if we look at many
 many proteins and whenever we know that there is a contact between two things
 and then we observe that the corresponding attention is very high
 then we can be pretty sure that BERT has learned something about contact between
 amino acids. The same goes for binding sites so if 4 here is a binding
 site and then all the connections all the attention that the higher layer
 gets from 4 so all the information routed away from 4 is very strong that
 means all these other tokens are paying special attention to the token number 4
 to this amino acid and if we find that there's a big correlation with this
 being a binding site then we can reasonably conclude that BERT has
 learned something about binding sites. So we're going to do a correlative
 analysis for proteins where we know the binding sites where we know the
 contacts. We can analyze them we can run simulations therefore we can know
 them. So we're going to look at this quantity right here which is simply a
 normalized quantity. So we're going to look at the attention in a given
 attention head so as you know BERT has many layers with many attention heads
 and we're going to look at whether or not this property is active and just
 normalize it by the total attention in that head so that we get some kind of a
 percentage number. That's the first task we're basically going to look at how
 does the attention correlate with these properties and the second task we're
 going to do is this probing task. So a probing task is like a linear probe in
 like a classifier so what we're going to do is we're going to take a layer right
 here and even though it's an intermediate layer we're simply going
 to run it through a linear classifier and then decide is this a binding site
 for example or not. Is a given amino acid a binding site or not? Is a
 given pair a contact or not? So this is kind of a linear probe but this sort of
 takes a backseat in this paper the analysis is really on the attention
 heads and what the attention heads learn. And that's already it they don't
 they take a pre-trained BERT model so there are these BERT models that are
 already trained on these protein databases and first they look simply can
 we find attention heads that correlate with a given amino acid. So here you see
 the attention to the amino acid this is proline I believe and this is phenylalanine
 is that the same in English? Yes phenylalanine and proline right here.
 So you can see that the the plots here are there's almost no attention pretty
 much throughout the network that pays special attention to the amino acid
 proline except this head right here seems to have if you look at the scale
 over like a 70% of attention always goes to proline in this particular head so
 this is layer 1 head number 11 focuses 78% of its attention on proline. Now this
 is not that special if you think about it because in language models as well in
 natural language models you might want to think that you have some mechanism in
 your neural network that's especially specialized on like a very particular
 word in the language because that might just be a often occurring very
 particular word for example in English maybe the is very important or the
 word what these these are like very indicative very often occurring words so
 it is reasonable to expect to find a an attention head that pays a lot of
 attention to these things especially here where our vocabulary size is 20
 instead of like 30,000 in natural language. And the same goes for this
 phenylalanine where you can see that in the layer in the last layer and in the
 first layer you have attention and also in the proline you have in the last
 layer so why does this make sense because what we would expect from like
 single tokens these are not interactions yet these are not biological functions
 yet so we know that in the lower layers of a neural network we have these kind
 of basic features basic feature extractors and here these basic feature
 extractors appear to be simply paying attention to one specific token in the
 vocabulary a lot okay so they kind of these heads sort of specialize for
 single for single amino acids and the same in the last layer so in the very
 last layer the the task of the very last layer is to prepare for the
 classification tasks so if you remember the BERT model you have layer layer
 layer layer and at the end you'll have to predict which ones are masked down
 here so at the end you'll have to predict single amino acids again so if
 there's a proline masked here you'll have to predict the proline so it also
 makes sense that the last layers would very much specialize to single tokens so
 this does make sense now our question is going to be do we find the biological
 function where where would you expect them we would expect the let's say the
 tertiary sorry the secondary structures which are sort of one level higher than
 the primary structure we would expect to find them maybe here and then we would
 expect to find the tertiary structures maybe somewhere here okay because these
 are most highest level and then it goes it goes back again or maybe it's like
 we find the tertiary structures rather here and here again and then in the
 middle we'll find the the most high-level the tertiary structure sorry
 yeah blue secondary this drawing is getting too too too weird but there
 there could be multiple scenarios but that could fit here but until now it
 sort of makes sense so they do it as an additional investigation where as I
 told you sometimes you can substitute an amino acid and nothing really happens
 right and in fact that this probably happens in you right now you probably
 might have some mutation that changed some amino acid and you don't even
 realize because it's just it's fine no notice so the biologists can build
 these matrices of how much you can substitute proteins with each other so
 here you see this blossom 62 substitution scores which are very I
 guess very high if you can substitute two protein two amino acids with each
 other and the effect is negligible and it's very low if it's the other way
 around now this is interesting so far but you compare this to this matrix
 right here this is the attention similarity so what we'll do is for each
 two amino acids we take those two attention things those two attention
 matrices and we'll calculate the correlation between the attention
 matrices and our hypothesis is that the more correlated the attention patterns
 are between the two amino acids the more likely we are to substitute them
 because as a direct result of our language model our language model is
 it's reconstructing right these things so our language model is going to treat
 if in natural language is like a synonym right is our language model is going to
 treat synonyms very similar to each other because they're synonyms they can
 be exchanged so a good language model should learn that they are almost the
 same and therefore the attention pattern is going to be almost the same so a high
 correlation we hypothesize is a means that the function of the amino acid is
 similar and therefore we can substitute it easily so this here is the matrix of
 the correlations between each two attention patterns of these amino acid
 and if you compare the two right here they look extremely similar just have a
 have a look for a little while and you'll see that the patterns they do not
 match perfectly but they are very very similar the dark spots are in the same
 places the light spots are in the same places so this already makes a good case
 that the language model here has learned something about biology now what we want
 to do is we want to investigate higher order higher order functions so here we
 are interested in these contact maps right so how likely is it that two amino
 acids are in contact and we'll look at it through the lens of attention as we
 did before so here you'll see percentage of each each head of each head's
 attention that is aligned with contact maps averaged over data set suggesting
 that had 12 for is uniquely specialized for contact prediction so look at this
 this this head here is just spiking so remember before we said our analysis is
 whenever whenever we're basically measuring the correlation of two things
 being in contact because we know it from our simulator or from our data set the
 correlation of that with an attention connection being particularly strong and
 will we find it in this attention head right here so this layer 12 head number
 four will always peek out whenever two things are in contact now you can see
 that it's it's not like always it's like 25% of its attention but significantly
 more than anything else right here in fact if you group the things by this
 attention you can build the following plot so you can see right here
 probability two amino acids are in contact as a function of attention
 between the amino acids in head 12 for showing attention approximates perfectly
 calibrated estimator which would be the green line so here we simply for each
 for each pairs to amino acids for each pair of amino acids we plot we make a
 histogram right here of what they're sorry not a histogram we plot the
 probability if they have the attention weight point nine we plot how likely is
 it that they are in contact so this is this if we just look at the data and we
 simply take this attention weight as a measure as a predictor of being in
 contact we get the blue curve and the green curve would be if we could
 perfectly predict from this attention head what the probability of contact
 would be and you can see that the fit is fairly good you can't predict with
 super high accuracy but the fit is fairly good and you can see that general
 trend that as the attention in this head rises the probability of the two amino
 acids being in contact with each other also rises so we can sort of confidently
 say that BERT has learned something about a higher level funk a higher level
 biological structure just from the language modeling objective how can we
 interpret this this must somehow mean that it is it is possible it is vital to
 it is vital for reconstructing the sequence from its surroundings so if we
 delete this right here if if this if these two are in contact in the 3d
 structure that makes probably probably means that this thing right here is a
 very good predictor of what was here right if we mask this out and we're asked
 to reconstruct which amino acid was there then it probably helps to look at
 its neighbors right it probably always helps to look at one's neighbors
 especially also in natural language but if these two are in contact they they
 have very special they have very special connection to each other it's very you
 can basically read out from this one which one this was this is sort of like
 if you have a sentence and you say does I don't know I can't come up with with
 one right now but if it's like da da da da da and then there is a name like mark
 and then da da da da da da and then there is him right and you would expect
 if I drop out okay let's do it the other way around if I drop out him then from
 the text right here you can probably determine that it is some sort of pronoun
 but then you go back and you see ah it's mark okay so it's not like it's
 not like it or some or or or she it's probably he or him this is sort of the
 analogous structure right here in biology the second thing we're looking
 at is these these binding sites now these are single properties of of of
 different amino acids and we're simply looking at all the incoming or sorry all
 the F all the other tokens that focuses their attention why is this important
 because these binding sites are central to the structure of the or to the
 function of the protein right if this here is a binding site then that's a
 very central important point of the protein so a lot of these other things
 are going to be determined by what this binding site is this binding site needs
 to have a very spurted particular function and therefore probably needs
 to be a very particular amino acid and the other things here are sort of
 supporting this binding site because they form the 3d structure around it
 and so on so you would expect a lot of attention to be put on this binding site
 and what do we find the we find that it's a bit more murky than before so you
 can see that the attention is kind of spread out percentage of each head's
 attention that focuses on binding sites especially in the deeper layers binding
 sites are targeted at much higher frequency than would occur by chance
 head 7 1 has the highest percentage with 34% so also here you can see that it is
 spread out but this is because multiple heads are now focusing on these binding
 sites because probably binding sites come in different variations so you'll
 have lots of heads specializing on attending to binding sites and they say
 it is much higher frequency than would occur by chance and you can see here
 this head is the highest with 34% of its attention focused on binding sites you
 can also see the general trend of the attention being rather in the later
 layers which we would expect from a tertiary structure now yeah it would be
 interesting here here you also see that actually most of the things are in the
 in the last layer which points to points to rather maybe lower lower level
 information because we reasoned before about the last layer or I was just wrong
 but also in a general trend you can see that the attention is rather shifted
 towards the later layers because this is sort of a higher order function okay if
 you look at the same calibration experiment you can see that the picture
 is not as clear there is the general trend at the beginning but then it sort
 of flattens out so you can sort of differentiate the very probably not a
 binding site from the somewhat probably a binding site but it's not a perfectly
 calibrated classifier and that might just be because there are many things
 specializing in different types of binding sites so you can't just go to
 this one head so this is just for this one head you can't just go to that one
 and expect that to classify all the binding sites because you might want to
 be you might want to combine all of the high ranking ones here to form a
 classifier the last experiment they do is these linear probes which where they
 just go and they just build classifiers from different parts of the network and
 you can see right here that what is predicted and how well they work so each
 bar here is going to be the difference of performance so this is differential
 performance of diagnostic classifier by layer sorted by task order in figure 8
 each plot shows the change in performance between the given layer and
 the previous layer okay so a bar up shows it's performing better than the
 previous layer bar down shows it's performing worse than the previous layer
 so you see right here that the these are the secondary structures right here and
 you can see that there is a lot of performance in the earlier layers right
 here and sort of not that high performance in the later layers whereas
 for the tertiary structures the binding site and the contact you can see that
 there is a bit of performance on in places but it sort of tends to be more
 towards the middle certainly more at towards the middle of the end of the
 network than the the secondary structures which sort of makes sense with
 our hypothesis you can also see this here where they show the percent of
 attention focused as a function of layer and the red is the center of mass and
 you can see that as the the secondary structures this their center of mass is
 at a lower layer in general than the tertiary functions all of this is not
 perfect of course but it's still an open question I guess whether or not it's not
 perfect because we haven't built a strong enough language model yet do I
 want to say GPT-4 is now for biology and not for language or is it because there
 is really you need you really can't very well predict the these things just from
 a language model I mean you should technically all the information is there
 but maybe the language model objective as such isn't able to capture that
 information so yeah this this was the paper it's pretty simple they have the
 in the appendix they have a lot of a lot of these additional experiments or full
 experiments I believe for all the amino acids and so on and I invite you to
 check that out in general I like this kind of work because it's very applied
 it's and it can you know tell us something about the nature of both these
 language models and the biological things that we that we care about in
 biology okay I'm just talking crap right now thanks for being here I hope you
 enjoyed it and bye bye