2024 in Post-Transformer Architectures: State Space Models, RWKV [Latent Space LIVE! @ NeurIPS 2024]

Yeah. So thanks so much for having us. Uh, so this is gonna be a little bit of a two-part presentation. Uh, my name is Dan. Um, I'm at Together AI, and I'll be joining UCSD as faculty in about a year.

Um, and Eugene, you wanna introduce yourself?

Yeah. I'm Eugene. I lead the Adaptivity team, and I- I'm CEO and co-founder of Featherless, and we both work on this new post-transformer architecture space.

Yeah. So, um, yeah, so today, uh, we're really excited to talk to you a little bit about that. Um, so first I'm gonna give a broad overview of kind of the last few years of progress in non-post-transformer architectures, and then afterwards, uh, Eugene will tell us a little bit about the latest and the greatest, and the latest frontier models, uh, in this space.

So the story starts with scaling. So this is probably a figure or something like this that you've seen, um, very recently. Over the last five to six years, we've seen models really scale up in parameter size, and that's brought with it a bunch of new capabilities, like the ability to talk to you and tell you, um, sometimes how to use your Colab and your AWS screens.

Um, but another place where we've seen scaling, uh, especially recently, is scaling in context length. Um, so this can mean just having more text inputs for, for your models, but it can also mean things like, uh, taking a lot of visual token inputs, uh, image inputs to your models, or generating, um, lots of outputs.

And one thing that's been really exciting over the last few months or so is that we're, we're seeing scaling not only during training time, but also during test time. So this is one of the, the... this is the iconic image from the OpenAI o1 release.

Not only are we starting to scale train time compute, but we're also starting to scale test time compute. Now, if you're familiar with our to- with our attention and our transformer architectures today, this graph on the right might look a little bit scary, and one of the reasons is that, uh, the implications are, are a little bit, um, uh, interesting.

So what does it mean if we want to continue having smarter and smarter models? Do we just need to start building bigger, bigger data centers, um, spending more FLOPS? Uh, is this, uh, the, this little DALL-E 3, "We need more FLOPS," guy, is this gonna be the future of all of AI, or is there a better way, um, another path forward?

Maybe we can get the same capabilities that we've g- gotten used to, but for a lot less compute, a lot less FLOPS. And one of the things that we're gonna talk about today is specifically looking at that core attention operator, uh, in some of these models.

And the reason is that, um, so th- this is just some, some basic, uh, you know, scaling curves, but attention has compute that scales quadratically in the context length. So that means that if you're doing something like test time compute, and you want to spend a bunch of tokens thinking about what comes next, the longer that that goes, um, the, the, the lo- the more tokens you spend on that, that compute grows quadratically in that.

One of the questions that we're interested in, in is, can we take that basic sequence model, that basic sequence primitive at the bottom, and get it to scale better? Can we scale in, let's say, N to the three halves or N log N?

And so in, in the first part of the talk, so we just went over the introduction. What I'm gonna do over the next few slides is just talk about some of the key advances and ideas that have, uh, shown over the past few years since, uh, maybe early 2020 to, to now, that shown promise that this might actually be possible, that you can actually get potentially the same quality that we want, um, while scale- while scaling better.

So to do that, we're, uh... and, and basically the, the story that we're gonna look is we're gonna start to see how, um, so this is a basic graph of ju- just the past couple years of progress of perplexity, where that blue line, that dotted blue line is attention.

It's your basic transformer, full dense attention, and then the dots coming down are some of the methods that you'll see in this presentation today. We're gonna turn the clock back all the way to 2020. So this, this, this question of can we make attention subquadratic?

Basically, as soon as we said attention is all you need, people started asking this question. So we have this at- attent- quadratic attention operator. Can we do better? I'll briefly talk about why attention is quadratic, and, and the basic thing that happens, um, if you're not familiar, is that you have these inputs, these keys and queries, and what you do in this attention matrix, this S matrix over here, is that you're using-- you're comparing every token in your input to every other token.

So when I try to do something like upload a whole book to Gemini, what happens beyond the... or maybe not Gemini, 'cause we don't necessarily know what architecture it is, but let's say we upload it to Llama. What happens beyond the scenes, behind the scenes, is that it's gonna take every single word in that book and compare it to every other word.

And this has been a really-- it's, it's led to some pretty impressive things, um, but it's kind of a brute forcing of the way that you would try to interpret a, um, interpret something. Uh, what attention does in particular is the, uh, and then what, uh, attention...

