Long Context Language Models and their Biological Applications with Eric Nguyen - #690

0:28All right, everyone. in biology, in particular hyena and hyena DNA, and most recently Evo. To get us started, I'd love to have you share a little bit about your background and research interests. Sure. So I am in the bioengineering department at Stanford, but my background is mostly in the machine learning side. I started off in computer vision and started working with state-space models and seeing if they can be used for some of these vision and language domains, and then started moving on to see if we can extrapolate them and use them for biology applications, in particular for their long context purposes.

1:04And so I've been exploring these applications of hyena in particular, which is a convolutional-based language model, to the domain of biology, in particular DNA. You know, let's start by talking a little bit about hyena and some of the motivations for that architecture in addressing longer context lengths for language modeling and other sequence models. Tell us about the big challenge there and the motivation. So our lab has been interested in this long sequence task for quite some number of years. And it started with Albert Gu, who was interested in state-space models at the time. And so initially their work in our lab was applying these models to continuous signals typically.

1:48So, like time series and audio. And so, when language models started to become really popular, we too were interested in applying it to language, but we noticed there was a gap in applying some of these early state space models. And so, the Hyena architecture was really focused as an offshoot of these state space models to see if we can get it to work on language. And so, we focused on getting it to work on these discrete signals that we think of as language. So, Michael Polley, the lead author in the Hyena architecture, and Stefano Mazzaroli, the other co-author on the Hyena paper, we started focusing on what's the gap and getting this gap to close between transformers.

2:24And so that was the initial work that they focused on. And it's when we started getting excited, we were able to sort of match similar quality, but then explore this opponent, which was being able to fit longer sequences. And so that's when we got started getting excited about new applications. So the obvious question is transformers have proven to work very well. Why the need for an alternative? What is the challenge that is presented with the transformer architecture in dealing with longer sequences? Yeah. So, you know, at the time when we were working on it, this was still in the age where not many folks were thinking about long context in particular.

3:01And so there's this key constraint with transformers, which as a sequence length grows, the time complexity, right, the algorithmic efficiency of the operation has this squaring law that grows. And so when you double the sequence length, the computational requirement quadruples. And so when the sequence gets very long, if you care about things like DNA or, you know, entire code bases, then it becomes intractable for these models. It can become intractable for, you know, these classic transformer models. And so the motivation is, is there a different type of operation or primitive that can be hopefully as expressive, but will reduce that complexity in computation for these longer sequences?

3:44And you mentioned one of the early attempts to try to address that in particular with state space models. I'll refer folks to my interview with Dan Fu from actually just about exactly a year ago. Language modeling with state space models was the title of that episode. And so this work builds on some of that state space work. Does it retain some of the kind of hierarchical approach that made that that made the state space work? Yeah, it uses a good portion of it in terms of its skeleton. So the differences that we took from that paper with Dan Fu on state-space models for languages was how we parameterize the convolutional kernel.

4:29So a key way to make convolutions work for long sequences in particular is to move on from this sort of explicit parameterization of kernels to an implicit version. And so what that allows you to do is essentially have a parameter efficient way to use convolutions. And so the state space model would use state space models for this kernel. And then the original hyena work we would use, we sort of generalized it to just use an MLP. So an MLP that would parameterize the filters for a convolutional kernel. And so the thinking was to have an unconstrained version that can learn perhaps a more expressive kernel for the convolution.

5:05And your research group published a really interesting blog post last year talking about this scope of research from deep to long learning is the post that I'm referring to. And in that post, there's a really interesting graphic, and we'll link to all of this in the show notes. But one of the elements discussed in that post and the centric to this graphic is the idea of the FFT as kind of a new primitive that's driving this alternative family of models. You know, talk about the role of the FFT in Hyena. So for Hyena, we rely on long convolutions. And so typical convolutions, they'll have short filters, right?

