#299 Jacob Buckman: Why the Future of AI Won't Be Built on Transformers

0:00Mamba and state-based models in general are what I would consider retention models. These models that have this sort of duality where they can be expressed either as recurrent models or as attention models. And once you have this, this model that can be expressed in both of these forms, then it also unlocks a secret third form of the model, which people call the chunked formulation, where you basically combine all the best parts of recurrence and attention to get an architecture that both has tractable growth in cost, It has linear growth in cost as you increase the context, but also is able to really saturate the GPU.

0:37I think that many tasks right now are better suited for in-context learning. Because current models have either very short context or, because of their architectures, a limited ability to actually use the thing in their context, it actually doesn't seem that most people are able to get enough value out of pure in-context learning. And so they have to essentially resort to fine tuning as a way of getting the information into the model. Build the future of multi-agent software with Agency. That's A-G-N-T-C-Y. Now an open source Linux foundation project, Agency is building the Internet of Agents, a collaborative layer where AI agents can discover, connect, and work across any framework.

1:28All the pieces engineers need to deploy multi-agent systems now belong to everyone who builds on agency, including robust identity and access management that ensures every agent is authenticated and trusted before interacting. Agency also provides open, standardized tools for agent discovery, seamless protocols for agent-to-agent communication, and modular components for scalable workflows. collaborate with developers from cisco dell technologies google cloud oracle red hat and more than 75 other supporting companies to build next generation ai infrastructure together agency is dropping code specs and services no strings attached visit agency.org to contribute that's A-G-N-T-C-Y dot O-R-G.

2:37Hey, I'm Jacob. I studied computer science at Carnegie Mellon and then went on to do a PhD focused on AI at Mila in Montreal. And I'm one of the founders of Manifest AI. We're an AI research lab, independent based in New York City, focused on pushing the frontiers of the current paradigm. Understanding that the stack that has taken us this far is not going to be the only pieces needed to get us all the way to the AGI future that everyone is hoping for and looking at the whole field from first principles to make the fundamental breakthroughs needed in order to get to that point. Just a quick question.

3:15How are you funded in that you're a lab without a product at the moment? We ultimately do plan to release products. Once we've made these breakthroughs, we're going to then build technology, build products around the technology that we develop. We're backed by Decibel and True, which are two highly technical VCs based in San Francisco. And we're looking forward to, you know, now that we've made a lot of progress on the architecture side, starting to build our first products around this technology. So we're going to talk about your new architecture that goes beyond transformer models and solves in particular the scaling problem with transformers as you scale the context window scales quadratically because every token has to record its position relative to other tokens.

4:18And as you scale, that becomes untenable or just needs a tremendous amount of compute. Your concept of power retention, which, as you were saying, is a combination of power attention and recurrence, not necessarily recurring neural networks. was that the concept that you started manifest with or is it something that developed once you had started the lab no the specific architecture is something that we developed only after really understanding what the sort of fundamentals of the problem were we started off with just a vision for the future of the field and an understanding of the first breakthroughs that would be needed at minimum to get there.

5:09And ultimately, the angle was based around the specific sort of like technical requirements to solve the problem. But that was sort of downstream. You know, the process of doing research is about understanding and discovering these things. So the insight that we started Manifest with was about scaling loss. And in particular, that certain scaling laws were sort of well-situated by the current paradigm. In particular, the parameter scaling laws is one example, where as you increase the parameter size, the model's performance improves, and the cost to do both training and inference on the model grows linearly with respect to the parameters.

5:49Linear growth is very tractable. This is what has allowed us to push the scale this far. Yes, it's expensive, but you get good returns on the investment because the growth in cost is consistent. But there's a different dimension of scaling. the input size, where this is not the case. And right now with transformers, as you scale up transformers to larger and larger inputs, you end up paying, as you mentioned earlier, a quadratic cost on the size of the input. So if we're talking LLMs, language models, the input is a big prompt, what's called the context window. It's all the text that the model is ingesting, synthesizing in order to produce a response.