Sorry, don't wanna... Okay. No, no laser pointer. What, what attention does afterwards is that, uh, instead of always operating in this quadratic thing, it takes a row-wise softmax over this matrix and then multiplies it by this values matrix.

So one of the key points to notice is that the output size is always gonna be the same as the inputs, at least in standard, uh, self-attention. So one of the first things that folks tried to do around 2020 is this thing called linear attention, which is just, just noticing that if we take out this softmax from here, if we take out this nonlinearity in the middle of the attention operation, and then if you compute the keys and the values operation first, you actually never hit this quadratic bottleneck.

Um, so that- That's potentially a way to get a lot more, uh, computationally, um, efficient. And there are various ways to do this by basically using feature maps or try to approximate this overall attention computation. But some of this work, uh, sort of started to hit a wall in 2020, and the basic challenges were, were two.

So one was quality. Um, it w- back then, it was kind of hard to, to get good quality with these linear attention operators. The other one was actually hardware efficiency. So these, this feature map that's just shown by a simple fi- uh, simple fi here actually ends up being quite computationally expensive if you just implement it naively.

So you started having these operators that not only were you sure, you're, you're, you're not really sure if they have the same quality, but also they're actually just wall clock slower. So you kind of end up getting the worst of both worlds.

So this was the, uh, the stage. So that kind of sets the stage for four years ago. Keep this in mind because linear attention is actually gonna come back, um, in a few years once we have a better understanding.

But one of the works that started kicking off this, this mini revolution in post-transformer architectures, uh, was this idea called state-space model. So here the seminal work is, is one by, uh, Werku in 2022. Um, and th-this, the, this piece of work really brought together a few ideas from, from some long-running, um, research, uh, research, um, lines of work.

The first one was, and, and this is really one of the keys to, to closing the gap in quality, um, was just using, uh, things that, that if you talk to a, a, a, an electrical engineer off the street, they might know, um, off, off the, like, the back of their hand.

But taking some of those properties with how we model dynamical systems in signal processing and then using those ideas to model the inputs, the, the text tokens in, for example, uh, a transformer like next token prediction architecture. So some of those early state-space model papers were looking at this relative, relatively simple recurrent update, uh, model that comes from maybe chapter one of a signal processing class.

Um, but then using some, uh, principle theory about how you should do that recurrent update, uh, in order to really get the most that you can out of your hidden state, out of your, uh, out of your sequence.

So that, that was one key idea for quality, and when, when this first, when, when this was eventually realized, you started to see a bunch of benchmarks that, that were pretty sticky for a few years, things like long range arena, some long, long sequence evaluation benchmarks.

Uh, there was stuff in time series, um, uh, time series analysis. Uh, they, they started to, to, you started to see the quality tick up in meaningful ways. But the other key thing that, uh, uh, was so influential about these state-space models is that they also had a key idea about how you can compute these, uh, things re- uh, efficiently.

So if you go back to your machine learning 101 class where you learned about RNNs, one thing that you may have learned is that they don't paralyze as well as attention, because if you just run them naively, you have to do this, uh, kind of sequential update to, to process new tokens, whereas in attention, you can process all the tokens in parallel at one time.

One of the key insights behind the S4 paper was that these recurrent models, you could take them, and you could also formulate them as a convolution. And in particular with the convolution, you could, instead of using a PyTorch conv1d operation, you can compute that with the FFT.

And that would give you N log N compute in the, uh, in the sequence length N with a operator that was relatively well-optimized for modern hardware. So those are really, I'd say, the two key ideas, um, in 2022 that started allowing these breakthroughs to happen in these non-transformer architectures.

So these ideas about how to principally model, um, sorry, how to model the recurrent updates of a mo- uh, of, of a sequence in a principled way, and also these key ideas in how you can compute it efficiently by turning it into a convolution and then scaling it up with the FFT.

Uh, along those same lines, so, uh, afterwards, uh, we started putting out some work on specialized kernels. So just like we have FlashAttention for transformers, we also have works like FlashFFTConv. And if you look at these lines of work, um, oftentimes when, whenever you see a new architecture, you see a new primitive, um, one of the, one of the table stakes now is do you have an efficient kernel so that you can actually get wall clock speed up?

So by 2022, 2023, we were starting to have these models that had promising quality, um, primitives, but, and, and also promising wall clock. So you could actually see regimes where they were better than transformers in meaningful ways. That being said, there were, there was still sometimes a quality gap, particularly for language modeling.