5:48And so the time complexity of having a kernel stride across an input is still basically a linear operation. But when you have a convolutional kernel like Hyena, the key difference is that we're going to have a filter the same length as the input. So a global convolution. Now, the challenge is the time complexity will increase. So you essentially get back a time complexity that's similar to a tension, like a n squared time complexity. The kernel is the same length as the input. And so you lose that benefit of a lower time complexity if you just apply the convolution in this standard way that people think of in the space domain.

6:27Instead, we'll take advantage of a convolutional theorem for folks familiar with signal processing, which states that a convolution done through the Fourier domain is the equivalent as an element-wise multiplication in that space to its convolution in the space domain. And so what we'll do is we'll use the FFT fast forward transform to essentially convert the input kernel to the frequency domain, do an element wise multiplication, and then bring it back to the space domain. And that whole operation allows us to reduce time complexity from an n squared to an n log n operation, which is for long sequences near linear.

7:06here. And do you retain the benefit of being able to accelerate using GPUs? Yeah, that's a good question. So in general, GPUs are a lot more optimized for matrix multiplication. And so there's probably some room for improvement in terms of fully utilizing the hardware. But in general, we can get, you know, fairly efficient convolutions using GPUs. And that's what we do for language model training for Haina. Now, there's follow-on work from Dan Fu, which I also helped out on, called Flash FFTConf, which focuses on trying to bring the... Flash attention plus this FFT work? Exactly. Bringing the sophistication from hardware-aware algorithm design from Flash attention to the FFT.

7:51And so that was really exciting work as well. And going from state-space modeling to the long-sequence convolution, Do you lose some explainability or mechanistic understanding in doing that? Having a states-based model that describes explicitly how the state evolves from one state to the next seems like it would have some of those benefits. Is that true at all? And do you kind of abstract away from some of that with a convolutional approach? Yeah, that's a good question. I haven't focused too much on the explainability of what these state-space models are learning in terms of their state-space representations.

8:34So it could be, but I think what we've focused on more so is using sort of attention matrix-like view of the convolutions. so even though we use convolutions we can also still produce an attention map you know quote unquote which allows us to essentially see what inputs are being attended to but but more so in convolutional like what's lit up right in terms of attributing weights essentially so that's not been a constraint so far for interpretability talk about how far you've pushed the hyena model in terms of input size or context length? Yeah. So in the initial work, we started off with 64 ,000 tokens for the hyena hierarchy paper on synthetic tasks.

9:21And then in a follow-on work, we applied it to DNA sequences in that it's in constant need of having long-range dependency modeling. When we created a model called Hyena DNA, which we scaled up to 1 million token context, which at the time was the longest language model context since then. Lots of folks have caught on and found this kind of work also interesting and also get Transformers to work as well with Long Context. So interesting to see the work evolve or the space evolve. Can you talk a little bit about the landscape of Long Context's enablement on the Transformer side? What are the approaches that folks are taking there and what's working?

10:04Do you have a sense for what Gemini is doing to, you know, get at 1 million, 2 million context windows and how those approaches will fundamentally differ from swapping out transformers for convolutions? Yeah, so I haven't. So, I think I've heard some rumors about how Gemini might do it. But, yeah, I actually don't have a strong idea myself. there's you know talk of using things like ring intention or some kind of sliding window attention where it's attending to a local uh window and then having a hierarchy where these local windows can get aggregated into more meaningful uh groups so i yeah i'm not sure actually but um i think I suppose, yeah, I think there's a lot more that could be explored.

11:00I think that sort of my takeaway is that folks have focused a lot on the memory reduction of using attention for long context. But there's sort of no way to get around, or at least I'm aware of, of reducing the time complexity. It's still going to be an N by N type of operation or N squared operation. and so you can reduce the hardware requirement but then the amount of time is still going to be pretty key at least with the current setups I've been seeing and so I think that's sort of motivated folks to look at these alternative architectures if they do care about long context including Mamba, obviously very popular and we're sort of placing our bets on a slightly different approach with hyena variants but in general I think quality of these long context model with a transformer or alternative architecture seems to be something that people are still trying to improve like you can fit it in context but does it you know generate high quality sequences throughout the long context for example still an open question of something we're observing and is hyena comparable to transformers from a memory utilization perspective?

