Scientist Stories: Kevin Yang, Multimodal Deep Learning for Protein Engineering

Hi everyone, welcome to this M2D2Talk. Um, our speaker today is Kevin Young, who is a senior researcher at Microsoft Research where he works on problems at the intersection of machine learning and biology. He did his PhD at Caltech with Francis Arnold on applying machine learning to protein engineering. And before joining MSR, he was a machine learning scientist at Generate Biomedicine, where he used machine learning to optimize proteins. Thank you so much, Kevin, for coming here today and to share your work with us. Really excited to, yeah, to learn today. So without further ado, the floor is yours.

Thanks for the introduction. So today I'll be talking about multimodal deep learning for protein engineering, and if you have questions, you can drop them in the chat, which I can actually see, which is rare for these online seminars. Or at the end of every kind of chunk of the talk, I'll pause to take questions as well. So proteins are biology's actuators. Each of our cells contain between 1 and 3 billion proteins each. And they're responsible for things like structure, metabolism, and signaling. Some examples include things like IgG, immunoglobulin, which is responsible for our adaptive immunity. Hemoglobin carries oxygen in your blood. Insulin regulates your blood glucose levels. The coronavirus spike protein helps the coronavirus infect your cells. And luciferase makes fireflies glow. And all of this diversity and function arises from 20 fairly simple chemical building blocks called amino acids. And they're called amino acids because on the left here you can see they have an amine group, that's the NH2. On the right there's a carboxylic acid group. And so those parts are shared by all 20 amino acids. And they each have a side chain which differentiates them and gives them their different chemical properties. And so the side chains can be very simple and small, like glycine has just a hydrogen there, or alanine has a methyl group. They have different properties, some of them are nucleophilic, like serine has a COH, threonine attaches a methyl group to that, uh, cysteine has a silicon group. And you can see there's diversity in their chemistries, and there's aromatic ones, there's hydrophobic ones, there's acidic ones, there's, um, nitrogen-containing ones, bases. And all this functional diversity arises from different combinations of these building blocks. And so why do we want to design proteins? Uh, one reason is we might want to do new chemistry. A few years ago, Arnold Lab came up with an enzyme that can catalyze the formation of carbon-silicon bonds, which are known to be done but not previously were not done by proteins or enzymes. Sometimes we engineer proteins or design proteins to understand how they work. These are— this is an adenovirus, adeno-associated virus. Which people are interested in for gene therapy, and people engineer the proteins on the surface to help it get into different kinds of cells. Sometimes we can use a protein directly as a therapeutic. The most common ones nowadays are antibodies, but we can also do these things like enzymes. This is an asparaginase enzyme which breaks down the amino acid asparagine. Turns out that Some cancer cells can't make their own asparagine, so if you break down the asparagine in this tumor environment, it kills off cancer cells faster than it kills their healthy cells. We also use proteins as molecular tools. This is GFP, perhaps the most famous molecular tool. It stands for green fluorescent protein, and as the name suggests, when activated with the correct wavelengths of light, close screen, you can use this to visualize constructs within the cell. And so the fundamental protein engineering problem is we know in the physical world your amino acid sequence, and here I'm representing the amino acids with a single letter each, so M stands for methionine, etc., and that there's a primary structure there which is the amino acid sequence, the order in which we covalently like those. That determines in most cases the 3D structure. How to get from sequence to structure is a hard problem, one that is— we're pretty good, getting pretty good at predicting that using machine learning with things like AlphaFold. And then we also want to know how structure determines function, which is even harder. But ultimately in protein engineering, what we want to do is find a sequence that would give us the function we want. We want to specify the function, get the sequence. And so today I'm going to talk about using multiple data modalities to discover and design proteins. So the focus today on the design— or the, I'm sorry, on the discover. So the general paradigm here is we want to go look in some databases or in the natural world for natural protein data. We want to train some initial model And these models can be multimodal. So the simplest thing you can use is obviously the protein sequence, the sequence of amino acids. But there are also things like annotations, functional annotations, or species annotations, etc. Many times we can get an experimental structure or predicted structure. You might have some chemistry information. You might have some free text descriptions. You might be able to get some energetic information, etc. So we're going to combine some subset of these into one model. Then once you've trained an initial model, there's kind of two ways we can use it. The first is to use it, think of it as a pre-trained fine-tuning paradigm where you've trained this model on a lot of data. Now you'd like to predict the fitness of individual proteins for the fitness on some task that you define for engineering. And what you'd like to do is feed in sequences and predict a number for how good it is at doing something, or that something could be how bright is your fluorescent protein, it could be how fast does your enzyme do this reaction, et cetera. And the hope there is that by pre-training on large amounts of data, you can make better predictions on small amounts of data that you collect in the lab. The other paradigm is this machine-generated data with proteins. So the, the pipe dream here is you want to specify desired function, pass that to a model which has seen a bunch of natural data, and that model then spits out sequences that hopefully do the desired function. And today we're gonna focus on this generation paradigm. Discovery paradigm. And so first I'm going to start with a short story of how to use sequences and annotations to generate just parts of proteins. So many proteins in nature have what are called signal peptides, and they're secretion signals. What that means is when a protein is made, it includes the rest of the protein, signal peptide, and the cytoplasm. That signal peptide helps it find the the channel and to be exported from the cell into the exterior. It's then cleaved and the mature protein then folds and goes to this thing and the signal peptide gets broken down. And we'd like to generate custom signal peptides for arbitrary proteins. I'll get into why I want to do that in the next slide, but this is a project that I did with a fairly large group of people because we did wet lab experiments, but the primary collaborators here were Alicia, who is now actually a software engineer at Microsoft. She was an undergrad at Caltech. And Zach Wu, who was in my PhD lab at the time. Now he's a research scientist at DeepMind. And we want to be able to generate artificial signal peptides because signal peptides simplify protein. So when we make proteins in cells, the— if there's no secretion, what happens is they build up in the cell, and that can be very bad for the cell, especially if you're overproducing an artificial protein or a protein that's not native to that species. So if you can get that protein out of the cell, it's a less burdensome to the cell and also makes your downstream separation easier. So what you'd like to do is you'd like the proteins to be secreted as sermates, and that means that later when you want to separate them instead of having to lyse the cells and do a bunch of steps, you can just spin them down. If their proteins are dissolved in the supernatant, cells are at the bottom, and you've got your first step for separation. Single peptides are structured in the mathematical sense, not in the sense that they have a 3D structure necessarily. But in the sense that they are not completely free form. So they start with a, a fatty at the N-terminus, and then there's usually 15 to 30— they're usually 15 to 30 amino acids long. The N region is generally a few positively charged residues. There's a hydrophobic core which often forms a helix, and there's a polar C-terminal region often ending with one of the few common motifs such as AXA where X is an amino acid or A or T followed by C, and then a cleavage site between it and the rest of the mature protein. And some more details, they're responsible for extracellular secretion of prokaryotes and secretion from the endoplasmic reticulum of eukaryotes. And the signal peptide sequence is dependent on both the species and the protein sequence. And these simple rules I've shown above are insufficient to generate new functional signal peptides. So we try to do this using a machine translation paradigm. So the idea is you feed in the sequence of your mature protein to an encoder stack, and we're using Transformers, and then you feed that to decoder, which spits out the signal peptide. We used a transformer architecture. We actually didn't have that much data. We trained this on 40K examples from UniProt. We actually only used SwissProt annotations, which is much smaller, and we got up to 40% test accuracy. Here are some examples of actual and predicted signal peptides. You can see some of them look pretty good. They've learned that things start with M. They've learned that there's a fair amount of Bs and Ls. And let's see how this looks in the lab or computationally first. So first we try to see if the things we generate can convince computational signal peptide predictors that they're signal peptides. And these are SignalP, is a state-of-the-art signal peptide predictor. 5.0 is state-of-the-art at the time. 6.0 is out now, which uses neural network language models, and they actually show as part of their work that 6.0 can distinguish artificial ones from real ones. But back in 5.0, we tricked it pretty well. So one kind of baseline is a hidden Markov model, which only looks at the previous state when generating the next state. And look at the hidden Markov model. For the most part, SynLP can tell that those are not real sounds, and the mean probability is about 60%. If we use some heuristics such as the ones I showed earlier, we actually do much better than the profile HMM. If we use a variational autoencoder, we do even better. And the transformer, which is conditioned on the mature protein, does about as well. The variational autoencoder is unconditional, so we're just generating— trying only on single peptides. It doesn't look at the mature proteins. The reason we ended up going with the transformer framework, the machine translation framework, is because the VAE sequences tend to be very low entropy and repetitive. And that's kind of analogous to how images generated with VAs tend to be blurry. Okay, so they work computationally. We also then took these generated single peptides and we start— we generated single peptides for several classes of enzymes that we're expressing heterologously in Bacillus subtilis. And so we want to show for each that we can get comparable or better performance using generated versus natural signal peptides. And so 79% of natural signal peptides we dosed in were functional, and about 48% of the generated ones are functional, or meaning that the enzyme is secreting has activity. However, of the ones that are functional, Um, of the ones that are functional, they have often have comparable activity, and some classes of enzymes are actually have much higher activity and some we have much lower. So that's pretty variable depending on the enzyme class. Of course, we could just be doing this by memorizing natural signal peptides and spitting them out, and we find that that's not exactly the case. So the ones we generated are functional, have on average 0.75% activity to the closest natural signal peptide that we have. And the nonfunctional ones are pretty similar. They're a little bit less, a little bit more diverse, but not by much. And if you just look at natural ones to see what the baseline diversity in signal peptides is, they tend to be about 82% identical to the closest natural— other natural. Some quick conclusions here, we can generate de novo signal peptides conditioned on the mature protein that looks like real signal peptides and result in functional secreted enzymes. Of course, this is a much simpler problem than generating entire proteins, which I'll talk about next. I'll pause here if anybody has questions.