And because language is, is so core to what we do in sequence modeling these days, um, uh, the, the next, the next key idea that I'm gonna talk about is this idea of selection mechanisms. And this is basically an idea of, so you have this recurrent state that you're keeping around that just summarizes everything that, that came before.

And to get a good sequence model, one of the things that you really need to be able to do is have the model learn what's the best way to pick out pieces from that recurrent state. So one of the, one of the major ideas here in, in a line of work called H3, Hungry, Hungry Hippos, and also these hyena models, um, were, uh, one way you can do this is by just adding some simple element-wise gates.

Um, so s- versions of these ideas have been around for decades. If you squint at the LSTM paper, um, you, you can probably find, find this gating mechanism. But turns out you can take those old ideas, add them into these new state-space models, and then you can see quality start to, start to pick up.

Um, if you've heard of the Mamba model, this, this also takes this selection to, to the next level by actually making some changes in that, uh, fundamental recurrent state space. So it's not only just this gating that happens around the SSM layer, but also, uh, you can actually make the A, B, C, D matrices of your state-space model, you can make them data-dependent, which will allow you to even better select out different pieces from your hidden state, depending on what you're seeing.

I'll also point out, if you look at the bottom right of this figure, there's this little triangle with a GPU SRAM, G- GPU HBM, and this, this is just continuing with that trend of when you have a new architecture, um, you, you, you also release it with a kernel to, to, to show that it is hardware efficient and that, that it can be hardware efficient on modern hardware.

Um, the, the... One of the next cool things that happened is once we had this understanding of these are the basic pieces, these are the basic principles behind some of the sequence models, uh, linear attention actually started to come back.

So in, uh, earlier this year, there was a model called Based, uh, the, from Simran Arora and, and some other folks, that combined a more principled version of linear attention that basically, uh, the, the, the, the two-second summaries that had used a Taylor approximation of the softmax attention, combined that with a simple sliding window attention, um, and was, uh, starting table, starting to be able to expand the Pareto frontier of how much data can you recall from your sequence versus how small is your recurrent state size.

So those orange dots are, at the top there, are just showing, uh, smaller sequences that can recall more memory.

And the, the, the last major idea, I think, that, that has been influential in this line of work, and is very relatively late-breaking, just a few months ago, is just the basic idea that when you have these models that are fundamentally more effici- more efficient in the sequence length, uh, you maybe don't want to prompt them or use them in exactly the same way.

So this was a really cool paper called Just Read Twice, also from Simran, um, that, that basically said, "Hey, all these efficient models can process tokens so much more efficiently than transformers that they can sometimes have unfair advantages compared to a simple transformer token."

So... Or sorry, a simple transformer model. So take, for example, uh, the standard, the standard use case of you have some long document, you're gonna pass it in as input, and then you're ask- gonna ask some question about it.

One problem you might imagine for a recurrent model where you, where you have a fixed state size is, let's say that your, uh, article is very long and you're trying to ask about some really niche thing, um, you can imagine it might be hard for the model to know ahead of time what information to put into the hidden state.

But these m- these, these models are so much more efficient that you can do something really stupid, like you can just put the document, uh, write down the document, write down the question, write down the document again, and then write down the question again, and then this time, the second time that you go over that document, you know exactly what to look for.

Um, and the cool thing about this is, so this, uh, this results in better quality, especially on these recall-intensive tasks. But the other interesting thing is, uh, it really takes advantage of the more efficient architectures that, that we're having here.

So one of the other, I think, influential ideas in this line of work is if you change the fundamental compute capabilities of your model and the way that it scales, you can actually start to query it at test time differently.

And this actually, of course, goes back to, to those slides on test time compute. So while everybody's looking at, say, test time compute for, uh, big transformer models, I think a potentially a really interesting research question is: how can you take those, and how does it change with this new next generation of models?

So the... I'll, I'll just briefly summarize what some of those key ideas were and then talk, and then show you briefly kind of what the state of the art is today. Um, so, so the four key ideas are instead of just doing a simple linear attention approximation, instead, uh, take ideas that we know from other fields, like signal processing, do a more principled approach to your modeling of the sequence.

Another key idea throughout all, all these lines of work is you really want hardware and kernel support from day one. Um, so, so even if your model is theoretically more efficient, if somebody goes and runs it and it's two times slower, uh, one of the things that, that we've learned is that if, if you're in that situation, it's, it's just gonna be dead on arrival.