12:19Or, well, I guess the implication is that it's much more efficient.

12:25Yeah, so in terms of, you know, if we compare to vanilla transformers, vanilla attention algorithm, yeah, we're definitely in reduction in terms of memory and time. The, I guess, bar is raised when you look at flash attention, you know, a very clever, sophisticated algorithm that's essentially learned how to linearize the memory by doing it in blocks, right? And so, I think that is a hard bar. Like, attention is super fast. And so, I think there's still room for improvement to get convolutional-based models to that degree of sophistication. But in certain... Yeah. So, the sort of answer is depends on which settings you care about.

13:04There's certain very long sequences that will still take very long on flash attention or transformers, which could be faster on the convolutions. But the memory side, the full benefit is not there yet because they probably have some more room in terms of the systems optimization. And you alluded to this a moment ago. In training Hyena, the training corpus was synthetic language as opposed to a traditional training data set. Is that am I remembering that correctly? And what does that mean? So, in the original hierarchy paper, we trained on language for, I believe, 2K context on the pile. So, like a pre-training perplexity comparison with transformers.

13:50And that essentially was stacking up favorably for Hyena. And then to push specifically the long context, at the time, there wasn't a lot of great benchmarks for analyzing long context or, you know, very useful ones that we thought. And so, we came up with or used synthetic tasks that focused on a specific task, which is recall. So, being able to look up portions, sort of like a key value dictionary lookup, look up values earlier in the sequence. And then we use a simple synthetic task of key value scores, like letters and numbers, essentially, just to see if we can do this simple task, but over long range.

14:31And so that sort of was a prototype for even fitting in lung context, but then that really motivated to look for useful applications and that kind of segue to the hyena DNA work. Got it, got it. So let's jump into the hyena DNA work. How did you identify biology as an application for this problem? So, part of it was I was always looking for an impactful application in healthcare or biology when I started my PhD, sort of like the right opportunity or right technology to dig in deep. And so, I sat in a number of different labs on campus at Stanford who work at the intersection of biology and computer science.

15:14And long sequences and long-range dependencies kept on coming up in the world of DNA where people can model it just like a language, like natural language in the sense that you have letters, you have a vocabulary, you can feed them to language models. And then you can ideally learn all the interactions and grammar associated with that language to do useful tasks like prediction or perhaps generation as well. Can you talk about how the long-range dependencies kind of manifest from a biological perspective? And how does that differ from natural language? Sure. So, I think the really neat thing that makes DNA an interesting problem area or application to work on for machine learning is that it does physically have these long-range dependencies.

16:03So, in DNA, you have this sequence of letters, essentially. They're actually chemical compounds, but we can represent them as letters, the four bases, A, C, T, and G. And in general, when you're modeling DNA, there's processes that where what you care about are essentially things binding to DNA and kickstarting the process of creating proteins. And so, the process of having something bind to DNA can be influenced by what we call motifs that are present in the DNA that will essentially alter the probability of things binding or not. Now, the interesting thing for long-range dependencies is that some of these interactions can occur over millions of base pairs away from each other.

16:51And so, having a model that could essentially identify the presence of certain motifs and how they affect downstream binding of certain molecules can mean the difference between someone expressing some kind of disease or altering functions or their traits, essentially. And so it's really desirable to capture these long-range dependencies within a single model. So can you talk a little bit about how one might use a long-range genomic foundation model and how that differs from the way someone might use something like an AlphaFold and kind of compare and contrast those approaches? Sure. So AlphaFold, AlphaFold 3 in particular, very powerful, very successful model and has probably the most impact with machine learning and biology so far.

17:39Now, as powerful as it is, it is a supervised task, which is specifically designed to predict structure, right? So, given a sequence of amino acids, and now more recently including DNA, they want to predict the 3D shape of that sequence because it has a 3D shape analog in physical space. For AlphaFold, it's specifically a supervised task. And so, it can be difficult to modify it for use of different purposes. So the nice thing about language models and foundation models, they're meant to be this general learning representation machine that could be modified to be applied to a number of different downstream tasks.