I have one question. Do you impose any limit on the size of the peptide or the composition?

No. So we actually can— I don't show it in the slide explicitly, but following the typical autoregressive transformer framework, we actually start with a start token. And have a stop token, and we just stop it when it hits the stop token. Okay, I see.

There's a few questions in the chat.

Right, so we did not try reinforcement learning. Do you know the difference between generated peptide and the test peptide correlates to the odds of being a functional peptide? I don't think it correlates very much. I'm pretty sure that's in the paper somewhere, but I, yeah, I don't think it correlates very much, but it's also hard to correlate there when we are purposely choosing things that are not super similar to the natural ones. Are these functional signal peptides across different enzymes for the very same similar protein signal peptide? So the functional ones in the plot I showed on the previous slide, are the ones that are here. So they're the ones that are in the teal. I think that answers the question. So here, the blue ones here are the teal ones from the last slide. The natural cyclopeptide here are the orange yellow ones from the last slide. The nonfunctional are the ones we generated that didn't work the last slide. The priority is— okay, so the priority is not your diversity or functionality or trade-off balance. I think that depends on your specific downstream task. So the way we originally motivated this was you would have a system where you don't— where the natural— you're expressing an enzyme heterologously, meaning not in its native host, the signal peptide from the native host doesn't work in that— in the heterologous host. So you can think about just trying to come up with anything that will work in the host. Would it be possible to use a self-supervised pre-trained model like ESM and fine-tune it? You could think of that, yeah. Or yes, you could do something like that. That's what SignalP 6.0 does, actually, to detect single peptides. Is the secretion signal more of an abstract task without specific focus on targets, transporters, cell types? Yeah, so the follow-up here would be to condition on the species or the cell type, all right, so that, uh, that gives you an easier way to map. For heterologous expression. Okay, so next section. So now we're going to talk about generating entire proteins. So earlier we talked about sequence determines structure determines function. Here we're, for this part of the talk, we're going to kind of not worry too much about the sequence part and just try to generate new structures. And the idea is that if we can generate new designable structures, we can expand function space. And by designable, what I mean is that we can find a sequence that will fold to that structure. So we're assuming that part is possible for well-designed structures. And the challenge is we want them to be diverse and designable. This is a project that This was originally an interim project with Kevin Wu, who's on the left there. Alex Lu is one of the other researchers here. James is Kevin's PhD advisor at Stanford, and then Rianne and Ava were also researchers at MSR working on this project. And so to understand this work, we have to dive a little bit into the protein biochemistry. So here's an amino acid as before. To make a protein, we link these together, and the way you do that is you take two amino acids. One of them loses the OH purple, another one is H, a condensation reaction, and what's left are the amino acid residues, and they're covalently linked. The water floats away, and, uh, to make a protein, you just do this many, many times. Different amino acids or the R groups differ, and the structure is determined by the bond angles of these residues. And which bond angles, what do I mean by that? So there's these, uh, you can rotate these bonds that I've highlighted in orange. You can, um, change these dihedral angles here that I've highlighted in purple. You can also rotate the side chain, and some of the side chains, depending on how big they are, have further bond angles. But, uh, the 6— these dihedral bond angles I've highlighted here— so we're going to ignore the side chain for now. These 6 angles, if you do this for every consecutive pair of amino acids, you specify those, fully determine what we call the backbone structure, which is the structure without the side chain. The reason we don't want the side chains right now is because we want to design the backbone first, then we want to— so we want to generate the backbone first, then we'll go back and design the amino acid identities, which determines what the side chains are, and hopefully your backbone can accommodate those side chains. And we're going to use a framework called the Denoising Diffusion Probabilistic Models to generate to generate backbone structures. And so the way diffusion models work is at generation time, you start from random noise, in this case random angles, and you want to go— you want to take away that noise step by step using the neural network to come up with realistic-looking angles, which hopefully give you a a realistic-looking structure. And the reason we want to use this angle formulation is because it kind of— it's an elegant way to kind of think about constraining the space of noise we want to add so that at the end, when we start with random noise, our atoms aren't all over the place. They're still in a chain, and that chain has just become unfolded. And so we can think of the generation process as folding that chain. And here in the circle, we kind of see the probabilistic model for diffusion. So in the forward— in the process going to the right, it's a generation process. If it goes to the left, it's a noise process. During training, we noise the angles, meaning we add random noise to them, and we try to predict the noise that was added in the previous timestep. We do that prediction with a vanilla transformer, which X_{t-1}, it just means the noise that was added from the previous timestep. So X_0 is the original, X_T is the final completely random set of angles, and in between we're adding a little bit of noise. Or taking away the limit suppose. We trained this model on 30,000 structures from CASP, so not a very big training set, and we're limiting ourselves to things between 40 and 120 residues. That's mostly for computational efficiency. Another reason to use this angle formulation is that diffusion on— if we diffuse directly in the 3D coordinates, we have to take into account the equivariances of of rotation translation, which means you have to use a much more complicated neural network than vanilla transformer. On the left, you can kind of see the diffusion process or the generation process for a point cloud diffusion. On the right, here's an example of generating from our full-length model where you start from a random string and you come up with more realistic angles. And we want to evaluate these generations at 3 angles— at 3 levels, sorry. We want to look at the generated angles themselves, the structural motifs, and the overall structures. So first we start from noise, we sample, we compare to the true distribution of angles. We're just looking at the individual angles. The way we do that is we, for each of those 6 angles, we look at the empirical PDF and compute a KL divergence between the test set and what we sample, and lower is better. So for phi, we can see visually, it matches pretty well, the KL divergence is pretty low. Same for psi and omega and all the thetas. So that's a good sign that our generated distributions match the natural distribution for the individual angles. But of course, we don't care just about individual angles. We want to capture the correlations between angles. So one canonical way to do this is look at the phi-psi co-occurrences, meaning the phi-psi in within the same residue because that largely determines secondary structure. So if we plot phi, psi against each other for the test set, you can see there's some beta sheets, there's some alpha helices, but mostly right-handed, not very many left-handed. And our generated ones match pretty well. So we can also generate the correlations that define common structural motifs. And if we also predict secondary structures, so we noise sample and then predict helices and sheets. So alpha helices look like this, beta sheets generally look like this. And so here's a true kind of co-occurrence matrix for alpha helices, beta sheets in our test set. And here's what the generative ones look like. So these also look pretty similar, which is good. And these are somewhat cherry-picked, but we can generate realistic-looking structures, which is also good. So on top you see generated ones, the bottom you see some training set structures. To kind of quantify how reasonable these structures are, we use something called designability. That's measured by a self-consistency TM score, which I think was first described by Brian Tripp and his backbone diffusion paper. So the idea here is to generate structure from pulling diff, you run it through a protein design model, this protein@pnn, which predicts the amino acid chain, the residues, run that through a structure prediction model like OmegaFold, and you compare how consistent that structure is with your original generated structure. And we're using the TM score, so higher is better. If they're exactly the same, the TM score is 1. And because two structures with a TM score above 0.5 are considered to be the same fold, people use 0.5 as a cutoff for whether things are desirable. You can see that for both generations between under 70 amino acids and over 70 amino acids, we can generate a fair number of designable backbones. Compared to a point cloud diffusion method, we do better at both short and long generations, so it's significant improvements over point cloud diffusion. Nowadays, there are— since in the like 6 months since we did this, There's probably been a half dozen other backbone diffusion methods that do better. People have come up with really creative ways to parameterize the backbone representations. This is a very fast-moving field. We do also show that these structures are better than a random baseline where you just sample sets of angles and shuffle them. The randomly sampled angles mostly do not generate designable backbones. We are learning higher-order correlations or distant correlations. Our structures are fairly diverse. On top, we have our generated structures. The second row shows the closest structure in the training set. And so sometimes they're pretty similar and sometimes they're less similar, but none of these are exactly the same. Highest TM here is 0.71. And then we also design using ProteinPMN and refold it using MegaFold to see if it's designable. And some of these are more designable than others. Another way to look at diversity is to cluster the backbones or to look at the distances between them. So on the left, we have natural protein backbones. In the middle, we have ProtDiff, which once again is point cloud diffusion. On the right, we have generative repolling diff. You can see using that angle formulation gives you much more diverse samples. Another way to look at this is to compare the self-consistency TM score versus the similarity, how similar you are to the closest training structure. Ideally, you want things that are very designable and very far from the training set. Unfortunately, it is still very hard to get there, but there is a pretty strong correlation between how close we are to the training set and how self-consistent we are according to the current learning models. And we also check that to make sure that the angle formulation combined with the vanilla transformer is sufficient. So here is, for example, the C-alpha, C-alpha distance map from a real protein. Here's something we— if you use a point cloud instead of the angle formulation in a transformer without equivariances, you don't generate any of these high-level, um, this contacts. And so FoldingDiff is a first step for centering functions. We generate protein backbones by varying the the folding process. We use this internal angle formulation. We can successfully generate structural motifs and realistic diverse samples. The code is available on both GitHub and Hugging Faces, and it might eventually be in BioNemo from NVIDIA. I see I have some questions in the chat. Um, okay. What is the point of diffusing over the bond angles? I would think this is not productive since they're fairly rigid, right? So for most proteins, once it's folded, those bond angles are fairly rigid, but we're not modeling a physical process here with our diffusion. We're using diffusion to generate new examples of proteins, of new structures, which in this case means new examples of sets of bond angles. Don't the errors of each model propagate down the pipeline? Or does the SCTM score prevent that from happening? So yes, I think what you're asking is when we do the self-consistency score, every model makes some mistakes, right? And the intuition here is that something with a high self-consistency TM score means that both protein and PNN and OmegaFold like that backbone structure, and in most cases that means it's going to work pretty well. Something with a low TM score could also work, but just be outside of what ProteinPNN and OmegaFold have seen, and it only takes one of them not making it, right? So high TM scores, I would be fairly confident. High SCTM scores, I'm fairly confident in. Low ones, there's probably some false negatives. Any info on doing diffusion in latent space, either alpha-fold MSA or pairwise embeddings instead of either 3D coordinates or sequence? So basically stable diffusion style. I'm sure there is somebody working on that. Uh, the next question. Just to clarify, you're generating structures first without knowing whether you're actually making amino acids during the batch? Yes. So this paradigm is you generate a structure, and this, uh, really works well if the function you want can be specified by the structure, right? And that includes things like binding, where if you can generate a low-energy structure that binds the thing you want to find, and then if— and if that structure is designable, meaning we can find an amino acid that folds to it, then you're 90% of the way there. Right now we do not condition the functions that we want. In the future, that's something we would like to work on. And there is some other work out there where instead of generating the entire backbone, they specify the position of certain residues that believe are, well, bind the thing you want to bind or do the reaction you want to do and generate the rest of the structure around it. How does your model perform in small chains? So anything that's too small isn't going to have a lot of tertiary structure, but it probably does okay. But we did not this and test that. All right, so we talked about generating in structure space, but as people point out, you still have to design the sequence from the structure you generate. So next we're going to talk about just generating directly in sequence space. So if we can generate new valid sequences, that expands functional space. The benefit of sequences is we have many more of those than we have structures, The downside is you're throwing away— the structure provides a lot of information. So the idea here is we can generate sequences, we can generate functions. The challenge is to generate valid, or some valid meaning it will fold and do something, and diverse sequences. And this also started as an intern project actually, but it's led by Sarah Alandari and Matthias. So Matthias was an undergrad intern. Last summer. Sarah is a data scientist here, with help from Nicola, Ryan, and Ava, who are all researchers here. And so we're going to use— before, we did continuous diffusion, meaning we take some data and we add Gaussian noise to it to reach a stationary state that's just a unit Gaussian noise in every dimension. Now those sequences are discrete, and we're going to keep them discrete. You You can— there's work out there that pretends they're continuous, or you can embed them in latent space and do continuous diffusion, but we're going to corrupt them directly in the discrete space. And there's a few different ways we can do this. We tried 2-ish, 2.5. One way is you can corrupt by adding— by replacing amino acids with a mask, with a special mask. So at every corruption step, we choose a position and we mask it. And at decoding step, we choose position and we unmask it. And we call this model order-agnostic autoregressive diffusion. And this is a really nice framework because it links diffusion in that you're corrupting something and uncorrupting something to generate. It links masked language models and autoregressive left-to-right language models. Here's how it does that. So in a masked language model, you take a sequence, you corrupt usually 10 or 15% of the residues, and you predict everything. You predict everything. You predict the masked positions from the unmasked positions. And you can think of that as learning a single— training a model that only learns a single diffusion step. All right. Um, whereas OARDM generalizes this by learning all the possible masked language models in a single model. But what also generalizes left-to-right autoregressive models, LRAR, which you can think of as you always mask everything to the right of what you've already generated and you want to generate, then you unmask the next position. So this is a left-to-right autoregressive model, only learns a single decoding order, whereas order-agnostic autoregressive model learns all possible diffusion, all possible decoding orders in one model. So OA-ARDM generalizes both masked and left-to-right autoregressive language models, and it's also a diffusion model because there's an absorbing state, which is all masked. The other way we consider is a framework called discrete diffusion probabilistic models, or D3PM. And the idea here is that every step, instead of taking everything to us, corrupting everything to a state where everything is masked, you corrupt to a state where every position is generated uniformly. And the way you do that is at every step you have some probability of mutating every position to something else. So here you can see the stuff I've highlighted, we've— two of the positions have mutated. You just keep doing that until you arrive at this random state. And you choose how to mutate things using what's called a transition matrix. And we consider both a uniform matrix where at every position any amino acid can mutate to any other one with equal probability. We also use one derived from BLOSUM62, which allows us to insert some evolutionary information. In that one, every residue does not have equal probability of mutating to every other residue at any given time. And we're going to train all three of these variants with the same dataset architecture so that we can make fair comparisons along with some baselines. And our base architecture is a 640 million parameter CNN. We're going to train on 42 million sequences from UniRef50, and we're going to look at how realistic are the generated residues individually, We find structural motifs in things we generate. How do the overall sequences and structures look? How realistic are they? How valid are they for some computational— computable metric for validness? And how well do we match the natural distribution? But first, we find that bigger models and diffusion instead of the baselines improve the generated amino acid distribution. So here what we do is for each model we generate sequences and we compare the amino acid usage in those sequences with the amino acid usage in the validation, in the test, in the training set. As a strong baseline, we look at the KL divergence between those distributions, the validation set and the training set, so that's going to be pretty low. And as a Negative baseline, we simply scramble, we simply generate randomly, and that gives you much higher KL divergence in the distributions of amino acids. So first, before we scaled up to 640 million parameters, we used some smaller models. So here we use, we say, we'd use both corruption schemes and the OA-RTM diffusion model. And you can see that in these small models, the OA-RTM model does a fair amount better than the D3PM. As baselines, we also look at a left-to-right autoregressive model, which is the same architecture, same CNN architecture, and trained on the same training set, and CARP, which is a masked language model also the same architecture, also the same dataset. And you can see that these do much worse. We then scale these up. So we have D3PM-Blossom and Uniform again and OA-RTM. So you can see that scaling up in most cases, well, for the Blossom, it doesn't really seem to improve. Sorry, for D3PM, it doesn't seem to improve the single-order distributions much. But OAA-RTM does improve a little bit. Scaling up also doesn't make the baselines much better. Left-right autoregressive CARP are still worse. We also try ESM-1B, which is a transformer mass language model. It's also trained in UNRF50, but they did a different train-test split. It's also slightly bigger, but it doesn't do much better, or actually does much worse at generation. So, an interesting thing is that the original OARDM paper does this for text, and they find that while OARDM in text does do better than D3PM for text, it does worse than left-to-right autoregressive models. We find that for protein sequences, learning all of the decoding orders gives you better generation than learning a single left-to-right decoding order. And as you expect, if you train your masked language model, which is essentially training a single masked coding step, which is from 10% masked to not masked, those are not as good a generation as if you train on every possible decoding set. So next we're going to look at structural motifs. We predict secondary structure for all these sequences. We first look at the distribution of strand versus helices in the validation set. And so ideally the things we generate would look like this. Like this distribution. You can see BLOSSOM kind of over-concentrates on QOCs. The D3PM Uniform and OA-ARDM both do a pretty good job. And here are the controls. The left-to-right autoregressive also does reasonably well, but not as well as the ordering-agnostic one. CARP basically doesn't generate strands, and random generation also doesn't do a good job. So we're beating the baselines, and visually they seem to do pretty well. To be a little bit more quantitative and to look at overall sequence and structure, we have a pipeline. It's very similar to that self-consistency TM score pipeline. So we generate sequences, we run them through OmegaFold to get structure predictions, and we can look at how confident those structure predictions predictions are using the predicted LDVT from OmegaFold. Then we can fold those using the SMIF or ProteinPNN and compare the sequence perplexity between the original generation and the prediction at the end of this pipeline. So here is the perplexity. You can see— and lower is better— on the left, you can see the gray line from the validation set. So that's kind of the thing we'd like to match. On the right is random, that's the bad thing. Everything is in between there as you would expect. OA-ARDM does do better than anything else here, followed by BLOSSOM and then Uniform, and then the left-to-right autoregressive and CART both do pretty badly. You can also look at the PLDDT, so this is how confident is OmegaFold in the predictions. Once again, this time higher is better. The validation set, C, is towards the right. On the left is random sequences, which are bad. Once again, OARDM does do better than the other models here. But everything— okay, so quick question. What sampling procedure are you using to sample sequences from masked language models? We're using the autoregressive, the ordering autoregressive diffusion sampling scheme. So we're masking everything, sampling one by one. Kevin, sorry, my question was about the baseline, ESM-1B. Yeah, so same thing. You would mask everything and You unmask one at a time. Other people, other people being me, have also tried using like a Gibbs sampling procedure from masked language models, where you mask some smaller percentage in general and kind of do a random walk. The difficulty with that is there's a lot of hyperparameters, because the number of steps, the number of places you mask and the starting sequences all matter there. So this was the simplest way to do it, although we might try some Gibbs sampling in the future as a stronger baseline, possibly. Does that answer your question? Yes, it does. In the first case, when you mask k tokens in parallel, you're saying your procedure is you unmask one, then you feed in the new sequence with the other mask and you repeat that? Okay. Yes. Yes. So at train time, you predict all the masked tokens, but at generation time, you choose one. You predict all of them, but you only unmask one. Okay. Thanks. We can also look at, uh, if we, if you want to set hard cutoffs for things, we can say anything that's better than the average for the validation set is foldable given the PLT or low complexity given the perplexity. And the intersect, if it's in the top quadrant for both of those, we're calling that a high-fidelity sample. So that's the first 3 columns here. You can see that the validation set, as you'd expect, about half are foldable, half are low complexity, and a quarter are high fidelity. That's based on definition of how we set those cutoffs. And the minimum Hamming is what's the— how far away are you from the closest training sequence? So higher is better. Higher means you're more dissimilar from the training set. And so the validation set here has some pretty similar sequences to the training set. And of course, random, it's very diverse, but not many things are well past these cutoffs. If we look at our diffusion models, once again, OA-ARDM does seem to do better here. It is both— it is slightly less diverse, actually, but you get much higher quality samples. And the baselines are more diverse, but the sample quality is much lower. And here are some cherry-picked examples from the OA-RTM model after putting them through OmegaFold. So you can see it is possible to generate things that are high confidence and have a variety of architectures. Finally, we want to look at the natural— can we recapitulate the natural sequence distribution? So here we generate from each model a bunch of sequences. We embed them using ProTransBERT, which is a pre-trained language model. Then we do two things with those embeddings. We can look at them using UMAP. So here they are. So gray is the validation set, or a sample from validation set. And the color is that specific model. You can see, and you want to match that gray distribution as well as possible. You can see that in terms of baselines in the bottom don't match this distribution as well. They tend to kind of cluster in one region, whereas the diffusion models seem to do a little bit better. And we can also quantify this using for saying inception distance. It's not really inception since we didn't use an inception model, we used ProTranspert, but it's analogous to that for images. And so here lower is better, and you can see that random is about 3, everything that's better than that, but OARDM is significantly lower than the others. And as a wrinkle on this, so right now we've been generating single sequences. We realize these frameworks would also work on multiple sequence alignments, and that would allow us to leverage a lot of evolutionary and structural information that's present in the alignment without having to start with structures. The way you get these alignments, you start with a query sequence, you do homology search, and you come up with dozens of thousands of evolutionarily related sequences, and the correlations between mutation positions here give you a lot of structural information. That's why all the best structure prediction algorithms right now actually use— do an MSA search step before they run their structure prediction. And so we're going to noise and denoise these exactly the same way we do with the sequences. So you can do OARDM, So now we just treat this as one big thing. You can mask anywhere at any step. And at denoising time, you can do unconditional generation where you just denoise the whole thing. You can also do D3PM where you can mutate any position in any sequence at any time. At the end, you get a bunch of random sequences. And so we're going to use the— once again, want to be consistent. In our architecture and training sets we can compare. And we're using 100 million parameter MSA transformer architecture, which is a 2D transformer, and we train these on 100,000 MSAs from the OpenFold training set. And here we're going to look at generated residues and the overall sequence instruction. So once again, we find that the order-agnostic autoregressive diffusion better recapitulates the natural amino acid distribution. So here we're generating alignment and looking at the natural amino acid usage. So validation set, we compute KL with and without gaps just because now we have gap tokens because of the alignment itself. And ESM-MSA is trained— it's the same architecture but it's trained on a much larger training set, but it's a masked language modeling test. And because of the masked language modeling task, it has a lot of trouble generating from scratch. Once again, we could do Gibbs sampling. That would probably give us slightly better samples, but once again, there's a lot of hyperparameters. We show that— so in our training, we have the random and the max are different ways to subsample the MSAs based on are we sub-sampling them randomly or sub-sampling them for maximum diversity. All the architectures are the same. We see that the two OA-ARDM variants generally outperform the other ones depending on which metric you use. So this is just looking at the individual residues. We would like to look at the overall sequence structure. And one way to do this is now the nice thing about this, the other nice thing about these masked diffusion models is you can do conditional generation just by specifying some residues and only masking the ones you wanna generate. So here we're gonna start with an alignment and regenerate the query sequence and original sequence. And first we'll look at diversity. So here, because this is kind of expensive to run, we only look at the OARDM, and then we do a POTS model, which is just, uh, it is explicitly trained on that MFA and it generates new sequences. And so similarity, lower is better, but lower is more diverse. Since the POTS model is significantly less diverse than the diffusion models, and then we can look at a few metrics. We can take that generated query, run it through OmegaFold, and get a structure where you can look at this, and then we take that structure and design the sequence again and look at the sequence perplexity. So here, lower is better, and you can see that even though the POTS model is more similar to the natural sequences, it— this pipeline doesn't like them as much. We can also look at just how confident OmegaFold is once again. And here, higher is better, so to the right is better. That's where the validation set is. And here, once again, OA are giving us slightly better than the POTS model. And we can also look at TM score between the structural prediction you get with the true query sequence versus the one we generated. And here, we see that the POS model actually does much worse than both of the diffusion models. All right, to summarize here, we've used diffusion models to generate structure, sequence, and MSA. Ultimately, no matter what you start with, eventually you'll be able to make things in the lab, always in sequences. This is why it's exciting that we can do this on sequences directly now. And of course, looking ahead, our goal is all is always to generate proteins with new and useful functions. And that means using different kinds of data like structures, annotations, and chemistry in the future. And we kind of think of natural sequences, biological sequences give you structure and then they do cellular functions, but we really want to use neural networks to generate new sequences that do engineered functions that people care about. And now I will take more questions. So I see one in the chat. Uh, do you have any notion as to why diffusion seems to be outperforming traditional MLMs in the context of proteins? So I think a better way to phrase this question is, why would order-agnostic autoregressive diffusion models outperform the left and right autoregressive models in context of proteins as opposed to NLP? And the answer here, the simple answer is that in text, left-to-right decoding is in the easier half. It's easier than the average decoding order, whereas in proteins, it's harder than the average decoding order. And the reason that the mechanism for that is kind of murky and up for debate, but intuitively we do kind of generate text from left to right, even though there are some higher-order dependencies, whereas apparently for proteins, you can— they're not specified entirely left to right, even if they are made left to right. Do you think generated proteins have a chance at being more specific in terms of interactions compared to natural ones? That's hard to say. It depends how you generate them. If you do like a conditional structural generation, like some of the stuff that Baker Lab does, then yeah, probably, because you can design the interaction you want and you can purposely destabilize the interactions you don't want and Although when it comes to enzymes, most natural enzymes are already very specific, so that's for their natural reactions, that's pretty high bar. Yeah, that's an open question.