Um, so you want to be designing your architectures with, with the hardware in mind. Um, one of the key, key machine learning ideas that has been important for the quality is just making sure that you encode different ways that you can select from your hidden state and, and really focus on that as a key decider of quality.

And finally, I think one of the, the, the emerging new, new things for, for this line of work, and something that's quite interesting, is what are the right test time paradigms for these models? How do they change relative to, uh, re- relative to what you might do for a standard transformer?

I'll briefly end this section. So I've dr- labeled this slide "where we are yesterday," um, because Eugene is gonna talk about some new models that, that he released literally this morning. Um, but as of yesterday, some of the really cool results out of, um, the, these efficient alternative models were, uh, so AI2 trained this hybrid MoE called Jamba, um, that, that, that seems, that is currently the state-of-the-art for these non-transformer architectures.

Um, there's this, uh, NVIDIA and MIT put out this new diffusion model called Sana recently, that one of their key, uh, key observations is that you can take a standard diffusion, uh, transformer diffusion model, replace the layers with linear attention, and then that lets you scale to much larger, uh, much larger images, much, much ba- uh, much larger sequences, um, more efficiently.

And, um, and one thing that I don't think anybody would've called when, uh, a, a few years ago is that, uh, one of those, uh, gated SSM, gated states-based models, um, ended up on the cover of Science because, uh, a great group of folks went and trained some DNA models.

So that's, uh, Michael Pali, Eric Yuan from, uh, from Stanford and the Arc Institute. Um, so it's, we're, we're really at an exciting time in 2024, where these non-transformer, post-transformer architectures are showing promise across a wide range, um, a- across a wide range of, of modalities, of applications, and, and of tasks.

And with that, I'll pass it on to Eugene, who can tell you a little bit about the latest and greatest with RWKV.

Yeah. So this is what? Yeah.

You talk into here.

Oh, I talk into here. Okay. So yeah, two streams. Uh, yeah. So I think one common questions that we tend to get, uh, asked, right, is what's the difference between RWKV and State Space? So, um, I think one of the key things to really understand, right, uh, the, the difference between the two groups, right, is that, um, we are actually more like an open source rando internet meets academia kind of situation.

Like, most of us never wrote any paper, but we, we basically look at RNNs and linear attention when attention is all unique came out. And then we decided to like, hey, um, there is a quadratic scaling problem. Why don't we try fixing that instead?

So, so, so we end up developing our own branch, but we end up sharing ideas back and forth. So and, and we do all this actively in Discord, GitHub, et cetera. This was so bad for a few years, right?

That, uh, basically the average group's H index was so close to zero, right? EleutherAI actually came in and helped us write our first paper. Great. Now our H index is now three, apparently. So, so, so but, but the thing is that a lot of these experiments led to results and, and essentially, essentially we, uh, we took the same ideas from linear attention, and we built on it.

So to take a step back into, like, how does RWKV handle its own attention mechanic, uh, and achieve the same goals of like ON compute respectively and, uh, in focus of our overall goal to make AI accessible to everyone regardless of language, nation or compute.

That's our open source goal. We actually train our models primarily on over a hundred language, which is another topic altogether. Uh, and our goal is to train to even two hundred languages to cover all languages in the world.

But at the same time, we work on this architecture to lower the compute cost so that people can run it on Raspberry Pis and on anything. So how did RWKV break the dependency of LSTM token flow? Because I think to understand our architecture, right, it's probably easier to understand it from the RNN lens, um, because that's where we built on.

Where, where State Space kind of like tried to, tried to start anew and took lessons from them instead. So there's a little bit of divergence there and aka this is our version of linear attention. So to take a step back, all foundation models, be it transformers or non-transformers at a very high level, right, pumps in a token, uh, I mean text that turns into embeddings and goes through a lot of layers, generate a lot of internal states, whether QKV cache or be it RNN states or RWKV states and outputs an embedding.

They are not the same thing, and we just stack more layers and more embeddings and somehow that magically works. So if you, if you remember your ancient RNN, uh, lessons, uh, which we, which we, which we, uh, we call less learning these days, um, the general idea is that you have the embedding information flowing all the way up and when-- a-and you take that information and you flow it back down, and then you process it as part of your LSTM layers.

So this is how it generally works. Uh, Karpathy is quoted saying that RNNs are actually unreasonably effective. The problem is this is not scalable. To start doing work on the second token, you need to wait for the first token, and then you need...