18:17And that could be more difficult or perhaps more challenging for an outfold model that's restricted. And so some of those tasks for a biological foundation model that's language-based can include the prediction tasks, right? So, let's say you have a sequence, what's its function, affect gene expression or how much of a certain protein does this particular sequence expect it to make, all the way to the generative side as well. And so, these language models, right, for natural language can generate lots of text and answer prompts and questions, right? But for DNA or proteins even, what you can imagine is, and what they're currently being used for, is to design novel drugs or novel molecules, proteins, or systems that could be designed just from the sequence space with these models.

19:09got it can you talk a little bit about um how the model is trained the human genome has what 3.2 billion nucleotides which is quite a bit less than the uh the corpus size for a state-of-the-art language model nowadays um but the language is a lot simpler in terms of the number of tokens like what are what's the training process and what are the um elements that come into play Yeah, so you're right. So our initial work started off with training on the human genome, which is made up of 3 billion base pairs or nucleotides. And the vocabulary, you know, depends on how you want to tokenize. But for our case, we use the simple vocabulary of just four bases, or you can consider like the bite level.

20:00And that can make it challenging for a number of ways. But yeah, for focusing on our initial use case, we'll train by just taking in an arbitrarily long sequence from this human genome and then do next token prediction or next nucleotide prediction, sort of a causal style. But we'll do this essentially billions of times, just like natural language. And so, yeah, that is a challenge for, I guess, scale, right? So, you know, initially, this could be seen as like a first attempt or prototype to focus on the human genome. But more desirably, we're following work for Evo in particular, we grabbed millions of genomes and millions of species to try to train a much larger DNA foundation model across different species.

20:52In that case, the good thing about DNA, it's one of the most common or most widely available sources of biological data out there. And so there's definitely trillions of tokens also in this domain that we can reach to to train a foundation model. Let's maybe take a step back and have you introduce EVO. Is it simply this idea of taking the hyena DNA approach and architecture and incorporating additional species? Or did the architecture itself evolve to enable you to do that? So, Evo is sort of our next stage in training large DNA foundation models, taking some of the learnings we had from hyena DNA and scaling it much larger.

21:37So, hyena DNA was about a 7 million parameter model, which in language models are very small. And then Evo is a 7 billion parameter model. So, it's about a thousand times bigger. And to get that to scale effectively, we did have some architecture improvements in this work. So, a couple of key ones is that it's a hybrid model. So, we actually mix in, reintroduce some attention layers with these high-end convolutional layers. And then thinking there is sort of to get the best of both worlds where convolutions can get, you know, the efficiency gain and then attention has somewhat stronger abilities in terms of recall ability.

22:18And so we'll take advantage of those as well. And so we also have some modifications in terms of how we parameterize the convolutional kernels themselves and then change some of the MLP blocks as well. But the overall training style and core of this convolutional style language model from Hyena is in large part used in Evo as well. Okay. And you talked about EVO, the EVO dataset spanning many species. Do you retain the human species from hyena DNA and then incorporate the phage and prokaryotic and those types of sequences? Or is it only the latter? I'm trying to wrap my head around, does it make sense to have all of that in one model?

23:13Yeah, that's a really good question. So, we actually made the sort of strategic choice to focus on just prokaryotic and phage genomes. And so, which, you know, folks who are not familiar with that biology, essentially microbial life. And the thinking there was we're going to focus on these genomes that are somewhat shorter. They're shorter and they're simpler life forms. And their DNA, the language of their, the grammar of their DNA follows more clear rules than humans or mammals, for example. Humans and mammals are quite complex organisms as we can imagine. And so, having a large model trained on humans and mammals or what we call eukaryotes is a level of complexity that we're still working our way up to and actually in follow-on working for EVO2 is what we plan to do next.

