Language Modeling With State Space Models with Dan Fu - #630

0:07All right, everyone, welcome to another episode of the Twimble. AI Podcast. I am your host, Sam Charrington, and today I'm joined by Dan Fu, a PhD student at Stanford. Before we get into today's conversation, be sure to take a moment to head over to Spotify, Apple Podcasts, or your listening platform of choice. And if you enjoy the show, please leave us a five-star rating and review. Dan, welcome to the podcast. Hi, thanks so much for having me. Excited to be on. I'm excited to have you on and looking forward to digging into our main topic, which is your ICLR spotlight paper, Hungry Hungry Hippos, towards language modeling with state space models.

0:46But before we dive into that topic, I'd love to have you share a little bit about your background and how you came to the field of machine learning. Yeah. So I'm a PhD student at Stanford. I've been in the PhD program for about five years now. I think I really got interested in machine learning during undergrad when some of these kind of machine learning models were just starting to kick on. I remember the ImageNet challenge the first time that a deep learning model was doing well, kind of doing image recognition. I was like, oh, that seems really interesting. Then when I got to Stanford, there were just so much exciting things going on.

1:23And obviously in the past few years, there's just been an explosion of ideas, of energy, and it's never been a better time to kind of be in this space and be doing machine learning. That's kind of how I got here. And it's been a wild ride, but I've been having fun. It's definitely been a wild ride and it's definitely an exciting time to be in the space. So good timing for sure. Tell us a little bit about your research interests. Yeah. So for the past few years, I've really been interested in seeing how can you increase the context length of these amazing models that we've been seeing. Language models in particular.

1:58Exactly. People language models. If you use chat GPT or GPT-3 or GPT-4, they can do amazing things. You can talk to them. Sometimes I have them write my emails. But one of the things that you really wish is you really wish they could kind of have more context about you, about your life. Like I wish it could read all my emails and then write exactly an email the way I would with my exclamation points and my commas and all those other things. Right now, there's really kind of some fundamental limitations to the context length. For the longest time, you could only do 2000 words, which is, you know, a couple pages, I'm sure, you know, this conversation will generate more text than that.

2:34When you read a book, there's a lot more in your mind than just two pages. So for the past few years, we've just been looking at how do you increase the context length. So that's everything from kind of optimizing attention, which is one of the building blocks to all the models these days. And we've seen, you know, our work be picked up by industry, OpenAI has used it to increase the context length of some of their recent models to some of the newer work and the work that I think we'll be talking about today, where we're just trying to see, is there another primitive, another building block that we could use to build these language models that where you can still kind of get these magical qualities, but then that doesn't scale quadratically in the sequence length so that it could actually be, you know, computationally feasible to show it all my emails and it say something, you know, meaningful about me.

3:19I was going to ask you to elaborate on the challenge a little bit. You touched on it just then. There's maybe a kind of naive question. Hey, can't we just throw compute at this thing? It's cheap. Why are we stuck at 4K or whatever the context length is? What are the challenges with increasing that? Yeah. So really the fundamental challenge is that attention, which is kind of this building block that we've used to build all these large language models, it grows quadratically in the sequence length. And there's really a pretty simple reason for that. So attention, it's kind of like a brute force approach to language.

3:57So the way it would work is you have a sentence, and then in order to figure out what a word means, you compare that word to every other word in the sentence. So in order for a GPT model to process this conversation, for example, you'd have to get every single word that you and I have said. And then in order to interpret the next word, go and compare every single word to everything that we've said. And obviously, this is not how humans do conversation. Like I can barely remember which words we started this podcast with. Maybe Sam, you know better because it sounds like you've given that spiel a bit.

4:29But in language, we don't need to remember all these things. That's kind of what gives us hope that there's something that is subquadratic out there that can do language. And just building on that, do you take inspiration from the biologic approaches in trying to come up with alternative methods? Just hearing you say that makes me think that maybe there's some meta attention mechanism that's like more abstract than attention or something like, but are you thinking about it from the perspective of like purely the optimization of computational architectures and structures? Or are you also thinking about it?