6:29And as you increase that context window, The cost to train these models grows quadratically instead of linearly. And it's basically unique among the possible axes of scaling that one might consider in this fact. And if we think that the neural networks of the future are going to be able to process and synthesize data from very large inputs, then just some quick napkin math will make you realize it's completely intractable to be processing these inputs at a quadratic cost. and understanding this and basically making the first sort of like insights towards what a solution would look like was where we were at when we started Manifest.

7:08And since then, it's been about two and a half years and it's been a journey of essentially, both from an empirical and theoretical perspective, understanding exactly what it means to have, where the quadratic cost comes from, why it's important, how it affects learning, how it affects training time, how it interacts with the hardware, right? we're training these models on GPUs. So you need to have something that is very hardware efficient as well. And power retention is the outcome of all of this research. It's not dissimilar in that you're dealing with a state space model, is that right? Then from Mamba, which I've done an episode on earlier, did you start by looking at Mamba or just were state space models attractive to you for other reasons?

7:57Yeah, Mamba and state-based models in general are what I would consider retention models. They're models, actually Mamba 1, less so, but Mamba 2 for sure, are these models that have this sort of duality where they can be expressed either as recurrent models or as attention models. And once you have this, this model that can be expressed in both of these forms, then it also unlocks a secret third form of the model, which people call the chunked formulation, where you basically combine all the best parts of recurrence and attention to get an architecture that both has tractable growth in cost, has linear growth in cost as you increase the context, but also is able to really saturate the GPU, which is what attention is able to do.

8:42And traditional recurrent models, like recurrent neural networks, like LSTMs from the previous generation of recurrent neural networks, were unable to do. So that's sort of the first piece of the puzzle was, and this wasn't a piece that was done only by us, this was done by many groups in parallel, including the Mamba group, to just realize that there exists this sort of like other family of recurrent models that have this attention form duality that enables their really efficient implementation. But that's not enough. The issue that we found with Mamba, and actually with basically all other modern RNNs or modern retention models that we were looking at as we started research is that the state size is too small.

9:23And in particular, the state weight ratio is highly skewed in the direction of the weights. In other words, the state is too small for the weight size of the model. And this is not true for transformers. Transformers, the state of a transformer, it's typically not thought of as a state in the sense of an RNN, but it actually is a very nice identification to understand that the state of a transformer is the KD cache. It's this sort of record of all the tokens that it's seen up until now. And if you consider the KD cache of a transformer to be the state, well, you realize that the KD cache is massive at long contexts.

10:01So transformers actually have a huge, huge state compared to RNNs, even modern RNNs. And this is the downfall of RNNs because state size has scaling laws just like parameter size does. Having a larger state translates into better performance. And so what you end up seeing with a lot of these retention architectures, including Mamba and including Gated DeltaNet and Rockov and various other subquadratic architectures that you might encounter, is that the speed-ups that you get, which you do get at long contexts, come into play at exactly the same point that performance relative to transformers takes a hit.

10:44In other words, by the time you're training on context long enough for them to be meaningfully fast, the performance, because the state is so small, is far worse than transformers anyways. So they end up only being sort of compute optimal at a very narrow range of flop budgets, like very, very small flop budgets. And you basically never want to actually train an architecture like this at scale. What power retention does differently is that the retention part is generic to all state-space models, including, you know, Mamba and whatever else. But the power part, that's the unique part. Power retention is a retention model that has an extra adjustable axis of scale.

11:29You can adjust the state size independently of the parameter count. And this allows us to exploit state scaling as an axis of scale, just like we've already been exploiting parameter scaling. And so we're able to construct models that are retention models. They are RNNs, but they have state sizes that are as large as we want them to be. And this allows us to basically pick the optimal state size always and have actual compute optimal performance for a very, very large range of floppy. Can we back up a little bit for the benefit of listeners that are not familiar with a lot of these terms? First of all, state space models, my understanding is that they record the state of the model at any point in time.