and likewise for the third token and fourth token, yada, yada. That is CPU land, not GPU land. So, so, so you can have a H100 and you can only use one percent of it. So, so that's kind of why RNNs didn't really take off, uh, in the direction that when we wanted like billions of parameter and becomes training.

So what did RWKV version zero do? We just did the dumbest, lamest thing. Uh, sorry, this is the, the bottleneck for RNN. We take the dumb thing of removing that line, and it kind of worked. It trained, it sucked, but it kind of worked.

Then, then we were like, hey, um, then no one cared because the loss was crap, but how do we improve that? And that's essentially where, where, where we move forward. Because if you see this kind of flow right, you can actually get the...

you, you can get your GPU saturated quickly, where it essentially cascades respectively. So I'm just waiting for this to loop again. So it's like once you get your first layer, your token to be, uh, computed finish, you start to cascade your compute all the way until you're, hey, I'm using hundred percent of GPU.

So, so we, we worked on it, and we started going along the principle of that as long as we keep this general architecture where, where we can cascade and, and be highly efficient with our architecture, nothing is sacred in our architecture.

And we have done some crazy ideas. In fact, you ask us-- if you ask me to explain some things in the paper, right? Officially in the paper, I'll say we had this idea and we wrote it this way.

The reality is someone came with the code, we tested it, it worked, and then we rationalized later. So, so the general idea behind RWKV is that we generally have two major blocks that we do. Uh, we call it Timemix and Channemix.

And Timemix generally handles, um, handles long-term memory state, uh, where essentially, where essentially, uh, where we apply the matrix multiplication and silo activation functions into processing an input embedding and an output embedding. Uh, I'm oversimplifying it because this, this calculation change every version.

Uh, and we have like version seven right now. Channemix is more-- is similar to Base in the sense that where it does shorter-term attention, where it just look at the s- the sister token or the token before it, because it, there's a shift in, in the token shift, uh, matrix.

I don't really want to go too much into the papers itself because, like, uh, we do have three papers on this. Um, basically, RWKV-RNN for the transformer era. Eagle and Finch, RWKV Matrix Value State. This is the updated version five, version six, and Goldfinch is our, is, is, is, is our hybrid, uh, uh, model respectively.

Uh, we are writing the paper already for V7 and, uh, which is, which is for RWKV7 codenamed Goose. All architectures are codenamed by a bird. And I'm going to cover as well Q-RWKV and MamaRWKV and RWKV MoE. So where did that lead to?

Okay, because we are all GPU poor, um, and to be clear, like, most of this research is done, like, only on a handful of H100s, which I had one Google researcher t-told me that was, like, his experiment budget for a single researcher.

So our entire organization has less compute than a single researcher in Google. Um, we, we... One of the things that we explored, uh, into was to how do we, uh, convert transformer models instead, because someone already paid that billion doll- uh, million dollars onto training, so why don't we take advantage of those, uh, weights?

And, and to- I believe Together AI worked on the Lolcat for, for the Mamba side of things and, and we took some ideas from there as well, and we essentially did that for RWKV. Um, and that led to, uh, Q-RWKV6, which we just dropped today, a 32B instruct preview model, where we took the Quen 32B instruct model, freeze the feed-forward layer, uh, uh, remove the QKV attention layer, and replace it with RWKV linear layers.

So, um, to be clear, this means we do not have the RWKV channel mix layer. We only have the time mix layer. But, uh, but once we do that, we, we train the RWKV layer. Important is that the feed-forward layer needs to be frozen so, so the new attention can be learned.

And then we unfreeze the feed-forward layer and train all the layers together with a custom learning rate schedule so that they can learn how to work together. The end result, surprisingly, and to be honest, to the frustration to- of the RWKV MWE team, which ended up releasing the model on the same day, was that with just a few hours of training on two nodes, we managed to get it to be on par kind of with the original Quen 32B model.

So in fact, when the first ran, right, that completely confused us was like... And I, and I was telling, uh, Daniel Goldstein, uh, uh, smoking food, who kind of leads most of our research coordination, "When you pitched me this idea, you told me at best it would get the same level of performance.

You didn't tell me the challenge and peak, uh, uh, score and peak, uh, we- I mean, the black score will shoot up. I don't know what is happening there." But it did. Um, MMLU score dropping, that was expected because if you think about it, when we were training all the layers, right, we were essentially, like, Frankensteining this thing, and we did brain damage to the feed-forward network layer to- with the new RWKV layers.