5:07Well, the brain is doing something, what's the brain doing? And how can we do something more like that? Right. So as an academic, I like to say we can play both sides. About a year ago, we really tried this computational optimization bit, and you can go pretty far. So I mentioned at some point that I think that for the longest time, the sequence things were at 2000. But now in the biggest models that are 8000 or maybe 32 ,000, partially because we were able to kind of some systems optimizations to make it so that you can at least fit that into memory. so for 32 000 for example if you're willing to pay a very very large cost you can kind of make that happen a couple years ago it would have just been completely infeasible but now there's kind of some ways to do it so that's kind of one approach kind of the systems approach but like you say there's a lot of things to take inspiration from so in hungry hungry hippos for example we took inspiration from a lot of kind of foundational work in signal processing so folks who have been analyzing time series for decades.

6:05So obviously, if you're trying to analyze an earthquake time series, for example, you're not going to something sampled at thousands of samples per second, you're not going to do this quadratic thing. So in Hungry Hungry Hippos, we kind of took some inspiration from those techniques. And a lot of it is kind of built on those fundamentals. I'm no biologist, I won't necessarily say there's direct biological connections, but certainly inspiration from great work that people have been doing in the sciences for years. Okay. Before we get to Hungry Hungry Hippos, you mentioned that we've been able to kind of brute force it using computational optimization to get to 32.

6:42Can you elaborate a little bit on some of what that looks like? What are the things that we've done that will get us there or that have gotten us there? Yeah. So there's one work that we put out almost a year ago now called Flash Attention. It was a great collaboration with some really smart folks at Stanford, there's a couple of pieces, but I think the most important one for going to long context was we made it possible to, even though the compute is still quadratic, you're still doing all this brute forcing, previous implementations, the memory requirement was also quadratic. So if you train models and you kind of like run your NVIDIA SMI or whatever, classically, you increase the sequence length, then those numbers blow up, you've run out of space on GPU and then, yeah.

7:23I was going to ask, GPU memory as opposed to CPU memory or storage. Exactly. And with Flash Attention, we were able to take that memory footprint and turn it linear, kind of using some classical systems ideas that have been around like the databases community for years. But we were able to adapt some of those techniques and make the memory footprint linear. So you can fit a sequence like 32K, 32 ,000 onto your GPU now. You're still going to pay that high compute cost, but at least it's feasible to take some existing model that you did. Most of your training, maybe on 2K, something cheaper, and then fine tune it for a longer sequence.

7:58So things become a little more feasible. But of course, if you want to go from 32 to 64, you're going to pay a four times cost there. It certainly moved where the state of the art was, but these fundamental limitations to the quadratic brute force approach. And you mentioned that that was kind of applying classical database optimization types of approaches. Things like hashing and binary tree or trees or things like that or other approaches? What specifically did you do? So in kind of the machine learning parlance, we call it tiling. It's actually even simpler than hashing. So this is, I think in a databases course, it's almost, I think we found it in lecture one.

8:39Like sharding? Not even sharding, just you're computing some like matrix multiplication. And then instead of materializing it, you compute it kind of block by block. And in databases, this is kind of classically, you're searching through some set of records, or you're doing a join between two tables. And then the naive way to do it is you write down like all the records that match, but instead, you can just do a block by block. So it's a lecture one type idea in databases, you know, not even sharding and hashing and all those things, but hadn't been applied to attention yet. But it's kind of inspiring, I guess, that we can take these 101, lecture one type ideas and apply them and make a lot of progress, right?

9:20Yeah, I agree. There's certainly a lot of lecture one ideas out there. Hopefully there are a bunch that are as fruitful. Yeah, I think people say that science happens or progress happens when you kind of take ideas from two fields and put them together. And I think this is a great example of the right person in the database community would have known this, but needed to be in the right place and understand kind of what the machine learning bottleneck is to kind of make it happen. And so Hungry Hungry Hippos, is it equally as simple an application of basic principles? Or is it we getting more sophisticated here?