12:25Tell me if I'm wrong, but the weights of all the nodes and a network and all of that is recorded as a state. And that state then goes through a transition as the model learns or operates in inference. Is that right? Can you describe that a little more precisely? Yeah, absolutely. So I think separating between the state and the weight is an important conceptual distinction here. There's sort of two things that inform a model's decision. When the model sort of gives you back a response, there's two things that affect it. One thing is whatever question you asked it. So if you ask a model a question, that is the context.

13:17All the tokens of the question get turned into context, and that context gets summarized into a state. The state is just sort of like a list of numbers that summarize all the context seen so far. And that state obviously impacts what the model's reply is going to be. Then, separately, there's a different source of information. This is basically the entire Internet. And now the entire internet affects the response of the model via the training process. Training is essentially about compressing the knowledge on the internet into the weights, which are just a list of numbers describing all the knowledge the model has about everything.

13:54So it's these two things, the state, which summarizes the specific information in the question, and the weights, which capture all of the knowledge the model has on the internet that combine to produce the response. And state-space models, they maintain a state, and as they see more tokens, for example, as your conversation continues and you go back and forth discussing some topic, they keep updating the state. I see. And the state is the context side. It's not the model weight side. So, yeah. Okay. And then recurrence. Can you do the same with recurrence? Yeah, they're the same. What I just described is the same exact concept of recurrence.

14:38Recurrence is just a function that you can apply over and over again. And that is just exactly what I described with a state. It's a function that you, the input is the state and the new token scene is the thing you're adding to the context. And then the output is the new state and the updated context. So yeah, instance case models are essentially just a modern rebranding of recurrent neural networks. And the power in power retention is, describe that again. The power in power retention refers to a sort of esoteric mathematical operation called the symmetric power. It's basically just a mathematical operation you can do to a list of numbers to turn them into a larger list of numbers with some exploitable symmetries.

15:27And this is basically the key insight is that you can use this operation to make a state larger and larger without introducing any extra parameters. So the details of how it's implemented, while very interesting, I think are not crucial to the larger picture. The important thing is it's just a lever, essentially, a hyperparameter that can be set to adjust the state size of a recurrent neural network. And what's crucial is that it's a very hardware-friendly one. So you can adjust it in a way that doesn't hurt the ability to exploit a GPU. Yeah. And in fact, GPUs have been tooled to work well with state-space models.

16:13Is that right? Yeah. State-space models, one of the main reasons why state-space models are sort of the most powerful family of RNNs is that you're able to get really good performance from GPUs on state-space models specifically. This is not true for all RNNs. For classic RNNs such as LSTM or GRU architectures, there's, for some deep reasons, it's actually basically impossible to get really, really good performance on a GPU. Just the way that GPUs are constructed, they're kind of good at just one thing big matrix multiplications it's not strictly true but first order approximation that's basically what they're good at and there's no way to express uh arbitrary rnm as a big matrix multiplication but for state space models or you know the germany is usually retention models retention models are a special family of rnm for which you can implement them as a couple really big matrix multiplications and therefore get really good hardware utilization.

17:21Through post-training and through inference, the model weights don't change. Is that right? I mean, they're fixed by the training. You operate, you perform inference from them, but the knowledge of the model doesn't change through operation in a transformer model. In this model, do the weights change over time? No, it's the exact same setup. Conceptually, the weights are almost by definition. The weights are basically defined as the things that get updated during training, and the state is in some sense defined as the thing that gets updated as you show it more tokens of context during inference.

18:14So this is basically like the way that these two concepts are partitioned. So no, the weights do not change during training. There might be some other research that blurs the boundaries here, but by and large, this is the set. Okay, so you've now built a model, this power retention model. And what's the name of the model? So we've currently released one model trained with one model whose architecture is based on power retention called Power Coder. And it's a coding assistant at the 3 billion parameter scale, so still fairly small. But what's really cool about it is it's sort of a proof of concept for this idea.