But 76%, hey, some of it is retained, and we can probably further train this. We, we have- we didn't even spend more than three days training this, so, so there's a, there's a lot more that can be done, hence the preview.

But this brings up a big question because we are already now in the process of converting the 70B. Uh, we are now... This is actually extremely compute efficient to test our attention mechanic. It's like it, it becomes a shortcut.

We can, uh, we are already planning to do our version seven and our hybrid architecture for it because we don't need to train from scratch, and we get a really good model out of it. And the other thing is that is, uh, uncomfortable to say is that because we are doing right now the 70B, is that if this scales correctly to 128K context length, I'm not even talking about a million, 128, majority of enterprise workload today is just on 70B at under 32K context length.

That means if this works and the benchmark, uh, matches it, it means we can replace the vast majority of current AI workload unless you want super long context. And then, uh, sorry, can someone give us more GPUs because we do need the VRAM for super long context, sadly.

So yeah, that's, that's what we, we are working on, and essentially, uh, we are excited about this to, to just push it further. And this conversion process, to be clear, I don't think it's going to be exclusive to RWKV, but, um, it probably will work for Mamba as well.

I don't see why not. And we'll probably see more ideas or more experiments or more hybrids or like, yeah. Like, one of the weirdest thing that, uh, I wanted to say out, right, and I confirmed this with the Black Mamba team and the Jamba team, which, um, because we did the Goldfinch hybrid model, is that none of us understand why a hy-hybrid with a state-based model, be it RK or state space, and transformer performs better when, than the baseline of both.

It's like, it's like when you train one, you expect... and then you replace, you expect the same results. That's our pitch. That's our claim. But somehow when we jam both together, it outperforms both. And that's, like, one area of abolition that, like, we only have four experiments across four teams that a lot more needs to be done.

But, but these are things that excite me essentially because that is what is potentially we can move ahead for. Which brings us to what comes next.

So th-this part is kind of just, uh, some... where we'll talk a little bit about stuff that, that we're excited about, um, maybe have some, uh, wild speculation on, on, um, what, what, what's, what's coming next. And of course, uh, this is also the part that I'll be more open to questions.

Um, so a couple things that, that I'm excited about is, uh, continued, um, hardware model co-design for, for these models. Um, so one of the things that we've put out recently is this library called Thunder kittens. It's a CUDA library.

Um, and one of the things that, that we found frustrating is every time that we built one of these new architectures, and I'm sure you had the, the exact same experience, we'd have to go and spend two months in CUDA land, like, writing these, these new efficient things.

And, um, if we decided to change one thing in PyTorch, like one line of PyTorch code is like a week of CUDA code at least. Um, so one of our goals with, with a library like Thunder Kitten, so we, we just broke down what are the key principles, what are the key hardware, um, things, uh, what are the key compute pieces that you get from the hardware.

So for example, on H100, um, everything is really re- revolves around a warp group, uh, matrix multiply operation. So you really want your operation to be able to split into a relatively small, uh, matrix, matrix multiply operation. So like multiplying two 64 by 64 matrices, for example.

Um, and so if you know that ahead of time when you're designing your model, that probably gives you, uh, you know, some information about how you set the state sizes, how you set the update, how you set the update function.

Um, so with Thunder Kittens, uh, we basically built a whole library just around this basic idea, um, that all, uh, your basic compute primitives should not be a float, but it should be a matrix, and everything should just be matrix compute.

Um, and we've been using that to, to try to both re-implement some existing architectures and also start to design some new ones that are really designed with this core, uh, with a tensor core primitive in mind. Um, another thing that, that we're, that at least I'm excited about is, um, we, over the last four or five years, we've really been looking at language models as the next thing.

But if you've been paying attention to Twitter, uh, there's been a bunch of new next generation models that are coming out. So, uh, there, there are video generation models that can run real time, so, that are supported by your mouse and your keyboard, um, that, uh, I- I'm told if you play with them, they, you know, they, they only have a few seconds of memory.

Can we take that model? Can we give it a very long context length so that you could actually maybe generate an entire game state at a time? What does that look like for the model? You're certainly not gonna do a giant quadratic attention computation to try to run that.

Um, maybe, maybe use some of these new models or some of these new video generation models that came out. So Sora came out, uh, I don't know, two days ago now, um, but with super long queue times and super long generation times.