9:55I think it's slightly more sophisticated to the average computer scientist. And I say computer scientist intentionally because I found out recently that some of the primitives that we're building on are actually in first or second year electrical engineering. It's actually another setting where the right person in the right place could have figured the ideas out. So I'll get into a little bit more details. But what we're building on in Hungry Hungry Hippos is a primitive call from signal processing called state space models. So state space models or state space representations. These are classically used to kind of model long time series.

10:30And they're the type of thing that if you take a signal processing class, actually are kind of lecture one of a signal processing class, have not been applied to language as much in the past. Like Markov chains, that kind of thing? A little bit. You can think about it kind of like some of those recurrent networks, like Markov chains. They have some properties of things like LSTMs as well, if you're familiar with those. But they have a nice mathematical property from the signal processing folks, where you can also think about it as a convolution. So it's one mathematical primitive that is both kind of recurrent and also a convolution.

11:04And that kind of lets us do a lot of nice things. So with LSTMs or recurrent RNNs, classically, the problem is that there's like these vanishing gradients kind of powering, like running this thing over and over. it can be hard to optimize, hard to learn. A states-based model, you can take a recurrent network and then turn it into a convolution and then compute it kind of with, in parallel with relatively high hardware utilization and it actually becomes more feasible to use to train a large model. Okay, and one more kind of high level kind of overall problem question, is the bottleneck, the challenge, is it primarily train time?

11:43Is it inference time? Is it both? what's been the big limitation? So for our work, it was actually a quality problem. So if you take kind of a standard state space model and just stick it into language, it doesn't perform super well. So for our work, it was actually mostly a problem of figuring out how can we adapt this kind of mathematical primitive and make it work for language. But even before we get to the state space solution that you've done in this paper, the general issue of quadratic relationship between the sequence length and compute and memory, does that become a limitation at train time primarily or inference time or something else?

12:25So it's actually both. So both train and inference. If you're training a model that scales quadratically in sequence length, in order to get it to kind of learn the long-term dependencies that you'd want from long sequence, you have to train it with very long sequences. And then that's a quadratic blow up. But it's also a problem during inference time. And what I mean by that is, if you take one of these models that has a longer context length and you feed it a very long document, you can actually watch the model slow down as it's trying to produce new tokens. And that's because for attention, every time you produce a new token, every time you say something new, you have to go back and compare it to everything that you said before, every single word.

13:04So the longer your document gets, the kind of the more expensive the inference is. So with these quadratic things, it's actually both. It's both training and inference kind of become the bottleneck for these models. Got it. So let's dig into the next level of detail on Hungry Hungry Hippos in this state-space model. So the inspiration kind of came from this second year EE or signal processing context. So you had this idea, right? What was the first step and where did you run into challenges? Right. So the first step is really, okay, we have this primitive that we think should work. You know, it can model time series well, it can model audio well.

13:43And I've had great landmates who have been kind of looking at states-based models for years. Albert Gu Tredow, who's a co-author of the H3 paper. Like we have people who have been using states-based models for deep learning for a few years now and been seeing success. But we always ran into this problem where if you just took a states-based model and swapped out attention in Transformer. So take your standard Transformer, swap out attention for a states-based model. The quality is kind of not up to snuff. To give some specifics, we all see gaps of five perplexity points. And for some context, that's the difference between like a 100 million parameter model and a 10 billion parameter model.

14:23So it's a big gap in quality and not something that we'd like to see. So there's nice computational properties, but it didn't quite give us the quality and the magic that we want out of these language models. And so something you said in there, the state space model is not supplementing the existing transformer architecture, but rather you're trying to rip out attention and plug the state space model in its place. And the idea being that you're hoping that the state space model accomplishes the same thing, like capturing these long term relationships between tokens, but more efficiently. Exactly.

15:02Yeah. So we want something that doesn't scale quadratically and can kind of still do these nice language tests, still model language well. Yeah. And so how did you overcome this quality gap that you ran into? Yeah. So what we did is we kind of went back to basics and tried to figure out what is a simple little toy example that we can play with to kind of understand what's happening here. And what we ended up doing is we kind of started putting together these very simple synthetic languages. So very simple languages like with maybe 20 words in them and sentences like A, 2, B, 3, kind of like there's these letters and these numbers and you try to each sentence will have some mapping from letters to numbers.