18:54We call it metamorphosis. You start off with a transformer, a pre-trained transformer. And by just doing a little bit of retraining, what you end up with is a power retention model of equivalent performance. What this means is you can actually get the benefits of power retention, such as fast linear cost inference and cheap post-training on large contexts, while still getting the benefits of the pre-trained transformer weights that, for example, large model providers will release. And so you're not actually forced to throw all that very expensive, very useful weight training away. You can actually exploit it even when you're using power retention.

19:35So the point of PowerCoder is sort of as just a proof of concept of this. It's a usable, good code model, very fast inference. You can download it from HuggingFace and play around with it right now. But I think what we're really hoping to do with this is showcase what it looks like to start from a really stalled transformer pre-trained weights and convert it into equivalently performant, but much faster and more flexible power retention there. And that's interesting. When you say start with a transformer model with pre-trained weights, will you use one of the open source models or do you build your own GPT and then work on it from there?

20:22No, we get to start with the open source models. So, for example, we could have started with LAMA. We didn't in this case, but we could have started with, for example, the LAMA 70 billion weights. This is a good, solid transformer model. It has its architecture, it has its weights, and you can use it for inference. And what you can do then, you just take a reasonably small number of GPUs, not nothing, but on the order of dozens of GPUs, not thousands of GPUs, and spend maybe six hours retraining Llama into a power retention variant. And all that means is you start with the regular LLAMA weights and the regular LLAMA architecture, go into the LLAMA architecture code, and remove the call to attention, and put in a call to power retention.

21:10Just that one line swap. And then just start training it. And what you'll see is within just a couple hours of retraining, we'll get performance back to where it originally was when we were just using the LLAMA model off the shelf. But now, of course, it's a power retention model. It's no longer a transformer. And that training, what do you call that? Post-training or fine-tuning? Yeah, we typically call it retraining. Fine-tuning often has a connotation of only training some of the weights. And we find that this works best when you train all of the weights. So we just call it retraining. Okay.

21:50But it doesn't affect the size of the model or the number of parameters. Correct. It's the exact same number for others. Yeah. I mean, what I'm sort of leading to is, and you're talking about how, you know, you're looking for new architectures to get us beyond transformers and toward the Holy Grail. And one of the blockers with transformers, aside from the scaling context window, is catastrophic forgetting. that you can't grow the model's knowledge beyond its original training. You can fine-tune, tweak a little bit, but if you try and do a full retraining or full, you lose a lot of the original knowledge.

22:44Does the power retention model have the same problem? Yeah, it doesn't address that aspect of things. That is just a property of the weight updates. Green descent is making local updates to the weights. And so that means if you make enough local updates to enough of the weights, you'll be in a very different sort of part of function space than where you started out. And I think these days, it's actually not such a problem. People have gotten very good at doing what they call like LoRa, like low rank updates. Or this is what I was referring to earlier with the fine tuning. Essentially, people have gotten good at starting with a base model, a pre-trained model, and only updating a few of the weights or only making small updates to the weights in order to specialize the model on some particular task or data set.

23:36And so doing, they avoid destroying the knowledge in the rest of the network. They're just adding a little bit of specialization for this one particular task that they care the most about. doing real meaningful additional pre-training, like adding a huge amount of extra knowledge to the network, that is still seemingly out of reach of standard techniques. You need to have the original data set still around. And as you do continued training, you would need to basically continue training on the original data set to preserve that knowledge. And then you can start to mix in some new data to obtain that knowledge.

24:14But if you just try to train only on the new data without also in parallel continuing to train on the original data distribution, you will typically forget a lot. But that only happens when you're doing a lot of training on all of the weights. But in your model, if you can increase the context window, I don't know, is it indefinitely? I mean, provided you have enough compute, you can add knowledge in the context window for the model to use an inference, right? Yeah, exactly. I think that many tasks right now are better suited for in-context learning, which is the process you just described, and fine-tuning, which is the process I described a moment ago.