So that's probably a quadratic attention operation at the, at the bottom of it. What if we could remove that and get the same quality but a lot faster generation time? Um, or some of the demos that, that we saw from Paige earlier today.

You know, if, if I have a super long, uh, conversation with, with my Gemini bot, um, what if I wanted to remember everything that it's seen in the last week? Um, I mean, maybe you don't for, for personal reasons, but, but what if I did, you know?

Um, what does that mean for the architecture? And I think, you know, that's certainly something I'm pretty excited about. Um, I'm sure you're excited about it too. Uh, I think we were supposed to have some hot takes, but I honestly don't remember what our hot, hot takes were.

Yeah. Including exciting.

Hot take.

Hot takes.

Yes, these are our hot takes. Um-

So I, I think the, the, the, the big, the big one on Twitter that we saw, uh, that we shared was the question is like, is RAG relevant in, in the case of, like, the future of, like, state-based models?

Let's see. I, I haven't played too much with RAG. Um, but when I have, I'll say I found it was a little bit challenging to do research on it, um, because we had this experience over and over again where you could have any, an embedding model of any quality.

So you could have a really, really bad embedding model, or you could have a really, really good one by an- by any measure of good. And for the final RAG application, it kind of didn't matter. Um, that's what I'll say about RAG while being recorded.

I know it doesn't actually answer the question, but-

Yeah. So, um, I think a lot of, uh, folks are, are, like, extremely excited of the idea of be it RWKV or state space potentially having infinite context. Um, but I think the reality is that when we say infinite context, we just mean a different kind of, uh, infinite context or y- or as, as previously covered, you need to test the model differently.

So think of it more along the lines of the human. Like, I don't remember what I eat for breakfast yesterday. Yeah, that's the statement that I'll say. And, and we humans are not quadratic transformers. Uh, if we did, if, let's say, we increased the, our brain size for every second we live, we would have exploded by the time we are five years old, something like that.

And, and I think, I think, uh, basically fundamentally for us, right, be it whether we... Regardless of whether RWKV, state space, XLSTM, et cetera, our general idea is that instead of that expanding state, that increase in computational cost, what if you have a fixed state size?

And information theory detects that that fixed state size will have a limit. Um, just how big of a limit is a question. Like, we, like RWKV is st- is running at forty megabytes for, for a state. Uh, its future version might run into four hundred megabytes.

That is like millions of tokens in, if you're talking about mathematically the maximum possibility. Um, it's just that I guess we are all more inefficient about it, so maybe we hit a hundred thousand. And that's kind of like the work we are doing, trying to, like, push it and maximize it.

And that's where the models will start differing, because it will choose to forget things, it will choose to remember things. And that's wh- that's why I think that there might be, uh, some element of RAG, but it may not be the same RAG.

It may be the model learn things, and it's like, "Hmm, I can't remember that, that article. Let me do a database search to search." Just like us humans, when we can't remember the article in the company, we do a search on Notion.

Yeah. I think something that would be really interesting is if you could have facts that are... So right now, the one intuition about language models is that all those parameters are around just to store random facts about the world.

And this intuition comes from the observation that if you take a really small language model, it can do things like talk to you or, uh, it kind of has like the, the style of conversation, it can learn that.

But where it will usually fall over compared to a much larger one is it'll just be a lot less factual about things that, that it knows or, or, or that it can do. Um, but that points to all those weights that we're, that we're spending, all that SGD that we're spending to, to train these models are just being used to store facts.

And we have things like databases that are pretty good at storing facts. So I think one thing that would be really interesting is if we could actually have some sort of outside data store that a language model can, can look at, um, that, that maybe is, uh, you know, has, has some sort of gradient descent in it, but, but would be, um, quite interesting.

And then maybe you could edit it, delete facts, you know, change who's president, um, so that it, it doesn't, uh, it doesn't get lost.

Can we open up Q&A and hot takes-

Sure

... to the audience?

I have hot take Q&A. Um, do these scale? When, when 405B state-space model RAG exists, no one does long context. Who's throwing in 2 million token questions? Hot takes.

The, the who's throwing in 2 million token question I think is, is a really good question. So, um, I actually, I was gonna offer that as a hot take. And my hot take was gonna be that long context doesn't matter.

I know I just gave a whole talk about it, but um, you know, what, what's the point of doing research if you can't, you know, play both sides? But I think one of the... So I, I think for both of us, the reason that we first got into this was just from the first principle questions of there's this quadratic thing.