15:46We call this particular language associative recall. The test was simple. It was just, okay, each sentence has this very simple mapping of letters to numbers. I can write a little piece of Python code to solve it. Can a states-based model solve it? And what we found is that transformers could do this fine. They would get to 100 % accuracy kind of every single time and things would work out really well. But when we swapped out the attention for the state space model, we found out that it couldn't solve these synthetic languages, which is kind of surprising because it's a very simple language. You would think that this powerful primitive can solve it.

16:20What we ended up finding was that there's kind of some fundamental limitations in the exact expression of the state space model. and in order to solve it, the fix was actually quite simple. We just took two state space models and stacked them together and then multiplied their outputs together. It was very simple, but with kind of taking two of them and stacking them together, we were able to figure out, oh, can I actually start doing this lookup? It can say at the end, you have this letter A that appears and you can go back and kind of look up what the right value is. Like if there was an A and a three that appeared earlier and then there's an A at the end, you can go back and do that lookup again.

16:55And when we started saying that, that O can actually turn on these synthetic languages, we then scaled it back up to kind of real language and find that it started to close the gap. So is there a specific intuition for why this stacking works? Like, is it, I guess to me, it maybe starts to seem like some kind of sliding window type of approach where you've got a buffer that you're able to kind of compare tokens? Yeah, so there's kind of two pieces in there. So the comparing tokens, your intuition is exactly right. That's kind of what those multiplicative elements are doing. And we can actually prove that standard state space model wouldn't capture some of these multiplicative elements.

17:36But there's another thing going on as well, which is with a state space model, we can create kind of a store for kind of global memory for the entire sequence. So the state in a state space model refers to kind of some hidden state that you can store. And with that, you can actually start to store some of these key value associations that you see and then recall them for the entire sequence. So previously a states-based model on its own, it can do something simple like, I will memorize the token that appeared at position two, but you can't kind of dynamically tell it where to look, what to look back up.

18:09And that's where the multiplicative interactions come in. So the multiplications, you can think of them as controlling like some sort of direction, some sort of like operator saying where to look. And then the states-based models with all this great work, all this great history and signal processing, You can actually kind of store a memory for a long sequence. And I don't know if this question makes any sense, but do you need the two state space models or is there some way to make it like self-referential, autoregressive or something like that with just one? It's something that we've looked into a little bit, but we haven't kind of fully explored.

18:45Actually, now that you mentioned in some follow-up work, we were able to kind of simplify the architecture a little bit. And if you read the Hungry Hungry Hippos papers, there's a lot of kind of stuff about exactly how we set up each state space model. But in some follow-up works, we're able to simplify it. And in one of those follow-ups, we actually kind of introduced some idea of a recurrence where we say one state space model and can just think of it as, okay, now it's going to multiply with the next one and kind of like that. There's a little bit of that going on. I'm sure now that it's out there, I'm sure there's a lot of very smart people who can take a look and see how they'd like to simplify it as well.

19:23What's the relationship between the number of states in the state space model and the sequence length? Is that a hyperparameter or is it one to one or what? We actually haven't seen that you need like a particularly large state for longer sequences. Maybe if we scale to hundreds of thousands or millions, the tune will change. But so far in what we've seen, the state vector size we use is 64. And that's the one that seems to be working so far. And so talk a little bit about results and benchmarking. You know, what context sizes did you look at? Sounds like you use, was it perplexity that you mentioned as the metric?

20:04Talk a little bit about what you saw. Right. So in the H3 paper proper, we kind of looked at language modeling at several different model sizes. So in the original paper, the one that we're talking about, we actually ended up keeping two attention layers in there because we saw some better performance. Kind of in follow-ups, we relaxed that to took out the last bits of attention. In the actual Hungry Hungry. And we're talking about the state space attention as opposed to the traditional transformer attention. Yeah, so in the HIPPO's paper, we have 30 layers in a transformer, rip out all the attention, and then replace it with the state, with the new H3 layer, except for like two attention layers.

