Scientist Stories: Kevin Wu, Learning the Grammar of T-cell receptors

Uh, so my name is Kevin, and I guess I'll be kicking off the rounds of orals for MLCB. I'm glad to be back in person, and today I'll be talking about a work that we've been doing called TCRBERT, which is sort of taking a lot of the same ethos and design principles that we already discussed for ESM but applying it to a special class of amino acids called T cell receptors. Now, before we get too far into this rabbit hole, I thought it'd be good to cover, you know, what is a T cell receptor, why do we care about it, and perhaps why do we even need some computational modeling to, to think about this system. So within all of our immune systems, there are these cells called T cells, and one of their key roles is to sort of mediate the immune response. They do this by recognizing what is called an antigen. These antigens are these short peptide barcodes that are presented to both healthy and diseased cells, sort of like a name tag or a grocery tag at the supermarket that says, I'm a healthy cell, I'm a diseased cell, or something along those lines. Now, when these T cells come along and recognize an antigen, they sort of may decide this is a bad cell, this is a good cell, and then recruit the immune system accordingly. Now, this binding recognition effect sort of mediated by this T cell receptor, this dimeric peptide chain. So we have the TRA chain, the TRB chain, which are more or less separate, and these two chains come to come along and sort of look at this antigen and make that decision. Now, since this T cell receptor and this T cell response in general is so important to how the immune system works, it's very important for us as computational scientists to be able to say, can we look at these systems computationally and predict how they might behave? And this is, you know, of course a huge challenge, but it's also very, very difficult to do. And there's a couple of different reasons for that. Perhaps the most pressing one is that these T cell receptor amino acid chains are not actually like most proteins you would find in a human body. So for example, if you were to take something like hemoglobin, which is sort of circulating in all our blood, you would find that, say, my copy of hemoglobin protein is probably almost exactly the same as anybody else's hemoglobin protein. And that's sort of this conservation is sort of the backbone of models like ESM, and as a backbone of many, many protein sort of analysis tools. The difference though is with TCRs, these are actually generated independently within your own body by this sort of very complex combinatorial process. And what that means is that every single one of us has orders of millions of different T cell receptor amino acid chains that are all valid, they all work,, but they're so incredibly diverse even within an individual. And that's not even to cover the fact that if you were to take my TCRs, for example, and anybody else's TCRs, there'd be very little, if any, real overlap between the exact sequences that are being produced that are functioning. So that's one challenge. The other one is that of what is called cross-reactivity, which is that there's no single clean mapping that says one TCR maps to this antigen, binds to that, recognizes that, and that's it. What we see in reality is that one antigen can be recognized by multiple TCRs, and one T cell receptor can then recognize multiple antigens. So this sort of very complex many-to-many mapping, that's much more difficult to understand from a sort of modeling perspective. So when we started looking at this problem, we sort of looked at what was being done in the literature up to this point. And what we saw was that prior to this work, a lot of the— methods tackling this problem of predicting the behavior of these T cell receptors were essentially these supervised methods, which I'm sure we're all aware basically just say we have this set of T cell receptor sequences. For example, we have some method to measure their binding affinities for specific antigens. We're going to train some sort of classifier on top of that. And this works reasonably well, but the critical aspect here is that this doesn't really capture a key aspect of the data, which is that we can actually profile a lot of T cell receptor sequences, but we don't know what they bind to. In other words, there's a huge pool of data that says we have TCRs, but we don't have labels for them. So that sort of significantly limits how much we can push these purely supervised methods in this area. So to sort of try and better tackle this problem, we sort of turn to the language modeling paradigm, which we have already discussed in depth this morning. Thank you, Alex. And from there, we aren't— make some additional tweaks though, because we don't just want to look at this unsupervised or self-supervised task of trying to learn the grammar of these sequences. We also want to figure out, we have this sort of set of subset of TCRs with known antigen behavior. Can we somehow integrate that known labels with the self-supervised aspect of these language models. So the solution we came up with is a sort of two-step sequential pre-training process. The first is this sort of masked amino acid prediction model that was common to models like ESM as well. And after we've applied this masked amino acid prediction model to a series of T-cell receptor sequences, we do a secondary pre-training task, which is saying, given a small set of sequences for which we have reasonably well-studied antigens binding affinities for. We want to predict, given this T-cell receptor sequence, which one of this panel of, say, 45 different antigens will it probably bind to. So it's sort of multi-class classification step in addition to this sort of pre-trained masked language modeling task that has been discussed already. So after we've done this combination of two pre-training tasks, you want to see, you know, does this actually make any difference? Does this do better than existing models. So the first thing to ask is, well, does this do better than supervised modeling? I mentioned earlier that these existing models do fine, but can we do better? The way you run this evaluation is we take a set of antigens, we take each antigen, we remove its data from training, we retrain this model, and then we take that resulting model and evaluate on that held-out antigen to make sure we're not sort of bleeding training and test set. We run this evaluation for many, many antigens, and what we see is that on the left-hand side, on the y-axis, we have TCRBERT's performance on the y-axis. And this is done by simply taking the TCRBERT model, freezing the model, taking embeddings out, and training an SVM on top of it. On the x-axis, we're comparing against two different supervised baselines. So we have a convolutional baseline on the left and a sort of K-mer-based baseline on the right. And in both cases, the performance of our TCRBERT model exceeds that of the supervised models in every single antigen that we evaluate in this sort of cross-validation fashion. Now, this in itself might not be very surprising given the fact that it's been pretty well shown at this point that pre-training in general tends to improve upon purely supervised learning. Now, the better question to ask then might be to say, Can we— does this model do better than other protein language models, right? Because after all, I'm proposing that we're going to train a different model for this very specific class of molecules. Is this more effective? To do this, we compare our embeddings against those created by two different pre-trained protein language models that have seen firstly much, much more sequence and are much, much larger models. So this is specifically the TAPE model, and also the ESM model that was discussed this morning. And we do the same procedure. We take these models, we generate embeddings that train a simple SVM on top of this to predict for each of these health-out antigens binding versus non-binding. And what we see is that for our model TCLBERT, we still outperform these general-purpose protein language models. And, you know, this is a good benchmark result, but I think the real takeaway here is that even in the era of sort of extremely large billion-parameter networks that are sort of open source and seems to be the backbone of everything, there is still space to think about, is there a targeted problem that I can do, a targeted dataset I can use to build a model that is specific to an area and may even outperform these large Swiss Army knife type models? So that was pretty encouraging to see. So beyond the benchmarking though, we also wanted to ask, can we sort of like apply this TCR birth model to various important problems? One of the most critical problems that is challenging in this space is this problem of cross-patient generalization. And that's this idea that since each of our bodies is independently generating these T cell receptor sequences, it's very hard to train a classifier that says we can look at one person's data and then generalize to a different person's data since that's a sort of different underlying problem. Process. We want to evaluate this, so we sort of evaluated TCRBERT's ability to do this cross-patient generalization as well as the ESM-based models, tape-based models, and various supervised baselines shown here in the AUPRC curves on the left-hand side. And what we see is that TCRBERT again gives us the best performance, but perhaps more importantly, we see that pre-trained models in general, shown in solid lines, are generally far outperforming these supervised models in dotted lines. And again, this might echo a lot of our intuition about how pre-training works as a paradigm. I think it's a powerful result nonetheless that shows that this pre-training method is powerful even under challenging generalization conditions. So that's also a very encouraging result to see. Beyond that, we also showed that on the right-hand side, if you have a very large set of data that's available. We can also fine-tune this model to an extent to also outperform existing models, but I won't dwell too much on that in the interest of time. So beyond these sort of problems that are asking, can we predict yes or no binding or non-binding for specific antigens, another challenging problem in the area of T cell receptor analysis is that of sort of unsupervised clustering. Since, as I mentioned before, you can often measure these T cell receptors but not necessarily have some sort of antigen-binding affinities. So we are often interested in this problem that says, can we identify groups of T cell receptor sequences that we think share similar antigen-binding patterns, even if we don't know exactly what that pattern might be? So essentially a clustering problem. Now, traditionally, this is done using various heuristic methods, such as saying, what is the edit distance between two sequences, for example? And for us, we might be able to approach this using our pre-trained TCRBERT model to generate embeddings and doing some sort of clustering on top of that. And we evaluate this sort of TCRBERT embedding clustering method versus the heuristic methods for two different datasets, one on the left and one on the right. And we essentially see that we see a better trade-off curve between accuracy on the left-hand side and percentage of meaningful cluster interactions on the x-axis. So that's also encouraging result and mirrors what we see in sort of improvements in supervised predictions. So all this is pretty, pretty encouraging so far, but one question to ask might be, how are we achieving this performance? Are we just memorizing training set, or are we truly learning something emergent about these T cell receptors, sort of similar to the structure predictions we saw out of ESM? Earlier this morning. So to do this, we ask a simple question, which is that if you recall in the very beginning of this slide deck, we have this antigen schematic and these T cell receptors that are touching it. The biological intuition here is that if these models are learning some— that in this system, the residues in the T cell receptors that are touching the antigen physically proximal should be the ones that are the most important for determining specificity. So what we did is for a system, we looked at the attentions, sort of the contributions of each amino acid to TCR-based predictions, and we sort of highlight that sort of interaction here. The green sort of region represents the major histochromatin complex that's basically just displaying the antigen. These sort of salmon-colored molecules is the antigen itself, the thing we want to recognize, and the pink and the yellow are the T cell receptor TCR sequences. Within the pink and the yellow, these solid-colored, sort of bolded molecules are the regions that have the highest attention, the most important things for predicting. And we see that the things that are most important are also the things that are most physically close. And that sort of mirrors our biological understanding of how this system should be behaving and provides some understanding of why this model might be working so well. So all that sort of is I hope, proof that this is a reasonable model, but I really want to close out with one example that shows we can sort of take these models and push them beyond just being able to get good benchmark results or good predictors or clustering or these classical tasks. Specifically, since these TCRs are so important for the immune system, one key question is, can we sort of engineer new TCR sequences to have sort of predetermined good specificity for targets that we have in mind? And the process we propose for doing this is really quite simple, as shown here on the sort of right-hand side. So let's say we start with a set of TCR sequences. We first want to sort of take these sequences and give it to our TCR BERT model to sort of rank them in order of which ones are potentially better binders than others. After we do this ranking, we sort of take the top 50% of the sequences, And then we give them back to TCRBERT, but this time instead of asking it to rank, we go back to the masked amino acid prediction results and ask it, can you sample for me a similar amino acid to what I have already that is sort of within the manifold of what you might expect naturally, but also different? And then we just re-rank again and we mutate, so on and so forth, this iterative cycle. And if we sort of run this cycle over iterations, we see that on the x-axis, as you run these iterations, on the y-axis for the center plot, these sort of predicted binding for our target of interest rises up over time. So in the very beginning, for the 0th iteration, we have basically no, not much binding, but after just 5 or 6, 7 iterations, we sort of converge upon a set that's predicted to have very good binding. And then on the right-hand side, I'm showing what is the sort of similarity to training set as we sort of run this design process. And this is an important plot because we don't want to just be memorizing training set sequences and spitting them back out. That's not very novel and not very interesting. But what we see is that we don't end up doing this sort of memorization modality. We are generating new sequences. One thing to point out here is that these T-cell receptor sequences are actually only about 15 amino acids. So you might say that 4 is a small number for edit distance, but 4 out of 15 is a relatively large number. Proportion. So that looks like it does work, but there's this key caveat here, which is that, okay, you're just optimizing some black box function, and yeah, you can optimize that, but we don't really know if that is working. We don't know if that gives us real hits, right? And perhaps the best way to evaluate this would be to put it in some sort of physical model and measure binding. But short of doing that, we can also do something computational that's much cheaper much faster to iterate. The basic premise here is that we take the sequences we have generated, we go to a database, and then we search it with BLAST and ask, are there experimentally measured sequences that are not in our training set but are hits to the things that we have generated? And it turns out that there are. So on the top row here, I'm showing a motif that sort of arises from our generated sequences. On the bottom is a motif that is derived from a set of matching sequences that we found through BLAST that is not seen in training, but has affinity experimentally determined to the antigen we were targeting. And while this isn't experimental validation by any means, this does suggest that we are sort of creating new sequences, at least new in terms of the model has never seen them. That have evidence that these should be experimentally binding if we're able to run that experiment. And with that, thank you for your attention. All our code is online, and thank you.