Clearly intelligence doesn't need to be quadratic. What is going on? Can we understand it better? Um, you know, since then, uh, it's, it's kind of turned into a race, which has been exciting to watch, like how much context you can take in.

But I think it's right. Nobody is actually putting in a 2 million context prompt into these models. Um, and, and, you know, if they are, maybe we can go, go, you know, design a, a better model to, to do that particular thing.

Um, yeah, what do you think about that? So you've also been working on this. Do you think long context matters?

So I'm gonna burn a bit. Uh, how many of you, uh, remember the news of Google Gemini is supporting 3 million context length? Raise your hand. Yeah.

2 million.

Oh, it's 2 million.

Yeah. How many of you actually tried that? See?

I use it a lot.

You, you are from MindsDB.

I use it a lot.

All right. So m- for some people there is use. Uh, and I think, I think that's the, that's might be, like, this is where my opinion starts to differ because, like, I think the big labs may have a bigger role in this because, like, even for RWKV, even when we train long context, the reason why I say VRAM is a problem is that because when we did the, we need to back prop against the states, we actually need to maintain the state in between the tokens by the token length.

So that means we need to actually roll out the whole 1 million context if we are actually training 1 million, which is the same for transformers actually, but it just means we don't magically reuse the VRAM consumption in the training time space.

So that is the, one of the VRAM bottlenecks, and, uh, I neither OpenAI nor Google, so donate GPUs if you have too much of them. But, but th- but then putting it back to another paradigm, right, is that I think O1 style reasoning might be actually pushing that direction downward.

I- in my opinion, this is my partial hot take, uh, is that if, let's say, you have a super big 400B model, and let's say you have a 70B model that may take double the tokens but gets the same result.

Strictly speaking, a 70B, and this is even for transformer or non-transformer, right, w- will take s- less than, le- uh, less resources than that 400B m- model, even if it did double the amount of thinking. And if that's the case, and we are still all trying to figure this out, maybe the direction for us is really getting the sub 200B to be as fast, as efficient as possible with a very efficient architecture that some folks happen to be working on, uh, to, to just reasoning out over larger and larger context length.

Yeah.

One thing I'm super interested in is models that can watch forever. Um, obviously you cannot train something on infinite context length. How are y'all thinking about that, where you run on a much lon- longer context length than is possible to train on?

Yeah, it's a, it's a, it's a great question. So I, I think when, uh, y- I think y- you guys probably had, had tweets along these lines too. When, when we first started doing these things, because these are all recurrent models, um, in theory you could just run it forever.

You, you could just run it forever, and at the v- very least it won't, it won't, like, error out on you or, or crash. Um, there's another question of whether it can actually use what it's seen in, in that infinite context.

And I think there, so one place where probably the research and architectures ran faster than, than other research is actually the benchmarks for long context. So you turn it on forever, you want to do everything or watch everything.

What is it that you actually wanted to do? Can we actually build some benchmarks for that, then measure what's happening, and then ask the question, can the models do it? Is there something else that they need? Um, yeah, I, I think that, that if, if I were to turn back the clock to 2022, that's probably one of the things I would've done differently, which would've been actually get some long context benchmarks out at the same time as we started pushing context length on all these models.

I will also say the use case, so, like, I think we both agree that there's no infinite memory, uh, and the model needs to be able to learn and decide. I think what we have observed for, I think this also with the state-space model, is that one of the key advantage of, uh, this alternate attention mechanic that is not based on token position is that the model don't suddenly become crazy when you go past the 8K training context length or a million, uh, uh, context length.

Uh, it's actually still stable. It's still, it's able to run, it's still be able to rationalize. It just starts forgetting things. Um, but some of these things are still there in latent memory. Some of these things are still somewhat there.

That's the whole point of why reading twice works, uh, things like that. And one of the biggest push in this direction is that I think both state-space and RWKV have separate papers, uh, by other researchers where they use this architecture for time series data, weather modeling.

So you're not asking what was the weather five days ago. You're asking what's the weather tomorrow based on the infinite length that we, as long as this Earth and the computer will keep running. So, so, and they found that it is, um, like, better than existing, like, be it transformer or existing architecture in modeling this weather data.

Controlled for the param size and stuff. I'm quite sure there are people with larger models. So, so there, so there are things that, that in this case, right, there is future applications if your question is just what's next and not what's 10 years ago.

Thanks so much for having us.