In the last part of the talk, you focus a lot on comparing the generated distribution to the natural distribution, and as I was wondering if that's not problematic by itself in the sense that you want this generative model to allow you to explore things that are not—

Yeah, so on the one hand, yeah, you want to recapitulate the natural distribution, right? On the other hand, you want to generate diverse things. The question is, can you generate diverse things by like filling in holes in the natural distribution, or does it have to go outside, and what does that even mean in high-dimensional space? Yeah, but I agree, that's only one thing we should measure.

It's not the end-all of quality measurements because there could be an argument made that if you just fill in the hole of the natural distribution, like any— you don't need an AI model to do that, right? If like maybe the benefit of the AI model is coming from exploring outside of natural distribution and finding things that are— might seem unnatural or kind of very weird for us, but yet work more effectively than, than the natural ones, right? I think we have the same problem in small molecule generation, right? Do we want to synthesize things that chemists, um, look to and are comfortable a weave of things that are very weird, but we work anyway.

Right. So Hugh asks, are you— is there a paper on the last part of the talk? There's not yet. It should be out in the next 2 weeks. We are desperately writing it as we speak between me making these slides. And if there's no more questions, just want to recap. That we did this. I did most of this work at MSRI New England in the BioLab group, which is just a wonderful place to work. And many of these people are already mentioned, but there's also Kristen and Lauren and Phil who I work with on other things, and Maddie, who are not in these particular projects. Oh, I have another question. Do you think there's a space to include an evolutionary representation in your model? More likely or similarly like the amino acid or protein large language model, like you guys thought? Yeah, so we kind of use BLOSUM62 as a transition matrix to try something simple and evolutionarily defined or inspired, but you could use a neural network transition matrix as long as it's doubly sarcastic. It will arrive at the same structure distribution. So going back to predicting the structure first, did I just miss how you're going back from structure to amino acid chain? I feel like this last part we just discussed is more about sequence to function. Right, so when we predict the structure first, there are now a fair number of methods that take in the backbone structure and will predict an amino acid sequence or set of amino acid sequences. That will fold to it. And they tend to— the state-of-the-art methods now tend to be graph neural networks. And protein PNN is generally the easiest one to use and very accurate. Ironically, you can think of protein PNN as a structure-conditioned, order-agnostic diffusion mass language model. It's basically OA RDM, but with structured permissions.