Nathan Frey presents discrete walk-jump sampling, a machine learning method for designing therapeutic antibodies. The approach redesigned trastuzumab's CDR H3 region and achieved 97% expression success with superior binding affinity, outperforming GPT-4 and diffusion models while requiring only 100-200 training examples.
Questions this video answers
There is essentially one hyperparameter: the noise scale. A theoretical derivation links it to data sparsity, but in practice σ=1 works very well for one‑hot vectors, so we typically leave it unchanged.
The method is not limited to one‑hot encodings; you can use any embedding space, including the latent space of a VAE. You would just adjust the noise scale to match the embedding’s scale.
The approach is highly sample‑efficient: you can train on as few as 100‑200 examples for a targeted distribution, or on the entire observed antibody or protein space, depending on your goal.
The JumpStep is a denoising network that learns the Bayes estimator for removing noise from the noisy one‑hot vectors. It can be any expressive neural network (e.g., ByteNet or Transformer). The theory from neural empirical Bayes ensures a single noise scale suffices.
Chapters
I'm Nathan. I'm a machine learning scientist and group leader in Prescient Design. And yeah, I think I'm going to jump right into it. Alex and Kevin gave some great context here. So I'm going to be talking about a new method that we developed for protein discovery. I'm going to focus on antibodies because this is an area where we run a lot of experiments in Genentech. I can show you some exciting wet lab results. For my purposes, an antibody is a very simple thing. It's just a protein that recognizes some other protein and binds to it. That's the antigen. And then kicks off, you know, some kind of response. So as we've seen already, large language models are all over the place in protein design lately. We've also been very interested in tracking a lot of progress in diffusion models as well. There's so much work in this space that they're now helpful lists all over GitHub where people try to keep track of all of the work that is going on in this space. But for us, therapeutic proteins present a lot of very unique challenges that make a lot of prior methods not perfect fits. So in particular for therapeutic antibodies, this space, like all protein space, it's very large. It's a discrete space. Space of choices of amino acids. And most importantly, it's very, very sparsely functional. So antibodies that can actually become drugs are very, very hard to find. Another interesting thing about antibodies, immune system proteins have very well-conserved regions and also very high entropy variable regions that undergo rapid evolution. In your body or in animal immune systems. So this makes things like evolutionary learning and large-scale training very difficult, just like we saw in the TCR space. So I'm going to lead off with the headline, which is that we have kind of a brand new method that's not a large language model, it's not a diffusion model, and we can generate functional binding antibodies to new targets., and how did we do this? The basic idea is to kind of reconsider what it means to generate antibody sequences with machine learning, with generative AI, and turn this very difficult problem of discrete modeling of sequences into a much easier smooth sampling problem. So when we think about generative AI, in diffusion models, we generally have, or other types of generative AI, we encounter energy-based models and score-based models a lot. For an EBM, the idea is to model the energy surface of the distribution that you're interested in. But for discrete spaces, these energy surfaces are very, very rough. This leads to brittle training and difficulty sampling. There are a lot of tricks to rescue EBM, but if you're a recent, I guess, convert to the ML space, you may not have even used an EBM before. They've been abandoned largely recently. Score-based models are probably much more familiar. These are underlying basically all of the diffusion models that people are using these days. And here, the idea is to learn the gradient or the score of a log probability density. In a score-based model, we have a noise schedule. So we're learning the gradient of a score function of noisy data over many, many, many noise scales. So from some kind of very simple noisy distribution to the actual things we're trying to model, this is also brittle. It's very hard to train a diffusion model in a new space if you're the first one doing it. And sampling is difficult as well. There are lots of tricks and cool things to make diffusion models work well. So we have kind of transformed this discrete modeling problem into a very simple one of trying to model smooth, noisy data. And we've introduced this method called discrete walk jump sampling. So the idea is that discrete protein sequences that form a distribution is very, very hard for a machine learning model to learn. So don't do that. Instead learn the surface of some kind of noisy smooth space, which are not proteins. They're, they're, it's a noisy distribution. Train a model, could be an energy-based model, could be a score-based model, to walk or sample this new smooth energy landscape and then sample from that. So you're sampling noisy data, not, not proteins yet. Then you have a completely separate, completely decoupled model that is trained to jump back from this noisy data back to real protein sequences. So these are totally separate, and the intuition is very similar to diffusion models. If I have my discrete space, my data density where I have very few examples, and I'm trying to learn this score function, it's only going to be accurate where I have data, and everywhere else, it's not going to work. So the easiest thing to do is to perturb that data density. The easiest way to do that is to add noise. This gives us support over the whole space, and we can learn a good score function over the whole space. This is just a picture of kind of what this score function is doing. This is kind of inspired by work called Neural Empirical Bayes, and it's very simple to understand. The Xs here are your noisy data that you're actually trying to model. Y is this smooth manifold that we've introduced. And when you train a model to denoise or to jump, it's kind of a compass. It's telling you via this learned score function how to return from this smooth space, which is really easy to learn and easy to sample from, how to return back to the space that you care about. I'm going to show this again because it's— and say it just in a couple of different ways because it's usually brand new and unintuitive. For a protein, this is really simple. You take a one-hot encoding of your sequence that just says what amino acids are present in which indices, and you add Gaussian noise to that one-hot. So it's no longer a simple representation of your protein, some noisy one-hot vector. We train, in our case, usually an energy-based model to learn the energy surface of those noisy vectors, that noisy vector space. And then we train a separate denoising model, usually a ByteNet, could be a transformer as well, to take samples, take these noisy samples, and jump back to clean proteins. So that's why we call it walk-jump. This is a Langevin MCMC, so it's a non-autoregressive method where you use the gradient from an energy function to sample. That's the Langevin MCMC, or the walking part. And then you use the score function to jump back to clean sequences. So one way to think about this is as a single-step diffusion. There's only a single noise scale in this method., and you can sample a new, a new noisy sample at every step and turn it back into a real protein instead of running this long reverse diffusion to get a single sample. So if you have seen like a molecular dynamic simulation or a structural diffusion, these sorts of methods let you explore a local conformational landscape, right? So you will see the atoms moving in space. WalkJump kind of enables you to do this, what we call an alchemical ensemble exploration. So at every single step of this chain, you can very easily jump back to a real protein. You can use a fast folding method. We use an internal method called Equifold. You could use ESMFold, anything that's really fast, faster than your sampling. And then so you get a structure. That is mutated at every single step. So this is like exploring the local mutational landscape around wherever your starting point is, which could be a sequence you're already interested in, or it could be sampling from total noise. I'll just go over this briefly. So the method, you know, it outperforms autoregressive methods, large language models, masked language models, GPT-3.5 and GPT-4. There's a kind of cute benchmark I did against GPT that if you're interested, you should check out the paper. It's actually these methods with good instruction fine-tuning are actually not so bad at protein design. It's of course much, much faster than diffusion. You get much faster sampling because of the single noise scale and better sequence diversity as well. So that's our in silico evaluation, but how does this perform in the lab? So we have a kind of toy example here on trastuzumab. So binding, trying to bind to HER2. There's a nice dataset of mutagenesis experiments that modify a particular slice of the sequence of trastuzumab that's most important for binding. It's a large search space, even if you're just trying to redesign this this one chunk, but 10 to the 13th dimensional search space. So the task is to generate novel segments, novel CDR H3 segments that are distinct from your training set that retain binding to this target. And so HER2 is not a particularly exciting target, but it's something we can talk about openly. There are also problems with it, but we also have comparisons against language models, structural diffusion, and discrete diffusion. And the kind of upshot is that discrete walkjump sampling has a much, much higher binding rate for redesign compared to any of these other methods. This is without any kind of oracle. There's no filtering. There's no additional model trying to predict binding here. We're just sampling from the model and testing in the lab. To see if we have binding. If we do apply structural filtering, like we try to retain the binding mode or the binding geometry of trastuzumab, the actual drug that we're basing these designs off of, we can improve the binding, we can improve the binding affinity, but it doesn't actually change the overall binding rate at all. This, I can't, share a ton of details yet, but just to say that, you know, this method doesn't only work on HER2, it works on a bunch of different antigens and a bunch of different epitopes. So different binding sites across those antigens. And we'll hopefully, you know, be able to share a lot more fun details about that very soon. So yeah, kind of in summary, I really like this picture from Brian's paper on you know, what are we trying to do with these generative models? We have this full mutational space of possible proteins. There's a submanifold of high intrinsic fitness, so high evolutionary plausibility, where these are high-likelihood sequences that will express in the lab. 97% of walk-jump designs expressed successfully on our very first attempt in lab experiments because the model is just, you know, very easily learning the real manifold. Of intrinsic fitness, then what we're really interested in is this extrinsic fitness. So function, some kind of objective values. And there again, we're designing real binders to real targets for actual drug discovery. Awesome. So with that, thank you very much. I'd like to thank especially Dan Barenberg and Said Seremi at Prescient who did this work with me. Feel free to reach out, check out our paper. All the code is open. Feel free to get in touch with me either at the conference— I'll be around— or shoot me an email.
Yeah, actually faster than you needed to be, so that's great. So any questions? I was wondering if you have any thoughts about balancing the sparsity of the landscape and the amount of noise that you use during the training.
Yeah, absolutely. So there's really only one hyperparameter in this model, and that's the noise scale, right? So we have kind of a toy derivation in the paper, if you're interested, of how to get an initial guess for what that noise scale might be based on the sparsity of your training data. But in practice, it's very— sigma equals 1, noise scale of 1 works really, really well for one-hot vectors. So we didn't really tune that at all. Really cool talk. Thank you. I was wondering if you guys tried operationalizing the Longvin equation as part of like a VAE and if you guys benchmarked against a VAE. Yeah, absolutely. So there's no real restriction to what kind of space you're operating on, right? We chose a one-hot encoding just because it's very, very simple, but actually you could use the latent space of a VAE or any other embedding that you want for the sequence. It would just be a simple matter of adjusting your noise scale to the scale of whatever space you're operating in.
I'll ask one if there's no more. So in terms of the training setup, are you training on like one protein family at a time or more broadly?
Yeah, so the method's super sample efficient. So if you want to, you can train on like 100 examples, 200 examples that you're trying to capture a really targeted distribution, or you can train against the entire observed antibody space or all protein space, just whatever distribution you're trying to model.
I think this is really cool. I'm just, I missed, maybe you went into it a little bit. I missed the, can you say a little more how the jump back model works and how it's trained, what the architecture is, that kind of thing?
Yeah, yeah, absolutely. Yeah, so the JumpStep is just a denoising network. So it's trying to learn a score function where it's very simple. It's just a really— it's taking these noisy representations of sequences, and it's just trying to learn to remove the noise. And the trick is that this is a result from neural empirical Bayes. If you set everything up just right, you only need this single noise scale. And you can basically write down what the exact right operator is to jump back to clean data. But in practice, it is just any expressive neural network that can represent that Bayes estimator.
Okay, great. Let's thank Nathan again. And next up, we have Yuhua Huang from Harvard. He'll be telling us about genomic language models.
This actually—
yeah, I wonder if this works. Yeah, oh, because it's Google Slides. Hopefully, I think the question is whether it, um, the question is figuring out with sharing. I think that should be the right one. Fingers crossed.
Okay, thanks. Right. Hello everyone, my name is Yuna and I'm a PhD student at Harvard, um, and I'm really excited to share