24:58But because current models have either very short context or because of their architectures, limited ability to actually use the thing in their context, it actually doesn't seem that most people are able to get enough value out of pure in-context learning. And so they have to essentially resort to fine tuning as a way of getting the information into the model. But one of those strengths of power retention is that a model trained end-to-end with power retention on long context still has a tractable cost. And these models are not suffering from the problem of only part of the context being genuinely useful.

25:39They're able to get value out of the entirety of their context. So we think that when you start really seeing these power retention models trained at scale on varying interesting data sets, we'll be able to get most of the value that we currently get from fine-tuning just off of pure in-context learning. Just show it some examples in the context at inference time of what you want the model to do and watch it sort of like learn on the fly how to do that correctly. Yeah, that's amazing. And what are some of the applications that you see with this? I'm just thinking people use RAG and all these different pre-training to get around hallucinations and have models make inferences based on proprietary data.

26:31But if you can put it all in the context window and then work with the model, you aren't faced with those engineering problems. Is that right? Yeah, I think it makes things a lot, a lot easier. RAG is essentially a technique for selecting the right tokens to put into the context. So the reason you need to do RAG is because you have to make basically hard decisions about what information do you actually want the model to look at, to synthesize when making its decision. And the larger the context window gets, the less you're forced to rely on RAG because the more you can just dump everything into the context.

27:11and let the model decide for itself what information is most relevant to the decision that it's about to make. Now, this doesn't mean that it will necessarily be able to memorize everything still, so there will probably still be a place for RAG or RAG-like methods going forward, but it will basically relieve a lot of the burden of having to design these essentially ad hoc external similarity functions or search heuristics in order to decide what stuff to put into the context. One of the things that I'm the most excited about actually within this space is information retrieval agents. So rather than having a rad heuristic to decide what information gets thrown to the agent, another way to handle really, really vast or private information sources that the model needs to look at before making its decision is to let the model use something like tool calls to decide for itself which pieces of information it wants to obtain.

28:05Now, the training process becomes one that you could do with, for example, reinforcement learning, of letting the agent learn, given a particular query, what are the right things to search? Given what I found in my search, what is the right thing to search next? What is the next sources of information I should pursue in order to get a complete picture of the response that I'm trying to give? And this is something that requires really, really long context length, because you want to have this entire trajectory be remembered by the agent instead of having to break it up into little pieces where the agent sort of like forgets what's going on in between.

28:41And I think once we move on to that sort of paradigm of information retrieval and we start asking questions to agents that can end to end represent their entire history and access vast troves of private information in order to come up with a good response, will really feel like we've stopped the issues with hallucination. They just won't come up anymore. Because like a human, the model will spend some time doing research before coming to a conclusion, and will actually be able to remember everything that it encounters during that whole research process before giving your response. Build the future of multi-agent software with Agency.

29:21That's A-G-N-T-C-Y. Now an open source Linux foundation project, Agency is building the internet of agents, a collaborative layer where AI agents can discover, connect, and work across any framework. All the pieces engineers need to deploy multi-agent systems now belong to everyone who builds on agency, including robust identity and access management that ensures every agent is authenticated and trusted before interacting. Agency also provides open, standardized tools for agent discovery, seamless protocols for agent-to-agent communication, and modular components for scalable workflows. collaborate with developers from cisco dell technologies google cloud oracle red hat and more than 75 other supporting companies to build next generation ai infrastructure together agency is dropping code specs and services no strings attached visit agency.org to contribute that's agntcy.org yeah but what about the problem of you know the longer the context window the less performant the model and the model gets lost or in the context or you know some of the knowledge in the context it doesn't pay attention to.

31:16What about that issue? I don't know what the word is for that issue. This is something that we attribute mostly to that self-same quadratic cost of long context, but indirectly. So let me explain. Given that model providers want to offer their users long context models, and it's intractably expensive to train a true transformer on long contexts for any meaningful amount of time. What a lot of model providers do is come up with a sort of in-between solution. What they do is, instead of using a true transformer, they use what's called a windowed transformer or a sparse transformer, something that only attends to a small local subset of context rather than actually attending to the entire context window.