20:43So one at the beginning, one in the middle that would kind of keep in. And we saw that there was better performance. There's like maybe a minor difference if you take the attention layers out. But the ones that we kind of evaluated at length were kind of this hybrid architecture. OK, so with this hybrid architecture, we found that we could actually outperform transformers in quality, which was kind of really exciting to us that we could rip out almost all the attention and get something that actually worked better. So some perplexity results for those who are listening. Perplexity, you can think of it as exponent of the loss.

21:13So it's kind of a measure of how well does your model learn how to model language. We saw that in terms of perplexity, the H3 hybrid models were outperforming transformers. We saw in some downstream evaluations, kind of zero shot performance on kind of NLP benchmarks like Superglue. We were seeing that they were outperforming transformers. That was really exciting to us. Anecdotally, we've heard people say that the models that we've trained seem to be more coherence a little bit than other models of the same size. So we've scaled up to 3 billion parameters so far, which is small relative to the open AIs of the world, but pretty decent for an academic setting.

21:52And when you say transformers, did you create a language model from a generic transformer architecture? Or are we talking about some kind of off-the-shelf named language model llama or whatever yeah so we took kind of a generic architecture and we trained our models from scratch so it's very similar to luther has a line of models called pythia gpt neo it was the same data set as that so they're kind of very comparable to that line of models okay so presumably then there wasn't a big focus on kind of optimization and tricks and things like that to you weren't necessarily going for kind of absolute performance relative to other language models.

22:42It was more base transformer than these various modifications and comparing the performance within that set. Is that right? Or did you also benchmark relative to other language models? We actually did benchmark relative to other language models. We compared against the original GPT-2 models, the GPT-neo models of the same size. When we did the paper, like Pythia, Llama, these other models kind of weren't out yet. But yeah, so we did kind of do a head-to-head comparison against those other models. And on average, we seem to be doing better than them. Okay. And the areas that you perform better, are there qualifications on that better performance beyond the kind of the architectural things that we talked about?

23:24Are there a certain set of cases where it performs better and cases where it performs worse or anything like that? So what we found is that in the kind of cases that we tested, it seemed to work pretty well. Since when we did this paper and now there's been a lot more kind of know-how about how to train these models in terms of fine-tuning them, using instruction tuning and doing things like that. So I think if you compare it to maybe a instruction tuned model of the same size, or one of these models kind of fine tuned for chat, the H3 model won't chat with you kind of as well, you kind of have to fiddle with it prompted in the right way.

24:02And I think the ways that we know how to train things now have evolved quite a bit. I'm sure that those things are hopefully also transferable. Yeah. Thank you for picking up on that. I was trying to get to like transferability of your results to other types of models and kind of the general problem domain of large language models and context length. And it sounds like you do expect transferability. In talking about the results, I don't recall you mentioning specific context lengths. Right. Yeah. Can you elaborate on that? I would expect that would be a big part of what you accomplished is increasing the maximum context length without unacceptable degradation and performance.

24:43Is that right? Yeah. So we have trained our own models up to length 8K, kind of a full training run. In some kind of smaller synthetic examples, we've gone up to 32K, 128K, and some follow-ups. We're working on trying to scale the full, take those sort of more synthetic examples, bring it to the full language modeling experience. I think, you know, hopefully within the next month or two, we'll try to release something that where the sequence length is much, much longer. We're actively working on it and, you know, very excited to kind of see what the models can do. Using the same state space approach or are there other approaches that you have moved on to?

25:22So one follow up that I'll mention briefly is called Hyena, where we. Just kind of going across the zoo. Exactly. We're looking for the next animal that starts with an H for the next version. So with hyena, I mentioned, I think briefly, that state space models, you can think about them as maybe a occurrence, maybe a convolution. In hyena, we said, okay, what if we just go with a convolution view completely? And we saw a little bit, we can eke out a little bit more performance. We don't need those last couple of tension layers in there with hyena. So it's further improvements. I think with Hyena, we're also trying to see if we can get kind of some of those recurrent properties back in there, all kind of orbiting around the same core ideas, the flavor of the week, the differences in architecture.