24:05But we sort of wanted to focus on simple life forms first because sort of the thinking so far before we started working on EVO was we were unsure that a large DNA model could work in general. We weren't sure if some of the tasks that we were going to do or go after could actually, specifically, to be able to generate and design DNA was a key task. So we wanted to start with simpler organisms first. Dig into that task in a little bit more detail. In terms of generating and designing DNA, what are some examples of reasons why you might want to do that and ways that it might be applied? Sure. Sure.

24:43So, one of the initial questions before we started the EVO project was, is it possible to have a language model or foundation model design an entire genome from scratch? And the motivation for that was a lot of work previously had, you know, or a lot of research in general allows us to write DNA, read DNA, edit any sequence of DNA. But to be able to design DNA, we still didn't know enough about how DNA worked to be able to design. And so, we wanted to use language models to see if we can help design through a foundation model. And so, you know, why would you want to do that if you could? the thinking was if we could design an entire genome then we can perhaps control the function of these genomes as well and so some of the ideas were can we target cancer cells if we can design organisms to to have that function can we make better biofuels can we make antibiotics these are some of the the sort of the moonshots we thought that could be used for um you know manufacturing or climate or therapeutic purposes, for example.

25:51And now that transformer-based approaches are accommodating longer context lengths, could those be as easily applied to these problems? We've seen attempts to apply them to these types of problems. Is it simply, quote-unquote, a choice of kind of training corpus and going through the motions, or are there unique adaptations of this approach that lends itself to the biological use case? So that's a good question. We get this question a lot. And it depends. You know, I think people have different answers. But I certainly, you know, having trained these DNA models for the past year and a half, certainly have my biases and my opinions, which I do believe that convolutional-based models have this inherent inductive bias advantage over attention models.

26:46And I think some of these reasons is that DNA is not natural language, for one. Natural language is a very dense, rich type of sequence or type of data distribution. In DNA, the information is far more sparse and over long range. It's noisier because there's, you know, like the biologists, there's a lot of talk about junk DNA in your genome where people don't know what it does or it might have no function in your genome. And so what we observed is that these convolutional models, we hypothesize that they're better at filtering out this noise and being able to parse the signals over the long ranges better than attention.

27:27perhaps because convolutions from signal processing and they're built to filter out information really well and let's use that analogy but yeah there's other things about modeling dna that make it challenging like being able to tokenize at the single character level or the single nucleotide level in general we've seen transformers just do a number of different works that are trying to model language at the character level and it usually underperforms and they've not been able to use long sequences, for example. And so in particular, it's important to have this single character tokenization for DNA because a single character change could mean the difference between having a disease or not.

28:07And so you want that sensitivity and resolution at that level that transformers, I haven't seen, I don't think anyone's seen a successful large scale language model that's tokenized at the single character level yet. Okay. One of the properties that is characteristic of transformer-based models is hallucination, of course. Is there anything kind of inherent to the hyena models that yields different characteristics in terms of propensity to hallucinate? Ah, that's interesting. I don't think I've played around with natural language enough to be able to identify human interpretable hallucinations, I suppose.

28:55But yeah, I mean, in general, it should be susceptible, I imagine, to these hallucinations, just like transformers. I guess maybe another way to ask the question is, is hallucination a function of the architecture more so than the probabilistic nature of the approach in general or the training approach? If you have a take on whether approaches like hyena or the hyena approach is for some reason less susceptible to it. And I think the context is that hallucination in language can be bad. hallucination in DNA generation sounds scary. Yeah, fair enough. Yeah, I'd say that these systems, they're probabilistic machines.

29:40You're teaching them the probability of the next token. So because it's probabilistic, these probabilistic errors could accumulate and then hallucinations can abound. And I don't think that's going to change by the architecture or the layer that we swap in. It's still fundamentally a next token prediction with a probabilistic mechanism. But yeah, you're right. So in terms of biology, this could be scary as one way to think of it, but it also could be a source of variability, right? Or like a controlled evolution, which that's usually the analogy that folks in machine learning try to describe these generative models in biology.