Yeah, we definitely have time for questions.

I don't know who was first.

Thank you. So thanks for the talk. You said that your attention predicts TCR-antigen contacts. I was wondering, to be clear, are you fine-tuning on the TCR-antigen pairs? Because you said that you earlier were just doing SVM classifications on the representations.

Or are you just feeding them in? That's a great catch. We are fine-tuning on TCR-antigen pairs for this specific analysis. So this basically takes from this part here where we do this fine-tuning process you described, and then we pull the attention out of that. Thank you.

Thank you.

And first, hey there, great talk. I just had a question. So you can predict the TCR, the antigen for a specific TCR, but have you tried to go to like TCR database for an entire patient's repertoire where now you have hundreds of thousands of TCRs and maybe 3 or 4 potential viral experiences and tried to deconvolute or predict the patient's history as sort of like a diagnostic?

Yeah, so we did try to do this. As you might imagine, fishing 1 or 2 hits out of 100,000 is incredibly difficult. So we sort of get better than random, but the AU PRC is like 1, 0.001, which is great compared to your baseline of 0.0001, but it wasn't really a result worth stowing crusade. Okay, the people who didn't ask questions first.

So the antibodies that attach to human proteins get filtered out during early childhood. So what happens if you present an antibody that would come from a human protein?

So one of these sort of interesting behaviors that emerge from this is that because we train on naturally observed human TCRs, you would expect that those TCRs have all been through this filtering process and should not sort of be self-reactive. It should not be like destroying human cells. So in the ideal case, your question of what happens, you present sort of human protein fragments. Hopefully they should not be reacted against. If all those things hold, hopefully.

Hey, thanks. That was really cool. I'm just curious about the general issue with all these large language models is, you know, the issues related to leakage, and you almost always end up in a sort of transductive setting, but particularly for this supervised task, the antigen recognition, have you investigated how the performance of your model varies as you sort of like you did in the last slide, sort of vary the similarity between the test example and members of the training set?

Yeah, absolutely. That's a really good question. So we did do analysis that says if we were to go through our evaluation set or test set and purposely remove anything within, I think, I believe it was 2 edits of anything we saw in training, do we sort of preserve performance? And it turns out that we get almost exactly the same performance numbers out of it, even with these sort of perhaps naive but very direct sequence similarity cutoffs.

Great.

Thanks. Last one.

Great talk. I was curious, when you trained your own specialized model and you showed the comparison to ESM, is there a reason you didn't try fine-tuning ESM?

At the time, my GPU did not fit ESM.

Ah, okay. No, it seems, yeah.

Okay, great. Let's thank Kevin again. So next up, we have Nathan Frey from Prescient, who's going to tell us about protein discovery. Great. Are you going off this laptop?

As well.

Okay. Wow, this guy is really small.

Where's your—

I'm not sure. I don't know why you was just following