Training Zamba: A Hybrid Model Master Class with Zyphra's Quentin Anthony

In this episode of The Cognitive Revolution, Nathan dives deep into the world of state space models with returning co-host Jason Meaux and special guest Quentin Anthony, Head of Model Training at Zyphra. Explore the cutting-edge Zamba 2-7b model, which combines selective state space and attention mechanisms. Uncover practical insights on model training, architectural choices, and the challenges of scaling AI. From learning schedules to hybrid architectures, loss metrics to context length extension, this technical discussion covers it all. Don't miss this in-depth conversation on the future of personalized, on-device AI. Check out more about Zyphra and Jason Meaux here: Zyphra's website: https://www.zyphra.com Zamba2-7B Blog: https://www.zyphra.com/post/zamba2-7b Zamba2 GitHub: https://github.com/Zyphra/Zamba2 Tree attention: https://www.zyphra.com/post/tree-attention-topology-aware-decoding-for-long-context-attention-on-gpu-clusters Jason's Meaux Twitter: https://x.com/KamaraiCode Jason's Meaux website: https://www.statespace.info Be notified early when Turpentine's drops new publication: https://www.turpentine.co/exclusiveaccess SPONSORS: Weights & Biases RAG++: Advanced training for building production-ready RAG applications. Learn from experts to overcome LLM challenges, evaluate systematically, and integrate advanced features. Includes free Cohere credits. Visit https://wandb.me/cr to start the RAG++ course today. Shopify: Shopify is the world's leading e-commerce platform, offering a market-leading checkout system and exclusive AI apps like Quikly. Nobody does selling better than Shopify. Get a $1 per month trial at https://shopify.com/cognitive Notion: Notion offers powerful workflow and automation templates, perfect for streamlining processes and laying the groundwork for AI-driven automation. With Notion AI, you can search across thousands of documents from various platforms, generating highly relevant analysis and content tailored just for you - try it for free at https://notion.com/cognitiverevolution LMNT: LMNT is a zero-sugar electrolyte drink mix that's redefining hydration and performance. Ideal for those who fast or anyone looking to optimize their electrolyte intake. Support the show and get a free sample pack with any purchase at https://drinklmnt.com/tcr CHAPTERS: (00:00:00) Teaser (00:00:42) About the Show (00:01:05) About the Episode (00:03:09) Introducing Zyphra (00:07:28) Personalization in AI (00:12:48) State Space Models & Efficiency (Part 1) (00:19:22) Sponsors: Weights & Biases RAG++ | Shopify (00:21:26) State Space Models & Efficiency (Part 2) (00:22:23) Dense Attention to Shared Attention (00:29:41) Zyphra's Early Bet on Mamba (Part 1) (00:33:18) Sponsors: Notion | LMNT (00:36:00) Zyphra's Early Bet on Mamba (Part 2) (00:37:22) Loss vs. Model Quality (00:44:53) Emergence & Grokking (00:50:06) Loss Landscapes & Convergence (00:56:55) Sophia, Distillation & Secrets (01:09:00) Competing with Big Tech (01:23:50) The Future of Model Training (01:30:02) Deep Dive into Zamba 1 (01:34:24) Zamba 2 and Mamba 2 (01:38:56) Context Extension & Memory (01:44:04) Sequence Parallelism (01:45:44) Zamba 2 Architecture (01:53:57) Mamba Attention Hybrids (02:00:00) Lock-in Effects (02:05:32) Mamba Hybrids in Robotics (02:07:07) Ease of Use & Compatibility (02:12:10) Tree Attention vs. Ring Attention (02:22:02) Zyphra's Vision & Goals (02:23:57) Outro SOCIAL LINKS: Website: https://www.cognitiverevolution.ai Twitter (Podcast): https://x.com/cogrev_podcast Twitter (Nathan): https://x.com/labenz LinkedIn: https://www.linkedin.com/in/nathanlabenz/

Speaker names marked ? are inferred from speaking time and the show’s known hosts — they are not voice-verified.

• 0:00Jason Meaux?

  The future of AGI will involve a combination of cloud and on-device deployment.

• 0:06Speaker 01

  These large modelic model companies, like Intensive OpenAI, just can't really specialize to every single person on the planet. We think that you need to have your own set of weights. And changing a system prompt per person is not enough, right? We want to actually bake into the weights. You can make the model simulate learning faster than it really is by doing activation steering. If the user tells the model, you're being too dry, then you can very quickly steer the activation to be a bit more fun. Until tonight, when you can bake into the model, I think it's got to be continual, and it's got to be per user. And the only way to do that is with weights on the phone.

• 0:42Nathan Labenz?

  Hello and welcome to The Cognitive Revolution, where we interview visionary researchers, entrepreneurs, and builders working on the frontier of artificial intelligence. Each week, we'll explore their revolutionary ideas, and together, we'll build a picture of how AI technology will transform work, life, and society in the coming years. I'm Nathan Labenz, joined by my co-host, Erik Torenberg. Hello and welcome back to The Cognitive Revolution. Today, we're once again going down the state-space model rabbit hole with returning co-host Jason Moe, who regular listeners will remember from our Mamba Palooza Literature Review and Albert Goo Interview episodes, and Quentin Anthony, Head of Model Training at Zifra, a large language model startup that's just released their Zamba 2.7B model,

• 1:26Nathan Labenz?

  which is built on a hybrid architecture that uses both the selective state-space mechanism and the traditional attention mechanism, albeit with some notable tweaks relative to the standard implementation.

• 1:37Nathan Labenz?

  In addition to sharing Zifra's high-level vision for highly personalized on-device AI, Quentin was super generous with both his time and knowledge, sharing a wealth of practical lessons learned from the front lines of model training. Over the next two hours, we will cover the delicate architectural choices that balance efficiency and capability, the many practical challenges of training at scale, including choosing the right learning schedules for different phases of training, the nitty-gritty details of training hybrid architectures, including why Zamba models don't need positional embeddings, the Zamba model's use of shared attention blocks and internal LoRa adapters to maximize performance on the edge, the not-so-simple relationship between loss metrics and model quality and capabilities, as well as the challenges of context length extension.

• 2:21Nathan Labenz?

  The Zephyr team's experiments with different optimizers and why they're sticking with Adam for now, Quentin's intuitions about the relationship between model scale and loss landscapes, and finally, even their recent published work on tree attention, which offers important advantages over ring attention for multi-node training. I have to say, I got a lot from this episode. And while it's technical enough that I wouldn't necessarily call it entertainment, I am confident that you will too. If so, we always appreciate it when folks take a moment to share the show with friends or write an online review on Apple Podcasts or Spotify. And we welcome your feedback via our website, cognitiverevolution.ai, or by DMing me on your favorite social network. Now, I hope you enjoy this highly technical

• 3:03Nathan Labenz?

  conversation with co-host Jason Moe and guest Quentin Anthony, model training lead at Zephyr. Jason Moe, returning guest and co-host and chronicler of state space models at statespace.info, and Quentin Anthony, model training lead at Zephyr, which has just released the new Zamba 7B SSM hybrid model. Welcome both of you to the Cognitive Revolution.

• 3:28Speaker 01

  Thanks for your time. Great to be here.

• 3:30Nathan Labenz?

  Jason, regular listeners will know as a partner in crime who's also obsessed with state space models and the potential that they have to unlock new capabilities in terms of potentially long-term memory, extreme efficiency, all these kind of interesting things. And we've been down the rabbit hole together on that a couple of times. So I was excited to have him back to help me with this conversation about the new Zamba model and all the ins and outs of that. So Jason, I'm going to ask you to lead the questioning today, and I'll be in the supporting role, but listening intently and probably jumping in with a few of my own follow-up questions along the way as well. How's that sound?

• 4:13Jason Meaux?

  Sounds good. Let's do it. So let's start with maybe a broad question first. So on the Zephyr website, it says that, quote, "The future of AGI will involve a combination of cloud and on-device deployment with an increasing shift towards local inference."

• 4:32Jason Meaux?

  So why is local inference going to be important? And how does this idea overall influence the direction of the company?

• 4:43Speaker 01

  Yeah, there's a few aspects, a few angles of this. The main one is personalizability. We think that doing system prompt packing, things like these large modelific model companies doing like in Dr. Grope and AI just can't really specialize to every single person on the planet. We think that you need to have your own set of weights and changing a system prompt per person is not enough, right? We want to actually bake into the weights. You like to be talked to this way. Your favorite restaurant is X, that sort of thing. A second thing here is privacy. So there's a lot of data about your laptops, your phones that you don't really want to be communicating to the cloud. There's a lot of enterprises out there who don't want to share all of their data with OpenAI, all of their proprietary code and things.

• 5:24Speaker 01

  So keeping that all on device, on prem in the organization is where we want to keep it. And then there's just some very practical challenges here of OpenAI spends, what was their CapEx realistically, right? In order to infer all these models, that's not really something that can keep up forever. It's being fueled by a lot of VC money. But if everyone's device is able to efficiently run their own models, then we offload that to the users. And then it's also just it's much faster. You can run it offline. All these small benefits come out from having the model locally in the device, but it's all steered by this global personalizability, privacy, these sort of North Stars. In terms of on the cloud, there are some model capabilities that are more challenging.

• 6:08Speaker 01

  Like they're too challenging to have on a one B or a two B or something like that. Most people we see actually, when they communicate with models, they just need these higher level, more simplistic tasks, just like basic chat, like factual recall. They're not asking like these O one type math questions, these Olympiad questions or anything to like their day to day on device model. But sometimes you will need it. And that's why we have also a cloud offering like in the seven B range. The seven B range is a turning point where you can fit it on like more powerful laptops, have long prem, but it's able to think through these more difficult problems more clearly. And we want those on device models to be able to know. Like when they need to go back and ask like an oracle model, something a bit more powerful.

• 6:48Speaker 01

  And then this is all seamless to the user and then you can pick and choose how often your model wants to talk to the cloud.

• 6:55Jason Meaux?

  It it sounds like it's when I hear local inference, it's often a privacy bullet point is the main bullet point that gets hit. But it sounds like the overall idea is it's some mix of capability and privacy. That if we had super efficient, super personalized localized models, that this is going to increase capability of the full stack. Is that correct?

• 7:17Speaker 01

  Definitely. Yeah, you can't forget personalizability. It's really hard to bake in like personalized outputs, personalized like how the model likes to talk to you in the same way. It's just a hard problem.

• 7:28Nathan Labenz?

  I would love to dig into that a little bit further.

• 7:31Speaker 01

  Yeah.

• 7:31Nathan Labenz?

  Because I'm working on a project right now to try to get a model to write as me. I think a white whale project for me for a while. I figure if a model can write as me in a way that's compelling to me on a consistent basis. It seems like that's a pretty big threshold. So I've been trying to see, can I make that happen? And the best readily available fine tunable model today is G P T four O. And so I've been working mostly with the OpenAI stack. They make it easy to do fine tuning. They don't tell you too much about exactly how it works. One presumes that it's a Laura type process that they're running. I think there's a lot of reasons to believe it's something like that.

• 8:14Nathan Labenz?

  It doesn't seem to learn facts very well. It does seem to learn like patterns and like stylistic layer, if you will, reasonably well. I don't wonder what you think is the right way to even approach personalization in the first place. Is it like when you envision doing this personalization for individuals on their machines? Are we talking Laura's? Are we talking like full training of a one or two or three B model? How many tokens does that take? Do we still need to like stuff huge amounts of context in there? Because that's another approach I've been taking where I'm like, maybe if I give it 50,000 tokens of everything it needs to know about me as the system prompt. Then it like won't need to hallucinate. And I still haven't quite figured that out. I'm still getting hallucinations.

• 8:58Nathan Labenz?

  So what is the the overall paradigm or path to this like deep personalization for individuals?

• 9:05Speaker 01

  Yeah, so the exact process will not be entirely clear until we have a bunch of people having, you know, Zamba on their phones, but we have a few like bets. One is that just we have a lot of expertise on continual learning and continual pre-training. So we envision doing like a wait update like overnight, maybe while your phone is plugged in, your laptop's not really actively being used on consumer hardware. It's really cheap to do this when the model is really small. Updating Laura's, as you mentioned, is another like that makes it even cheaper. But just continual learning, I think, is the way to go. In terms of your GPT four O problem, I think I would need to know more details about the the fine tuning data set and the model, which I'm just never

• 9:47Speaker 01

  gonna get. We find continual learning on a small model per person is really great. There's some sort of you can make the model simulate learning faster than it really is by doing things like activation steering. Like you remember Golden Gate Bridge Claude and stuff. If the user tells the model you're being too dry or something like that, then you can like very quickly steer the activation to be a bit more fun. Until tonight, when you can bake into the model, okay, this person likes generally to be a bit less dry and to talk in a more fun way. I think it's gotta be continual learning. It's gotta be per user. And the only way to do that is with weights on the phone.

• 10:21Nathan Labenz?

  But how much of that exists today? Like, can I get a Zamba? First of all, like which like scale of model would you recommend that I bring locally? And can I do this continual learning? Can I like start to feed it my emails and my Slack messages? I assume I would have to mix in some like general purpose pre-training mix as well. I wouldn't just go all in on like my content. So tell me a little bit more about what that looks like practically. I'm taking notes.

• 10:51Speaker 01

  Yeah, absolutely. So right now, we have the 1.2B, the 2.7B, and the 7B that we just dropped, all of the our Zamba two series of models. The 1.2B is really good if you're like really edge constrained, like you don't want your phone to use much battery, you want it on a phone in the first place. In terms of whether you can do this yourself, right now you can. It's just a little less user-friendly than it will be someday. But right now, we have our own Hugging Face fork for Zamba2. So you would just download the model weights and then you would just have to up tokenize and actually just train on whatever personal emails or whatever else you want baked into the model. The product vision here though for Maya is that you have like a Maya cloud across all of your devices.

• 11:33Speaker 01

  And that's just some like social aspects here. You can upload pictures. It'll be multimodal. Your conversations and such, you can decide how much you want to share, how much you want Maya to be personalized to you. But training on that cloud of your own data, just as a continual pre-training on whatever edge hardware you have is the high level. I'll start here.

• 11:54Nathan Labenz?

  Are you doing this today? Like, how far have you guys taken this internally? Do you have a Quentin model that drafts your emails or what's the frontier in application?

• 12:04Speaker 01

  So earlier you said the 1.2B is good for edge constraint. The 2.7B is really good for quality. I'm always surprised by it. We do have some fine tunes for that. So we have some role-play models, summarization models, if you want to summarize emails, summarize meetings and such. We have some very early results on audio as well. So generating audio and the style of people you talk to, note-taking audio for the meetings you're attending. It is definitely more mature in-house, but we really want this to be really seamless before we give it to everybody else. So I would say that the 2.7B, I've seen really amazing results on in terms of summarization and like these tasks that I was mentioning. And it's all from this Hugging Face, just fine-tuning our DPO on whatever my personal data is.

• 12:48Nathan Labenz?

  Cool. I'll probably have some more questions about that as we go. But maybe a good transition to, as we think about this on-edge paradigm.

• 12:59Nathan Labenz?

  Jason, you want to take the baton and talk about why state-based models are the, or hybrids at least, are the attractive approach to take there?

• 13:08Jason Meaux?

  Yeah, absolutely. I guess when you think of what kind of models can run locally, you get hit by transformers in two places pretty hard. The first place is just memory. Anyone who's done their own work with a limited compute budget knows the classic out-of-memory error. So you quickly, whether that's too large of a model in terms of parameter count or if your KV cache is growing and at some point you run out of memory. So transformers, because of the way they work, they're famously pretty memory intensive. The KV cache scales as you scale your sequence. And then, of course, secondly is just computation.

• 13:50Jason Meaux?

  Transformer has amazing properties, but it is pretty computationally intensive, especially when compared to other architectures.

• 14:00Jason Meaux?

  So when you start looking at Mamba, that's one of the exciting things about it. It has properties that are much more memory efficient. And computationally, it's able to be more efficient as it processes through sequences due to its fixed state size. And so you even see that play out in the literature. There's been some early work with Mamba models with a couple of papers on robotics. These are models that would have to run locally. You could not do cloud API calls for split-second decisions that would be made in robotics. So you can just imagine, finally, maybe we have an architecture that helps inform what the ultimate models running locally could look like for these kind of applications.

• 14:43Jason Meaux?

  So how does that vision look like for Zifra and the Zamba architecture? What does the ideal local inference strategy look like as it relates to Mamba?

• 14:54Speaker 01

  Yeah, so we love Mamba for its systems properties. As a guiding principle, in terms of model training, we try to train the models that are perfect for inference, because that's where we're going to need them to perform the most. We can eat some training cost, but at the end of the day, it needs to be blazing fast for inference. This translates to a few different things for us. For one, attention for dense transformers, like a PyTIP model, is just not going to happen for us. We don't think that this is efficient on a phone. Earlier, I mentioned, like, we want text and email summarization. Sometimes you want news article summarization. Those inputs are too long for the KV cache. Like, it's just going to grow outside of memory for most phones today. Attention, though, is really nice. So you get these exact cross-sequence dependencies that we think are just required for these some specific tasks.

• 15:39Speaker 01

  Remember, I say specific, so not every task needs attention. But there are some tasks like in-context learning and long sequence dependencies for which attention is just necessary for you to be performing on. So eventually, if you train on several trillion tokens, then eventually Mamba becomes somewhat performant at this on benchmarks. But I've never played with a pure SSM, a pure RNN model that's able to really like... To speak well in an in-context learning regime. And I think that's a common intuition that people have when they actually sit down and play with the models. So pure SSMs have quality issues from what I've seen. And then pure attention dense transformers have, you know, performance issues.

• 16:21Speaker 01

  So you can't get fire on a phone realistically to, you know, to run the amount of time with the fastest time for a token that the users actually want. So there's some other architectural changes here. Instead of having independent attention blocks, we have this global shared attention block. And the reason for this is that us and a few other groups have independently found that attention blocks are highly correlated across depth. So they're telling the MLPs in a dense transformer, these are the important tokens. These are still the important tokens. These are still the important tokens. There's some small changes, but most of the specialization across depth is coming from the MLPs. So if you tie all of those weights together, then you're able to allocate more parameters

• 17:04Speaker 01

  towards the, for example, the Mamba block for us. If you were on a dense transformer, you could apply more to the MLP blocks or something. So now we just have, like for the Zamba 1 series of models, you just have a single attention block and you apply it every six Mamba blocks across depth of the network. We improved on this a bit in Zamba 2 by seeing, okay, a little bit of specialization across depth is helpful. So we added LoRa back in, and those LoRa's are independent across depth. So now you have this single global attention block in the 1.2B case. And in the 2.7B and 7B case for Zamba 2, you have inter- inter leaves. So you have attention block 1, attention block 2, attention block 1, attention block 2, all across depth. And then individual LoRa's across depth as well.

• 17:49Speaker 01

  So you get a little bit of specialization. But back to the systems angle, you're not really applying very much attention. So your slope of like your KV cache memory, here we've seen this in all of our plots, is something more like 10 or 13 invocations rather than 30, 32 in like a dense transformer, because we're just invoking far fewer attention blocks. Now, the very first criticism people should be thinking is that, okay, you're applying for the same parameter budget, you're applying parameters over and over, the total number of flops is going to increase. And this is true. We're applying blocks more often than like a typical 7B. But the Mamba blocks, and now the Mamba 2 blocks, are so high throughput on parallel

• 18:32Speaker 01

  hardware in general, that this totally makes up for it. And we end up with a net improvement of about 20 to 30% in time to first token and time per output token. So it makes up for our sins of sharing these attention blocks. But yeah, this, both of these translates to significantly lower memory overhead, significantly lower time to first token at the pre-fill stage, so the model responds really fast initially. And then time for output token is also 20, 30% beat compared to similar dense transformer models.

• 19:00Speaker 03

  Hey, we'll continue our interview in a moment after a word from our sponsors.

• 19:04Jason Meaux?

  The choice for using global shared attention blocks is really interesting to build the hybrid model that way. I guess it's in alignment, I saw on Zyfer's site, the desire to maximize performance per parameter. That seems like a choice that obviously is doing that as well as maybe some other architectural choices. How important is it in the Zamba family of models to focus on performance per parameter?

• 19:28Speaker 01

  It's pretty important, right? At the end of the day, most of your memory overhead is coming from just storing the parameters. We can get away with a lot of quantization, but each parameter is really painful when you're deploying it on a phone. So trying to make all of those parameters maximally useful for the attention block, as in tying it all together and making that attention block really well-trained, yeah. Maximizing the performance of each one of those parameters, we've got to do. We don't want this really sparse model where most of what we're storing is not doing much. But yeah, very important.

• 20:02Nathan Labenz?

  I think there's like a broader kind of trend here of like dense attention being the original thing. And then there's been obviously many attempts to make that more efficient with like linear attention or things that are like not quadratic. My sense is that those have all been pretty good, but somehow still fall short. Then you have these other techniques of shared query attention. Maybe you can help us understand in a technical sense, like exactly what that is. And then I'm also aware of shared query key attention. And now you're going all the way up to shared full attention blocks. So help us map that space and just understand like, what are the different options and trade offs? And is this like the

• 20:46Nathan Labenz?

  end of that paradigm? Or is there even like still further that we could go on the spectrum from full dense to full shared? Is there even further on that continuum?

• 20:57Speaker 01

  So we use grouped query attention, car upcoming attention blocks. I would say I haven't seen like a non quadratic attention that really performs the same and does what we wanted to. Frankly, I would say that there is a very interesting emerging area of work where people try to remove the attention from attention models. So that maybe the theory there is just more of a play from, for example, Liquid AI sort of player. You take a model that's already been trained with attention, and then you try to linearize that attention, get a final post training step. And the hope there is that the model's already baked in the ability to do these long sequence dependencies in context learning. And then you can just sort of copy it over to much lower cost blocks. It's a potential play.

• 21:43Nathan Labenz?

  Can you give it a, help us develop an intuition for, this is like really a question I want to, I'm asking for me, and hopefully other people find it useful. I have a intuition for why the full attention would be qualitatively better versus like a state space mechanism or some approximation of attention. For me, that kind of boils down to, we compute a relationship between every token and every other token. We don't forget about any tokens. It's all there, and it's all, when a new token comes in, we can see how that relates to everything that came before with no compromises. When we have that, we try to map that onto a linear thing. What's the intuition for why that like would work or wouldn't work?

• 22:26Nathan Labenz?

  Maybe there is no intuition and people are just trying stuff and seeing what works. But I'm grasping for understanding of what that transformation is and why we would. What should we expect and what should surprise us?

• 22:38Speaker 01

  Okay, my first order intuition here is that attention is really good. Because the point of the attention block is to note, take as much as possible about the incoming sequence. So you have this really long sequence, let's say. And if you're taking the dependency from every token to every other token, you're giving MLPs or whatever it is that you're projecting as much information as possible in the sequence. And you have multiple heads, each of them saying, okay, they are each looking at the sentence in a slightly different way. And then if you have large heads, they're keeping a lot of detail about their own specific angle of the sequence. Any sort of attempt to linearize this is saying, okay, how much of that detail do we actually need?

• 23:19Speaker 01

  So the original attention paper was so useful because a lot of that, we found that MLPs could make use of a lot of that detail. Because MLPs are just like trying to find the deeper meaning on whatever those notes that the attention took. So attention takes notes about this, as many notes as possible, but very computationally heavy. And then MLPs try to find deeper meaning from those relationships. They sort of mix the heads, the new cross heads. So for linear attention, it's all just trying to say, okay, are all those heads doing something important? Are we able to cut away a lot of them? Do you really need every relationship for every other token? And this is a very theoretical question, right? So it's not super clear how many notes, whatever you need.

• 24:02Speaker 01

  You also have different ways of projecting across these heads. So MLPs are a very powerful way to mix across heads and make use of these notes that are being taken. But like a Mamba block has much smaller projectors. It's a very different story there of like how much they can really do from the notes taken by the selective scan versus an attention block. Maybe if you had like bigger MLPs for a linearized attention, maybe that would make up for it. Maybe you still need like more notes if you had a Mamba block without like powerful projectors and then an MLP. This is really like something I can't really say. If I did, then I would make the model. But I think all of these people are motivated by the fact of,

• 24:47Speaker 01

  how much can we throw away from these notes and still get performance? Because do you really need n squared dependencies? Does this make sense?

• 24:55Nathan Labenz?

  Yeah, and so coming to the shared attention, I guess I'm not 100% sure what exactly all is shared. But if I understand correctly, it's like you are literally using the same attention weights over and over again at every layer, but you are recomputing the effect of that attention at each, or you know, every six layers because the computation has proceeded. And so each time the attention block is used, you're able to use the same parameters. That saves space, but it doesn't save like runtime computation per se, because you're

• 25:35Nathan Labenz?

  still doing that attention on the input as it's been processed to that stage through the model so far. Is that right?

• 25:43Speaker 01

  Exactly. Yes.

• 25:44Nathan Labenz?

  Okay, cool. And that's different from, so when you use a reuse an entire attention block, that's different from just reusing the queries or just reusing the queries and keys, you're reusing queries, keys and value.

• 25:58Speaker 01

  Yeah, so yeah, we have taken it to the extreme of saying these attention blocks are pretty much just reminding MLPs anyway. So let's just stop storing all of these independent parameters that eventually learn to do the same thing anyway. So let's just get them, when you're reapplying, you also get really accurate gradients, right, at training time? Because it's applied six times instead of, you know, one if they're all independent. So we get this really highly trained, highly saturated attention block and it's completely tied to query key value. And then we add a little bit more specialization back in for Zamba 2 by having alternating layers. So now you have two attention blocks, they're each tied, but then we apply them one after another. And then we just have some depth-wise LoRa's that are themselves independent. So you get a little bit of specialization across

• 26:43Speaker 01

  depth, but not too much.

• 26:45Nathan Labenz?

  Yeah, interesting. The fact that it's alternating is somewhat counterintuitive. When you say there's specialization across depth, I would have thought you'd like maybe segmented into halves and have a first half and a second half or a first, second, third blocks, but you're still staggering them. And then I guess with the LoRa adapter, you have some kind of depth specific thing that is fully localized to that place in the model.

• 27:10Jason Meaux?

  Yep.

• 27:11Nathan Labenz?

  Cool. Fascinating. All right. I'm talking more than I intended to. Jason, let's get back to your main line of questioning on the technical details.

• 27:19Jason Meaux?

  Yeah, I guess if I could just get, set the history straight real quick, we can go back to February of this year. You actually released, because this is how Nathan and I actually came from there with Zyphra, was not the release of Zamba, but the release of Black Mamba, a mixture of experts model that used Mamba. This was really early in the literature. I think Nathan and I were counting papers at that point. This may have been among the first 10 Mamba papers published. We've since lost count. Yeah, well over 300 now. I really don't know the full count. So what was it about Mamba in particular that got Zyphra making this early bet and going forward with Black Mamba?

• 28:03Speaker 01

  Yeah, so you're taking me back. When I first joined Zyphra, we were think, we're all about MOEs. We're like a mixture of experts with very vanilla attention blocks, no exotic shared blocks or anything like that. It was the way forward for on-device. We had some ideas for like reducing the inference time, memory and compute for these, and we were all in on training. We trained some good models actually at the time for MOEs. Then Mamba came out and we were like, okay, wow, this is amazing for on-device. But the question is, what does quality look like? Because we already knew that like attention-free models had this sort of problem of quality. Our motivating figure was in the original Mamba paper, there's a figure where they try pure Mamba versus Mamba

• 28:47Speaker 01

  plus attention versus Mamba plus MLP. And they have the loss per flop graph for all of them. And they don't show the absolute loss. So you don't know what the actual quality is at the very end. So the MLP plus Mamba hybrid.

• 29:04Speaker 01

  They saw what had like was really expensive, had a lot of flops for equivalent loss. And we thought, we don't know what the actual quality is here. So maybe the quality is actually really good. And that you in fact, don't need all of these notes from attention. All these really rich cross sequence dependencies, which ended up being the wrong, the exact wrong thing. And we replaced these really expensive attention blocks in our current MOEs. Because they're the ones that's really get, really giving us so much trouble. And then just do Mamba and then a mixture of experts MLP block. And just do that over and over. That would be really inference efficient for what we were preparing for. So we trained it. So we set everything up. We trained the model. We did indeed see that in terms of long sequences, it was amazing for overhead.

• 29:46Speaker 01

  But Mamba on its own can't really make up without these, the attention cross sequence dependencies in context. Like those exact same things we learned. So the model was not great on MMLU, for example, which is like your early indicator. Not that you want to saturate MMLU. But that's an alarm bell of, oh, your model is not really able to learn in context. So that was the early prep. That was our early lesson. And that's why we didn't really proceed with doing Mamba without attention anymore. MOEs, maybe they'll come back in the future. But yeah, that was the initial model. Any questions or thoughts on it? It was an experiment that we shared with everyone else, rather than something that we wanted to like actually like deploy.

• 30:24Speaker 03

  Hey, we'll continue our interview in a moment after a word from our sponsors.

• 30:29Jason Meaux?

  It's a huge change in strategy to pursue, to go from an MOE to something like ZAMPA. I mean, you know, something that's, MOE model is really hard to run. The parameter count's huge. If you scale that up in terms of just memory, you've got to store a lot more weights. And then you go to ZAMPA one, way more parameter efficient. So I guess just the thinking of what drove that, drove you'd go the other way. I understand the attention part with in context learning, that makes sense. But I guess like the local inference part.

• 30:58Nathan Labenz?

  Sort of as well a little bit?

• 31:00Speaker 01

  Oh, definitely. Like in terms, there's a reason all of our black Mamba plots are like in terms of efficiency. It is such an efficient model because there's no attention anywhere, there's no quadratic complexities, there's no KV cache at all. Like it's just a constant state that Mamba has to require. So your memory overhead is on the floor. Oh, I also forgot to mention that motivating figure. So yeah, the reason we were doing MOEs was that we were trying, we knew that like the, the flops per loss of like MOEs is much more amenable because you're only routing to a few of these, a few of these experts at inference time. So we were trying to hope that we would like dip that line of loss per flop quite a bit lower than you've both the attention plus Mamba

• 31:41Speaker 01

  and the MLP plus Mamba. And in terms of pure loss, oh yeah, we did, we definitely did that. But in terms of model quality, it just wasn't where we wanted it to be.

• 31:50Nathan Labenz?

  Can you talk more about that divergence between pure loss and model quality? I think that is one of the biggest picture questions that the entire field is struggling with, right? Like from a sort of super big picture, like when does AGI arrive? And also people worry about this in the context of AI safety with emergence. We believe or we feel pretty confident we can predict a relatively narrow range of what the loss value is going to be as we, you know, extrapolate scaling laws out, but then everybody's, but what does that turn into? And so you're highlighting here one moment where you have, or one example where you have seen this divergence between

• 32:32Nathan Labenz?

  loss and quality in a negative way where like the quality is less than you might have hoped based on the loss. Presumably it could go the other way too, right? It doesn't seem like these are, I think more guidelines than actual laws is probably the right way to think about the scaling laws. So yeah, tell us basically, I'd be interested in everything about how these divergences between loss and quality open up.

• 32:55Speaker 01

  Yeah, so loss is not a very useful indicator, like at all. So early, I use it to detect spikes, obviously, so it tells you, okay, your model died. And we definitely saw those, especially with Black Mamba, is a very unstable model to train, actually. In terms of model quality, though, it's not a great, like, generally, you want it to go down. Do you want it to go down to three or two? It's dataset dependent, the noise of your data strongly determines this. I would say one thing I check it for is that it's a very early indicator of whether your model is saturating, or doesn't have the high enough learning, like whether it's progressing. So your loss decrease is obviously going to slow like way down.

• 33:35Speaker 01

  As you hit a stable point of training, like you've reached like the slope of you optimize your landscape and you're going down. So that's great. But if your learning rate, for example, was not aggressive enough, then you might just completely plateau in your loss. And this is kind of like danger, like you've done something wrong. So that's really all it can tell you. Everything else is coming from evaluations. Now evaluations are also tough because if you maximize evaluations of the cost of everything else, which if I wanted to do that, I would just do synthetic textbooks for my entire training data, and I would blow up evals. Great. But when you actually talk to the model, it's super dry. You can't really fine tune it very well because it's assumed that it's exact very clean distribution of data.

• 34:20Speaker 01

  It doesn't really respond to unclean data very well. So what you need is to understand what each eval is actually doing and providing for you. For example, I believe it's Hella Swag is one of the evals. And that one is like general language capability we're finding, like does your model just model very simple sequences? And that's like a very smooth eval, it's like a better loss almost, where it just like linearly increases across training. Then you have weirder evals like MMU that are important indicators. For one, they tell you how contaminated your data is. MMU explodes, you probably done something wrong. You probably downloaded some data that includes some MMU questions. And then obviously you also, like, we do very careful test overlap with our training sets as well to also indicate this.

• 35:05Speaker 01

  But MMU is a really good indicator of whether your model has reached this emergence of in-context learning. And it really is an emergence, so it doesn't show a signal for a very long time. And then near the end, it starts to grow up because your model is starting to grow up this in-context learning task, which is difficult for it. Finally, there's a bunch of us like we test like factual understanding with the ARC. Easy and hard challenging, or challenge evals. So those really show us how well the model like is understanding factual recall, how much Wikipedia we were put in, how many textbooks. And then lastly, you really just need vibes, man. We have this like a set of prompts that show us really key, like is the model fun? Like things like, is the model able to role play like as a pirate?

• 35:49Speaker 01

  Is like legitimately something that's really important to see because if you give a model like phi, you tell it to act like a pirate, it'll ignore you. Like it'll stick with it a little bit and then just immediately drop it. And that's important to test. We also test things like not quite strawberry, but like how many times like a flower bloom? If it blooms this many times like in a year, like there's a lot of vibes questions that we ask the model across training. And at the end of the day, you want your model to perform well on things you actually want it to do, not just benchmarks. And those are the questions that we're asking, like personal assistant questions and this sort of thing. And that's what actually indicates whether the model is becoming quality. Anything there you want me to drill down on? That's kind of our process.

• 36:33Nathan Labenz?

  That's fascinating. Moments of, I think more anecdotes, if any, on like moments of emergence are interesting. It sounded like you were saying MMLU kind of you understand to be a reflection of in context learning. Is that because you're like doing few shot testing? Like I'm a little bit unclear on what basis you would equate MMLU with within context learning? And also just interested in what the two curves look like as you go through training. I'm taking away, I think that like the loss curve is just doing its smooth decay thing. And then, but you're seeing these like relatively sudden spikes in performance on.

• 37:16Nathan Labenz?

  Key tasks, is that right?

• 37:19Speaker 01

  This is exactly true. It's also like all of the vibes questions are also like this where it's like, oh, the model suddenly gets it like early on and even halfway through the model keeps getting tripped up on silly things, but when the model gets it, we're like, okay, we've reached this key milestone. So yeah, loss is always you just zoom in and you're like, okay, is there a slope? Is it negative? Yes, move on. The reason MLU is hard and the reason it tests like in context learning is that its format like it requires a very strict format. That's really it. And also very weekly test factual recall, because these are all college based questions, but it's answer as a letter, a multi-choice, A, B, C, or D. So the model is used to filling in the next token. It's not really used to just this big, long question. Oh, format your answer like this. The model is not used to just saying C.

• 38:03Speaker 01

  It's used to saying, oh, astrophysics is this field with all these things that used to actually using text. And for the model to understand, you are supposed to answer with a single letter and that letter is supposed to encapsulate all of the background knowledge of whatever the question is trying to get you to answer. That is something that shows up very late. So yeah, MLU is like random for most of training. And then after it, it depends on the architecture, actually. So for a dense transformer, you need to be like later on in your learning rate schedule. So your learning rate has to have settled down a little bit and then you have to have seen something like several hundred billion tokens. And then suddenly MLU just starts to climb pretty linearly until you set the rate. So it's like flat and then

• 38:47Speaker 01

  it grows up and then it flattens again. Once you've seen your learning rate is too low and you've seen enough tokens already. For pure SSMs, actually this number is really high. Falcon and Mamba, this was like three or four trillion tokens, I think. It's multiple trillion before and we theorize this is because of in context learning, like it takes a long time. Like you need a little bit of attention to kick the model into understanding. Early on in the question, it told me answer with a letter. And that sort of relationship gets lost in a pure SSM block. But yeah, this is what this is all it looks like.

• 39:24Nathan Labenz?

  So I often think back to that paper that was given paper of the year at one of the major conferences. Emergence is a mirage, basically, was the sort of counter narrative to the earlier narrative of emergence is a big deal. It sounds like you're on the side of emergence is a big deal, and I'm basically with you. But to try to steelman the idea of emergence as mirage, I think one of the more powerful points that they made was like, a lot of these graphs that we see are logarithmic, obviously. And so when a thing is presented as a spike, that is still often over like, actually, an order of magnitude of of data, right? You're like 10 to the whatever flops, 10 to the one, N

• 40:08Nathan Labenz?

  plus one, 10 to the N plus two flops. And you're like, oh, my, that came up fast. But it's actually, that was 10 times more than everything that came before in that in that plot. Are you seeing things emerge? Is that consistent with what you're seeing in terms of emergence? Does that like, change how you think about it at all? Or are you perhaps seeing things that are spiking even faster as you really drill into to what's happening? And I guess one more, if I could tack on even one more question there. You said the term grocking. Is that basically your mental model for this, that these things are learning the right algorithm or the right conceptual representation to make a phase change?

• 40:46Speaker 01

  Yeah, it's hard. Like, I have, everything I'm about to say is speculation, as just a very empirical one in my case. So it's definitely something to do with learning rate. So there's papers like mini CPM, where they just keep the constant learning rate for a long time. And then all of a sudden, they anneal it really fast to zero. And you get all of the quality benefit during that annealing phase at the very end, when you're annealing it very close to zero. And also many tokens. So it requires a lot of tokens. And then also something about their stage and the learning rate schedule. My sort of theory here is that the model needs to see a lot of tokens. Like, that's how I'm saying grocking, like it needs to grok. Okay, this type of type of data that requires

• 41:30Speaker 01

  me to follow instructions exists. That requires a lot of tokens for the model to learn. So that's the grocking part. In terms of learning rate, I think. It might be you need to settle more into the loss basin in order to like make these really fine grains changes to how the model views the world. It really makes me feel gross to say something that generic, but it's some sort of relationship between the two of these. But both of them, I would call it rocking together.

• 41:58Nathan Labenz?

  Yeah, the learning rate I do think is a or the learning schedule is a probably generally underappreciated part of this. It's not something I'm like super deep on either, but I've been watching for Sophia. I don't know if you've looked into that like alternative to Adam. Do you guys use that?

• 42:17Speaker 01

  We do not.

• 42:18Nathan Labenz?

  Interesting. Have you tried?

• 42:21Speaker 01

  Yes, we think that this is partially because of the learning rate schedule. So more aggressive learning rate schedules make a mirage of increased quality when you're really just playing with learning rate. So we stuck with Adam for now. We might be looking at like shampoo here soon, but not Sophia.

• 42:39Nathan Labenz?

  Okay, I want to understand that a little bit better, but just to try to set it up for you. If we envision a, of course, we're all going to do this in different ways in our own minds. For me, it's a 3D landscape, even though, of course, this is huge numbers of dimensions in reality. But I subscribe to the school of thought where we visualize things in 3D and then say, and dimensions really hard in our head to allude to the famous cartoon about that. So these are like very strange landscapes, right? Where we're seeing an observation, computing a gradient, how would we have done better on this observation? Let's move the weights a little bit in that direction and be at a slightly lower loss as a result of having made that change. And then in theory, that gives us better performance.

• 43:24Nathan Labenz?

  Although, of course, it's a super noisy thing because an update that made us work better on one particular observation might not work better on another observation. There's backing that goes into this and aggregating a bunch of observations at once and making an update. And this is basically a very black magic sort of thing still where nobody really has a great sense for what the true nature of these landscapes are or how best to navigate them. But when you talk about aggressive learning rate versus annealing, that's basically saying, how big is the step size that we're going to make, right? From one moment in weight space to the next moment, like, how big is that delta? And if you find that the value Is really happening most in this annealing phase where you're shrinking the step size.

• 44:09Nathan Labenz?

  This sort of suggests to me like a, almost like an Antarctic landscape where there's like lots of little crevices that you could be like around a good space and then you really have to drill into this very, very local space in the landscape to find something that really works well. And then that would suggest, because obviously we see a million things happening in the broader ecosystem, that would suggest that there's probably a lot of them and that's where I get to this crevices notion. But then I wonder there, I'm like, I also see papers like the platonic representation hypothesis where at the same time that's happening, there's this other line of thinking that's at scale, things seem to be converging. The basic notion of that paper, I'm sure you've at least

• 44:51Nathan Labenz?

  come across it a little bit, is with greater scale, with more modalities, it seems like the internal representations of concepts are converging in some like high level statistical sense. So how do you make sense of, can you begin to make sense of all that where there's like seemingly a lot of different like local points that you could drill into, but there's also this sense of like global convergence. Can those be compatible? Does one have to be wrong? I find myself confused.

• 45:24Speaker 01

  Um, for the pocket idea, I don't, it's definitely not a flat area with lots of pockets. If we were to go with that, it would definitely be a general slope downwards with pockets along the way. And then you can settle into one of them as you go. This is just because even in the constant phase when you're taking those big steps, people like mini CPM authors still saw improvement. Just less improvement. I definitely know that your lost landscape is much smoother for larger models. You can more richly like model the space and therefore just increasing the scale makes it much easier for your optimizer to make progress. This is why things like novel optimizers like AdaFactor or Sophia or like second order optimizers.

• 46:08Speaker 01

  It's like the bumpier the landscape is, the more accurate your step should be. So that you don't get kicked out of like the gradual slopes that you want to reach. It's also like why we're seeing. A lot of the conventional wisdom is assuming big models because big tech assumes big models. So the reason people do dense transformers with like large scale optimizers. The reason why people say, oh, shampoo doesn't scale and all this kind of thing. Uh, it's because when you're at large, your step is going to be broadly accurate anyway. When you're in this giant, smooth loss basin, like, why bother with any of these things? And I agree with them. You shouldn't bother. But when you have a three B novel architecture model with a really bumpy loss landscape,

• 46:50Speaker 01

  that accuracy really matters. That also translates into things like batch size. So I trained the Pythia models back with a Luther. And we were seeing that, for example, like large batch sizes was causing instabilities. And it really confused us for a long time because the higher your batch size is, the more accurate your gradient should be, right? So if you have these super accurate gradients, uh, this is even at the cost of seeing less data, right? As you increase batch size, you'll see less data. But now I'm, I'm fairly certain that by looking back at the data, if you have like bad samples and as we were increasing the batch size, we were also increasing the learning rate. And if, uh, there's a few, just a few outliers in that batch that sort of kick your gradient in the wrong direction.

• 47:35Speaker 01

  And then you have a large learning rate, you take a big step in that wrong direction, then larger batch sizes were actually bad because you were reaching too many of those outliers per batch. And you reached too many bad batches in a row. And now you're going completely the wrong direction. You just diverge entirely. So we actually had to use lower batch sizes, but it's complex. You also need to, this is why people clip their gradients also. Because if you take one step in the wrong direction, make sure it's not a big step. And then hope that big step can be reversed by the next couple of steps. So there's a lot of like weird alchemy going on to get models to actually converge. I've noticed that every run has its own like cursed way. And the reason people like really want to use Adam, the reason we're using Adam

• 48:17Speaker 01

  is because Adam is extremely robust. That's the main reason. So that's why people can publish papers like Sophia and AdaFactor. If you look enough, then you'll find cases where they're better. But Adam is good at everything. It's not the best at some things, but it's good. It doesn't have any like glaring weak points, like optimizers like AdaFactor and such that might just not work. And then you spend three weeks like trying to get your initial stability for this weird combination of hyperparameters and model size to actually work. That's how I look at it, though.

• 48:49Nathan Labenz?

  Somewhat novel intuitions for me, but I want to just confirm with you that they match yours. Overall, you're like, it's easier to get big models to to work. It's harder with small models. That seems to be maybe a reflection of just like how overloaded the small models are. If I think about just a super dense, everything packed in, super tight sort of network, then I would imagine that it's just easier to mess things up if I take a step in the wrong direction. Whereas if I go to a trillion parameters at the very far end of reportedly a GPT-4 type architecture, then I can spread this out over many experts and things aren't so overloaded.

• 49:32Nathan Labenz?

  And one wrong step in one place might not even affect most of the network because it's localized to experts or what have you. And so then I can be more aggressive in my step size and I can get something like a Sophia to work well in those bigger contexts because I just have more room for error essentially in my training process that I just don't have that luxury if I'm training a smaller model. Am I interpreting you correctly there?

• 49:58Speaker 01

  You're close. So like for larger models, you want actually like lower learning rates because even though it's like a smoother loss landscape, you don't need the small model ability to really jump over these crazy hurdles and stuff. But think of Adam, right? You have a variance term, you have a momentum term, and that's to tell you if there are bumps along the way downwards, then you still make progress downwards because you still have momentum. You have faith that this way is probably like the best for the loss landscape despite some bad gradients. For larger models, despite the the smooth loss landscape, you want a lower learning rate just because one of instabilities that's in general tend to come up is because you have so many parameters that go beyond the point of no return.

• 50:42Nathan Labenz?

  So I heard the Adam one, the momentum is there to say, we're not just going to look at this particular batch of updates and.

• 50:54Nathan Labenz?

  Calculate naively what would have made us better here because we know there's a lot of randomness and so we're going to keep track of the general direction we're going, which is the momentum term, to make sure that we're like not getting too randomly sent off in weird directions. We're going to try to have some, you know, general sense of what direction we're going. I'd be interested to hear the sort of similar narrative form for Sophia, if you have one there that would explain like why that's different and maybe why it's better in some places than others. And then I was also going to ask about distillation. It seems like the big trend with the frontier developers has been go huge and then distill. I wonder if that's something that you guys might also be thinking about if there's like

• 51:38Nathan Labenz?

  a Zamba 70B or even 400 and however many B in the background somewhere.

• 51:43Speaker 01

  So in terms of Sophia, I don't have a good intuition for this, actually. I've explained why Adam is great and why I think Adam is good at everything. It's like very robust, but I can't really explain why everything else is brittle in its own way. And this is not really just to say that like everything else is bad. It's just to say I haven't put the time in to any other optimizer just because I've been burned enough times by like alternative training, like strategies that just end up not being generalizable. I don't want a bunch of like a toolkit that's specialized to each training run. That's not the point because then I have to spend as much compute as a training run to try and figure out how to get this specific thing to converge. I want it to just converge. And even if it doesn't

• 52:25Speaker 01

  converge quite as well, I'll train on more tokens because I was able to kick it off day one and I'll still end up at a lower loss. On your MOE point, so MOEs are themselves like a really tricky thing to optimize because like for one, the attention block is not divided. The attention block sees everything, but for the expert blocks, you're subdividing the batch across everybody. So like they're effectively seeing like a different gradient noise scale than the attention block. So trying to find the right batch size is clearly much higher, but do you scale it by number of experts higher for the batch size because you've subdivided? That intuitively makes sense, but the problem with this is that now the attention block is getting this massive batch size. This is part of why MOEs are

• 53:10Speaker 01

  super unstable, I think, to train. Like Black Mamba was like a pain to train because of this. Attention blocks are seeing more accurate gradients is one way you can think of it than like the experts. Trying to port new transfer is something that I'm really looking into here. Actually, the people at the team at Cerebrus is really deep on this. They're really good, so we're trying to crack this, mostly them, but yeah, we're working together. But in general, trying to find for these hybrid architectures that we're seeing, for example, that's kind of true for us too. Because we have this shared attention block, which is going to have more accurate gradients and see more data. Finding something that converges is really difficult. And like being able to navigate the hyperparameter soup really determines if you live or die.

• 53:53Nathan Labenz?

  Fascinating. What about, did we get to distillation? Are you guys going to pursue a sort of big to small strategy as well?

• 54:00Speaker 01

  Yeah, yes, that is the plan. The problem here is that you need a lot of compute to train those large models. So it's something we're scaling up for now. In terms of the actual distillation step itself, unfortunately, big labs are very tight lipped about this. So like even annealing, it started emerging. Oh, annealing, by the way, it's like the continual pre-training step that we picked up early. Where you first do like a standard cosine decay, you warmth learning rates. And then you slowly decay over time to some high value, higher than normal. And that's like your phase one data. So this is like more like web data, it's a bit noisier. It's still important to have so that the model can like knows what this data looks like, knows how to respond.

• 54:43Nathan Labenz?

  Talk like a pirate.

• 54:45Speaker 01

  Yeah, exactly, exactly. But later on, you want to bake in more intensely a higher quality subset of your data. The higher quality doesn't mean textbooks again, doesn't mean you want your model to be dry. But you like, you can only look closely at so much of your data. Like everyone says, look at your data, but you can't look at three trillion tokens. You can look at like 50 billion, 100 billion tokens and apply some like really detailed filtering and deduplication to those. And make sure it's the highest quality, even if it's a distribution that represents all of your data. So anyway, you decay down to a relatively high learning rate.

• 55:23Speaker 01

  And then on your first phase, and that's where most of your tokens are coming from. And then in the second phase, you very quickly re-warm your learning rate. You say, okay, I want the model to respond really quickly again. We're in a new loss landscape. We have a new distribution of data. And then very aggressively decay down to zero. You say, okay, I want my model to respond quickly to this data. And then settle very quickly into the new loss basin that emerges from the high quality data. And also you're able to stomach this new high learning rate again, because your data is very clean. You don't have this problem I was mentioning earlier with Fittia, where you have these crazy outliers that kick your gradient in some different direction. So you're able to go to, with a model that's already settled into a good place, you're able to increase the learning rate a lot on this high quality data.

• 56:07Speaker 01

  And then really bake it in, so it's top of mind to the model. This also increases the flexibility. I see it as a branching path. The base model is the base, the trunk. And then you can instruct tune and then fine tune so that your model is becoming progressively narrower along one branch. But if you have multiple anneals, that's like multiple sub-trunks, right? So you can anneal on role play kind of data. You can anneal on, like, if you want to try, like, very factual model. You can kneel on textbooks and such. But then from there, you can fine tune so that the role play model is role played on like specific characters from the annealed role play checkpoint. And that is a much more flexible scheme than just having a single base model, a bunch of fine tunes from that base model.

• 56:51Nathan Labenz?

  Can you define anneal? I've learned this in the context of like biology and I'm looking it up online. It's also used in glass and metalwork. Maybe this is wrong, but I guess I think of it as hardening into its final shape. Or with DNA, it's actually, you know, combining, coming back together and reaching some like lower energy stable point, basically. Is that the model here too?

• 57:14Speaker 01

  Yeah, we put very detailed instructions and explanations on it's called Letter Z and then Cookbook. The Z Cookbook. If you look it up on GitHub, we have a whole section on what is annealing, what it means for us, what it means for your model. But yeah, annealing for us is like trying to harden the model a little bit, the base model into a pre-fintuned checkpoint that has some type of data that we want top of mind. And practically, this means rewarming a learning rate on the new distribution. It means having some replay, so you still have to feed some of the tokens from your batch from the phase one data so that you don't have catastrophic forgetting.

• 57:52Speaker 01

  And then that replay plus new data distribution, then you very aggressively anneal the learning rate down so that your model settles like very definitively into a loss basin for you. This all getting back to the reason I brought up annealing was you were asking about distillation. The sense of the GPT three paper, the big labs have learned be a bit less open with people because then they'll actually take compute and try and do what you're doing. That's why annealing for a long time was like a hidden trade secret. Like the Lama team cracked this with Lama two. They split into the Mistral team, they split off, they knew this sort of secret. And then it became more and more public over time and now everyone publishes it.

• 58:32Speaker 01

  Distillation has been the same sort of regime where the big labs have cracked it, but they're staying tight lipped about it so that they can try and extract all the value before their people get poached somewhere else and then it ends up being public anyway. We don't really know the secret to distillation, actually, but synthetic data from large models is a proxy for this, right? So we're trying to to look at this at least so that we can steer our annealing and post-draining strategies to to more approximate the output of large models via synthetic data. And then trying to actually use the logits during training to stabilize. So remember earlier, we were talking about can distillation help you navigate the bumpy loss landscape by having

• 59:13Speaker 01

  basically a big atom that's holding your hand and saying, okay, no, this is the right direction because this is what the big model would do. The big model is able to see much better that this is just a little bump and you should keep going down. Using those logits at training time is something that we're still trying to crack. We see generally positive things from a model convergence standpoint, but it's still really inefficient to try and have a live inferred big model at training time. So training, pre-training, scale, distillation, we have not cracked yet, but I'm confident that we will.

• 59:46Nathan Labenz?

  So the basic idea that I have of distillation is training a small model to reproduce the outputs of a big model.

• 59:55Speaker 01

  Mm hmm.

• 59:57Nathan Labenz?

  Not so much at the token level, but at the either the logits level or potentially if you had like direct access to the large model, then you could perhaps tie it like multiple places, right? You could say, I've got my trillion parameter model here. It has a hundred layers here. I've got my seven B. It's got however many layers. But I'm in math. Okay. At this point, I wanted to tie up to the, the big model and try to figure out a way to imitate the internal representations. Am I off base conceptually there at all?

• 1:00:27Speaker 01

  That's exactly right. And maybe I'll expand on what I was saying to tie it to what you are. So when, when I'm saying synthetic data, it's like a poor man's way of doing distillation. That's you're just choosing the max of the logits, right? So like you're only training on what the big model is paying attention to. And you can do that also with close source models. We could do that from GPT four or whatever else. Now, logits are important, uh, because those are like the mo, the large model saying exactly how it would respond to the current batch of data. That's why I'm saying it's kind of like holding your hand in terms of the optimizer. Logits are also super important because you also want to know what the large model is not paying attention to. So those tiny zero values in the logits, like the whole distribution helps the small model,

• 1:01:10Speaker 01

  not just the large values that, uh, you could get if you were to like approximate with just the synthetic data. But yes, you need access to the model for these logits, which is a big hurdle. It's helped now by big like Llama 3, 4, or 5B. But when you're running inference of Llama 4 or 5B every single training step, that's why I was saying that at pre-training scale, this becomes prohibitively expensive. And trying to get around that expense hurdle while still getting logits per step is what we're trying to figure out. We think there are some cheaper ways to do this, and I think we will get them figured out. But this is why it's not like popping up all over for open source people. Because it's either prohibitively expensive or synthetic data is just not like rich enough of a representation

• 1:01:54Speaker 01

  to really help you steer in the right angle.

• 1:01:57Nathan Labenz?

  And so the, just to make sure I understand what makes it prohibitively expensive. It's basically that running 405B is like some fraction as expensive as training 405B. So if you want to train your small model, but also have to run 405B, it's you're going to be spending as much compute as if you were just training a much bigger model in the first place.

• 1:02:17Speaker 01

  Exactly, yeah. In terms of practical sense, if I have a 7B, then I'm already doing parallelism across GPUs to get it to fit in vRAM. And then if I have 400 billion parameters that I also need to store, even if it's for inference, you don't need to store optimizer states, gradients, and all this for the big model. You still need to fit it somewhere. You still need to use the flops to infer it. You need to synchronize. It's a forward pass with the forward pass of your own small model. Why don't I just train like a 4B from the start? And why even deal with distillation? Is it really worth it? So far, these practical trade-offs have not been cracked for us. That's what I'm saying.

• 1:02:53Nathan Labenz?

  Yeah, gotcha.

• 1:02:54Speaker 01

  It's not practical. It's probably a slight improvement, but a lot of things are slight improvements. And just seeing more data is always a very like direct way to improve your model. And I could see more data or have a bigger, more parameters. And that's a more direct way to improve model quality right now than trying to get logits from Llama 4 or 5B. That's just the practicality.

• 1:03:14Nathan Labenz?

  Do you think it's possible that the big labs maybe don't necessarily have a huge secret, but what they maybe have is just a different cost benefit analysis where they're running so much inference that it's worth it for them to do this?

• 1:03:28Speaker 01

  Oh, I think this is likely. Yeah, this is likely. I think that if you're Microsoft working on five, then just use GPT four. It's on our own systems. We have infinite funding. So making it, there's two possibilities, right? Either they've cracked a cheaper way to do it and we have to find it, or they have more money and we have to find the cheaper way, crack it and then do it. And both of those cases require us to crack or find some cheaper way, unfortunately.

• 1:03:56Nathan Labenz?

  Are there other things that you would put in the bucket of frontier labs are sort of generally understood to have a solution that nobody else hasn't diffused to the rest of the community yet?

• 1:04:07Speaker 01

  Data is a big one. Data and data cleaning requires time, patience, and really just manpower that open source does not have. So when you're open source, then you're just beholden to whatever data set other people are able to like completely focus on. And not a lot of people focus on data. So it will always lag the big tech companies. Also, big tech has access to their own data, which is also very clean. It's in infinite abundance and we're definitely behind. So yeah, clean data, clean data pipelines. This is what we talked about. We just released Zeta two, which is our Zephyr dataset two. There's some new data sets in the open source regime called DCLM and FineWeb and Dalma and Zeta

• 1:04:47Speaker 01

  one. And cleaning them up and producing like an actual like at scale data set with Nemo data curator, which is Nvidia's like data processing tool. We found has really lowered the bar for entry for data. But yeah, data's one. Another one is things like hyper perimeter transfer. Unfortunately, has not been picked up a lot by the community. You might recognize this as muP or muTransfer. This like, I think it's modular. It's a new paper that's also found similar things. But it's like 30, 40 pages of math for these papers. And it's not approachable for open source people or small labs and such. So only big labs have the compute funding, like to put one person, okay, go read the muP paper for a month and

• 1:05:29Speaker 01

  try and port it to our ginormous monolithic code base and get it to actually work. So there's a lot of trade secrets, I would say. Distillation is another one.

• 1:05:38Nathan Labenz?

  Give us a little bit more on that one, because I honestly don't have a lot of intuition for what that. Is adding to the overall stack.

• 1:05:45Speaker 01

  Repeat?

• 1:05:45Nathan Labenz?

  Yeah, exactly.

• 1:05:46Speaker 01

  Oh, yeah. Oh, so mu transfer is a way of initialization where across it's like width scaling. So you do all of your hyperparameter search on some really low width model. It's really cheap to run training on and you do all of your learning rate sweeps. For example, of your batch size sweeps. And then if you literally just divide by a scaling factor of your low width by the high width of the actual model that you want to train, both in terms of the learning rates of your hidden layers, like your embedding scaling, you can multiply by some new hyperparameter. Then you just get a perfect transfer from the small model to the big model. And as in the learning rate that I found was optimal for my small model will also be the optimal learning rate for my big model.

• 1:06:29Speaker 01

  And this is all just making sure activations don't explode by all of these small scaling factors of width. This helps you in terms of both you're able to more accurately like search the loss landscape because the compute is so cheap on the small model that you're able to really like search a width or a higher amount, search base of learning rates and batch sizes and stuff. It also helps the just the general stability of your model because your activations aren't exploding and your model's not learning on the fly how to account for its own explosions. So you're able to also deal with a much higher learning rate. Alternative architectures and optimizers are also much easier for this because you can just directly, for example, the mixture of experts.

• 1:07:13Speaker 01

  If new transfer is cracked, then you know exactly how to ensure both the attention block and the expert blocks are both seeing relatively equal magnitude weight updates, which is what you want for stability. These sort of details tend to be hoarded in in big tech right now and slowly diffuse out to everyone else just by people going to different organizations in big tech, people publishing, people in open source kind of cracking it, that sort of thing.

• 1:07:38Nathan Labenz?

  If there's others in that category, keep running them down.

• 1:07:42Speaker 01

  It's endless, man. So this is why Eluther AI is so important. I mentioned them before. An open source collective where everything start to finish is open source. The data set, the training framework, the actual produced models. If you don't see all of it, like for example, Llama 1 was great because they shared everything. They shared all of their data pipeline and people made red pajama, which is a full reproduction. Llama 2 was they stopped sharing the data set. And as I told you before, this is because of a kneeling. It's just too tempting to keep the important parts to yourself and publish a paper that's just vague and glosses over it. But keep going, what else? So for example, the second order optimizer thing that I told you for small models specifically

• 1:08:24Speaker 01

  is definitely something hoarded. So that's like a trade secret. I mean, Gemma also did it. So they just say, we use a second order optimizer. Which one? How do you distribute it is a big one? Because if you have a second order optimizer, you have a bunch more states that you have to store in memory. And how you split those across keep views is non-trivial. So that's another engineering detail that they want other people to like catch up on while they extract the information from. So second order optimizers are small models that are distributed in a specific way. Now we're cracking it with Shampoo. And then all of these things that I mentioned all interplay. If you have a second order optimizer, how do you port mu P to it, mu transfer to it? Right? Because now it's not Adam. You have to reapply the theory. Do you spend another month trying to figure out the new theory?

• 1:09:08Speaker 01

  I would say, curiously, alternative architectures like Zamba are not hoarded by big tech because big tech saw the quality issue with pure SSMs in a lot of ways and said, okay, quality is most important to us and we have infinite inference time compute. So let's not bother with it at all and sure, we have the highest quality model. And that's more for us, right? But I think they were much slower to catch on to SSMs and alternative architectures just because it was not super clear to them at the time that you could solve this accuracy problem by some exotic architectures and by being really data driven. What else? Paralism schemes for all of these, like I was starting to mention with the optimizers,

• 1:09:49Speaker 01

  but paralism schemes for new model architectures and optimizers and combinations, those are, tend to be proprietary. So Lama shared that they had like a 4D paralism sort of scheme. But how you set up that topology on your own, that's all like empirical knowledge that you have to get from training a model. Same thing with like arcane things like checkpoint restart. So like a lot of these bigger runs, they use like dedicated nodes that are just being used to save checkpoints. And then when you have thousands of GPUs, one of them dies all the time. So you have to keep slotting in the checkpoints from those backup nodes. That's something that's really easy conceptually, but hard to set up and in some sort of proprietary big tech stack. If everything is what I'm getting to, there's a lot of points, especially about your gritty

• 1:10:34Speaker 01

  details, engineering details, scalability details that are very closely kept tied to whatever organization has them. And they also know that even if you say worked on one of those big tech teams and you know the secret and you go somewhere else, like porting it to their stack is a very non-trivial thing, right? It's completely built off of the proprietary stack of whatever company came up with the idea. So do you remember it well enough to port it to this new stack? Uh, maybe not. So that's really a hardening of ideas and training processes that are exist in each organization and that is increasing over time.

• 1:11:10Nathan Labenz?

  So I think we have a few more technical questions. One obvious kind of pulling up question for a second would be like, the story you're telling sounds pretty consistent with the allegedly leaked Anthropic fundraising deck in which they said that the leading companies in 25, 26 might get so far ahead that nobody can catch up. And I think we've understood that, or I've understood that broadly as sort of a accumulation of these secrets. And then also just like having the best models. We now have this, I think, pretty credible reporting that OpenAI is using O1 to generate the data for the next, you know, scale up again. And that seems like, you know, there's not just a moat, but there's like multiple moats in depth there.

• 1:11:56Speaker 01

  Yeah.

• 1:11:56Nathan Labenz?

  And yet you guys at Zyphra are like a small company and trying to figure out how to compete in this space. So what is your sense of how you and other companies that are not infinitely resourced, even if you're like reasonably well resourced? How can you compete? Is it about finding a niche that they're not going to play in? Or is it about some sort of breakthrough insight that is going to be needed? How do they not just come to dominate the entire space?

• 1:12:27Speaker 01

  Yeah, I will say small models is overlooked by big companies because like OpenAI, XAI, Anthropic, all of them, their job is producing the highest quality model. And yeah, sure, they wanted to still into a smaller model for inference reasons. But if it came down to cracking distillation or making an even better, bigger model, that is, it only needs to be a little bit better. Claude Sonnet 3.5 only had to be a little bit better than ChatGPT to take a large portion of market share. So accuracy really takes them. It's also why they don't want to like experiment too much with these weird hybrid architectures, because if it affects model quality and you only determine that when you train a big model, it's a lot of cost when you know the dense transformers work the way that you expect.

• 1:13:12Speaker 01

  Yeah, in terms of how we can compete with them, I would say it's a mixture of this. They're focused on large scale inference, large scale training. Definitely makes them, even if they have things like distillation, that's only one part of the process, right? So like we had to crack annealing, we had to crack architecture, we had to crack data, the entire like training process from start to finish. It wasn't just like the distillation step, which is only helped when you're at a big company, right? We're very nimble in that way. We're definitely in a unique time where if you're nimble enough and you're focused on small models, like you can very clearly have an advantage over big tech. Over time, as more of them use distillation such that if we were to not move at all,

• 1:13:54Speaker 01

  then we would definitely vaporize compared to their resources. So the way we're staying ahead is by continually pushing onto what the next free lunch is for small models. Uh, next like second order optimizer is is one. Getting a cheap way to do distillation is one that will protect us from that angle. Like we're clearly very willing to make big bets on the alternative ways of training.

• 1:14:18Speaker 01

  And big tech just is not really looking very closely at this sort of regime. And that's why we have the best model now with seven B is the best because of this exact reason.

• 1:14:27Nathan Labenz?

  Are you studying the entropics project that has recently been catching a lot of buzz or similar?

• 1:14:34Speaker 01

  Very loosely, I'm not an expert. I would kind of categorize it under the definitely doesn't scale. So I believe it was Tim Detmers tried this at scale with like larger parameter counts. And it was like no better than Adam. I know the thing is that I kind of view it as kind of specific. So whether it's robust or not, or whether it's very specific to the training setup on like the GPT is yet to be determined. It's true with this for a lot of things. You can show like huge speed ups on even though, what is it like crowd sourced compute, like pre-training compute? But it's hard to show things that are robust. Like that is the actual test. Is this robust? And I'm waiting for someone else to spend the compute to determine if this entropics is robust.

• 1:15:18Nathan Labenz?

  You're describing it in a different way than I expected.

• 1:15:21Speaker 01

  Okay.

• 1:15:22Nathan Labenz?

  Uh, the one that I have been at least momentarily fascinated by recently is essentially a improved sampler. So it is like trying to differentiate between like when the model is confident and should just like pick the most likely next token and everything's kind of hunky dory versus when it's out of distribution. And then when it might need to like, it does interesting things like kind of dynamically inject a additional clarifying question. If it detects that like generally we're confused, let's maybe like just burn some more cycles. How can we do that? Let's artificially input a clarifying question. And then sometimes it'll even like back up and say, we seem to have kind of lost the thread here. Let's like go back to an earlier point and rerun inference again.

• 1:16:08Nathan Labenz?

  So I've thought about this entirely as an inference thing. You were going toward the training side. Are we talking about the same project? And if so, I missed the training side of it.

• 1:16:17Speaker 01

  You're making me question myself. Maybe I'm confusing it.

• 1:16:20Nathan Labenz?

  Well, for what it's worth in terms of like a possible next free lunch. The one thing you did say that matched my intuition, even though I've read like a little bit that like it might not be working as well at larger scale models. I wouldn't say that's like well established in my mind, but if there was a story as to why it might be, it would be like the larger models are maybe like.

• 1:16:43Nathan Labenz?

  Grocking into some version of that on their own through all the, especially through all the reinforcement learning that they're doing in late-stage training, whereas the small models are maybe not getting there. And so layering on these sampling heuristics could be like just much more beneficial at a small scale. This naturally has me thinking too, what about learning a sampling function as part of the overall architecture? That feels like a candidate for, if any of this is real, I haven't quite got to the point where I'm like convinced that there's a real there, but it seems like there is. Perhaps like learning a sampling function could be another relatively like small amount of parameters that could lead to some like qualitative change.

• 1:17:27Speaker 01

  I agree, it definitely could. I also agreed on the point you made about like larger models doing this on their own. Larger models, for example, data quality matters a little less because larger models learn to self-select what is important, what's not. There's a reason that like when you go to a large enough model scale, you just scale it up and just churn it out, churn out more and more data on more and more parameters. And the model kind of is so robust that it's able to learn these properties that we would need explicit training properties and processes for at the smaller scale. So that is one thing that big tech kind of lies on as a crutch, that we are able to find the explicit process that is able to emulate the same behavior from the model a bit faster because they

• 1:18:09Speaker 01

  can just rely on the model figuring it out on its own. But yeah, I do agree that it could help us, but we haven't looked too deeply into this just yet.

• 1:18:18Nathan Labenz?

  I'll ask one more and then Jason, let's come back to you and you can map some of this stuff onto the specifics of Zamba and state space hybrids in particular. How would you describe like where we are in this overall process? We always hear we're still so early.

• 1:18:37Nathan Labenz?

  That basically seemed, that would be like my vibe level takeaway from your comments. There's a lot still that maybe hasn't even been properly investigated, let alone understood.

• 1:18:48Speaker 01

  I would agree with this. I definitely don't think this is a solved problem. I don't think we've found the limits of scaling. We clearly have not found the limits of architecture. Whether there's some groundbreaking thing tomorrow or like Attention 2 comes out and everything is thrown away. I can't really say that, but I can definitely say that we've moved very quickly into a search space that's massive. There's a lot of very low hanging fruit that is just obvious, like it's everywhere that we have not picked up. And these training processes and stability, like there, there's still so much out there. Even if you were to just take pre-training model papers and compare them, like I'm sure there's a lot of things that people are reinventing.

• 1:19:28Speaker 01

  There's a lot of trends like across scale and training process that are to be found, that people have like implicitly found, but either not known or purposely not shared with everyone else that are just lying around.

• 1:19:40Nathan Labenz?

  That's my sense too. It doesn't seem like we're at the end of this cycle or anywhere particularly close, which is pretty crazy to contemplate because it's already come pretty far.

• 1:19:51Speaker 01

  Yeah, one really low hanging fruit example of this is like the GPT-3 paper had like really specific like hidden dimensions. And this was both true in terms of quality, but it was also in terms of like kernel paths. So for example, there's like lots of powers of two in their hidden dimensions, but not quite enough. And everyone else copied the GPT-3 model configurations for a long time without understanding that, oh, these are more efficient on GPUs in addition to be approximately the right size. So this is another thing we bake into our Zamba models is that since we're trying to make the most inference efficient thing on any parallel hardware, if you do very simple tricks, put lots of powers of two in your hidden dimension, like your sizing of these different blocks,

• 1:20:31Speaker 01

  just by rounding very slightly to good kernel sizes will give you speed ups on any parallel hardware. So whether it's multi-core CPU on a phone or something like an MPU on Intel, like an M1 for Apple or something. If you have parallelizable model sizes, then for the Lifetime, the entire lifetime of the model, you're more efficient. And this is something that big tech just says Llama 3 has this size. But they don't show you that they went per block and found the efficient sizes for kernels for each of these. And then baked it into the model. Same with the vocab size, the vocabs are all rounded to some factor of 64. For the same reason, like there are just more efficient kernel paths you can go down with that.

• 1:21:11Nathan Labenz?

  And my very simple like intuition for that would be just like if you want to break something out into two, four, or eight GPUs, then if you have a ninth thing, then you're probably like, you could, if you go eight wide, the difference between eight and nine is one and two in terms of how long you have to wait. Are there like a lot of different versions of that you see? Or is it really pretty much that simple?

• 1:21:39Speaker 01

  It's not that simple. So you can decide how baked in your model size is going to be. So for example, this, I mentioned a few times, there's powers of two in your hidden dimension. This is true for any parallel hardware. Like if you train on AMD MI 300X GPU and you have lots of powers of two, it's probably also going to be efficient on Nvidia H100, it's going to be efficient on your phone. But you can get really deep with this. There's something called wave quantization, where let's say you have, I don't know, 100 scaling multi processors on H100. If you have 101 units of work, then you need to do two time steps, right? First one with 100 and the second one with one. And the throughput will be roughly half because that second time step is pretty much empty.

• 1:22:22Speaker 01

  About the same amount of work, but two total time steps throughput is half. Whereas if you had 100 units of work, you just have to do one time step. So then on your graph of throughput on the Y axis and the problem size on the X axis, you're going to get waves. And each time you kick down on the wave, like you have lower throughput, you are on like a new time step. This is baked into your models. If I was Nvidia and I wanted to train a model that's only efficient on Nvidia hardware, I would bake in that this model is, has very specific sizes to fit on the number of CPU multi-processors on H100, has the amount of SRAM, whereas on an AMD like competitor, like an AMD MI300X, on that one, like it'll require two waves

• 1:23:04Speaker 01

  or something like that, or SRAM will be like undersaturated with the model sizes that I've put in. And like the bad version of this feature is that everyone is building their own chips and everyone is building their own models, then everyone would produce a model that's only efficient on their own chip. And it's like specifically inefficient for everyone else. So like, there's all, it's unclear whether we'll go this direction of like perfect specialization per company, or if companies will say, okay, I can't crack all of model training and I need my hardware to be able to run competitor models that we like run on our own data. So maybe in that sort of case, you would want everyone, like the good future of everyone provides similar hardware that can run everyone else. It's unclear which way this goes, but um, this is just another example of like, there's

• 1:23:48Speaker 01

  a lot of low hanging fruit for one. And for two, there's a lot of return that you can get by just like getting out a magnifying glass and looking at every single step of your training process. So like extract everything you can from it.

• 1:23:59Nathan Labenz?

  Yeah. You follow Greg Brockman on Twitter and you realize how much of that he has personally done over the last couple of years at OpenAI and you're like, yeah, there's a reason that they're obviously extremely good at what they do. And you can see that the sort of the out, rough outlines of how that work is unfolding there in, in his Twitter feed in particular. So recommend that as a not very open window into what's going on, but at least a tiny little peek.

• 1:24:27Speaker 01

  Yeah.

• 1:24:29Nathan Labenz?

  Jason, what's on your mind?

• 1:24:30Jason Meaux?

  Yes, just a pleasure. A super interesting conversation to hear. The ins and outs, the inside scoop on training. I guess I would really be remiss not to do a couple of deep dives with Quentin while we have him on Zamba 1, Zamba 2. Maybe let's just start with a deep dive on Zamba 1, if that's okay. And in particular, you look at the architecture, seems pretty straightforward. There are six Mamba layers to start off that the input sequence goes into. And then this interesting thing happens that I'd like to talk about. So there's this step before the attention block of concatenating the residual stream coming out of the Mamba block with the original input embedding. So this is very interesting.

• 1:25:14Jason Meaux?

  So the attention block gets to see the Mamba residual stream plus completely unmodified input embedding. It feels right that if you're running the attention computation, you'd want to do this. But at the same time, it's interesting. How does one arrive at a feature like that? Is this something that your team is, okay, it's intuitive, so let's try it? Or is this something that through many, many different experiments that just rises to the top and it's more experimentally found?

• 1:25:40Speaker 01

  Okay, on the specific concat, if I recall, we did run a lot of experiments, but it also was intuitively, right? Our experiments matched our intuition, which is where we want to be, that we want the attention block to be able to see like the entire input and residual. There's really not much more to that. Empirically works and intuitively makes sense, so we stuck with it. In terms of Zamba 1 in general, we kind of crystallized more and more of our thoughts on how a inference optimal model should look like. So that's when we discovered that, yes, indeed, attention blocks are correlated. Yes, indeed, you can replace them with just one and get pretty much all of the same benefits of attention. Yes, indeed, hybrid models with a rough six Mamba blocks to one attention block was like a good split.

• 1:26:25Speaker 01

  Let's see, other things. Everyone likes to talk about model architecture, but it's one step of the three holy, the holy trinity, right? So it's model architecture is a big one and we really optimized that for inference efficiency and like performance per parameter, as we discussed earlier. So number two is like training framework and process. So we learned from the mistakes of Black Mamba of like how to affect for optimizers and stability, how to tune hyperparameters for exotic model architectures. We baked all those in. And then number three is data. So Zamba 1 also came out with ZYTA 1, which was our training dataset at the time. And really getting quality baked in at scale for this size of model was really important.

• 1:27:09Speaker 01

  We have some nice graphs where like at 1 trillion tokens, we're able to perform similar to models trained on 5 trillion, 10 trillion, 15 trillion. If you really optimize for quality at this smaller model scale, it really has effects. And Nealing was first found with Zamba 1. So this process of two-phase training, we discovered it with this. So this is like the first pass of everything in our overall training process. Any thoughts or anything you want me to drill down on more?

• 1:27:36Jason Meaux?

  Yeah, I guess if we could just walk through the the layers, they're pretty straightforward. Six Mamba layers. And then just a follow-up question on the on concatenation. So I guess first thing is because you're using concatenation. So you have to pay for that computationally. I guess the dimension is expanding. Obviously, the trade-off's worth it if you're getting a bump in performance. And then I guess at some point, it's not entirely clear. At some point, I guess you're projecting back to the original dimension. So if you could just step through, imagine there was a visual on the screen, step through the Zamba 1 architecture from beginning in.

• 1:28:13Speaker 01

  So yeah, no positional encodings. Then you have an embedding block. And then you have this alternating between Mamba 1 blocks because it was only Mamba 1 at the time was out. So six Mamba 1 blocks, global attention plus MLP. So there's also an MLP in that global block. That's important for people to know. Six Mamba 1s, attention plus MLP. Six Mamba 1s, then the same attention MLP together. In terms of the concat, yes. Empirically, it is better to do a concat than to try and like, I think we tried to like project it back down and not eat that computational cost. Nothing really worked very well, so it's very empirical.

• 1:28:52Jason Meaux?

  So I guess in the literature, there's a range of results of Mamba 2 attention layers. Zamba lands on 6:1. That's consistent with Zamba 2. So it seems that don't fix what's already working. But some of the literature suggests wide ranges up to 7:1, 10:1. How did the Zephyr team arrive at the 6:1 ratio? Was that experimental?

• 1:29:16Speaker 01

  That was purely experimental. I also suspect that we have to apply our attention block more often because it's shared. So there's some sort of interplay there. Yeah, I mean, 7 in 5 are slightly worse. That's really all it is for us. I think it's similar to everyone else. Zamba, mixture of experts, clearly has some sort of interplay with how often you apply attention. I'm sure there's some, there's some consistent rules here that are being applied to everyone. But the way it's manifesting is we all are empiric. We don't have time to make these consistent rules play nicely. So empirically, we find that, yeah, for us, it's every 6.

• 1:29:52Jason Meaux?

  Excellent. So I guess, jumping from Zamba 1 to Zamba 2. So the Mamba 2 paper comes out in May. Both of the original authors published that, that paper, Albert Gu and Tree Dao. The main result being the SSD algorithm is adding a restriction, adding structure to the A matrix. That's the part of the algorithm that influences how the state changes over time. And so now we can do matrix multiplication and take advantage of modern GPUs, get those tensor cores going. So now we can train much faster and as a consequence, much larger state sizes. When that paper came out, was this immediately on the minds of the team

• 1:30:34Jason Meaux?

  at Zephyr? The kind of impact it could have? Was Zamba 2 immediately planned? And what were the features of the Mamba 2, the SSD, that gave your team confidence it was worth coming up with the follow-up model?

• 1:30:47Speaker 01

  Yeah, we, after Mamba 1, like, we evaluated a lot of architectures at the time. And since Mamba 1 gained our confidence initially, like, Mamba 2 was something that we initially, immediately wanted to test out. We ran a lot of oblations of, like, pure Mamba, we found, oh, pure Mamba 2, that is, we found, okay, it has the same problems with, like, in-context learning, long sequence dependencies. So we can't just make a pure Mamba 2 model, still. But we also were able to verify that we get significantly higher throughput on H100s with our training hardware for no model quality degradation compared to our initial Mamba 1 tests. We also found something a little different from Jamba. So Jamba 1.5, or it was either 1.5 or the initial 1, I don't remember which, but they found Mamba 1 was better in a hybrid case.

• 1:31:31Speaker 01

  So they did, like, Mamba 1 plus attention. We didn't find that, so we found they really performed the same. So we stuck with Mamba 2 plus attention. But yeah, the Mamba 1 to Mamba 2 change in, from Zamba 1 to Zamba 2, is literally just throughput. So we get much faster models, both the training and the print time.

• 1:31:51Nathan Labenz?

  I'm not sure if this is the best time to interject, but I also, my ears perked up around no positional embeddings. And I was wondering how you understand that working? Is that an artifact of the, or a sort of consequence of the inherently sequential nature of the Mamba blocks that they create?

• 1:32:12Speaker 01

  Right.

• 1:32:13Nathan Labenz?

  Sequential understanding where the transformer doesn't have it natively, or the attention mechanism doesn't have it as natively?

• 1:32:20Speaker 01

  That is our exact intuition, yeah. So if you, we found that with and without them, with Mamba, wasn't much of a performance difference for the first model, so we're just, okay, without them. For the second model, for Zamba 2, we actually put rotary positional embeddings in. This is because we found there are some cases where you can make rotary help you a little bit, so we kind of had to figure out how to ask the right questions. If I recall correctly, rotary helps us a bit with long-range dependencies.

• 1:32:50Speaker 01

  But it hurts us in terms of context length extension. Rotary is very inflexible with its context length. And so that definitely did bite us. In the future, whether we continue with Rotary, whether we get rid of embedding or positional embeddings again, whether we go with some like alibi or learn positional embeddings instead, I'm very uncertain with which one we do. But I will say that they're all very slight effects. And this is because you correctly state that Mamba kind of handles a lot of the position, like already encodes the position. So like you don't live and die by the positional embeddings, that is.

• 1:33:24Nathan Labenz?

  Yeah. Context extension is another thing that I think is really interesting. And that was one of the biggest things that kind of inspired me about the Mamba architecture in the first place. At the simplest level, the state size doesn't grow with sequence. So that means we can, in theory, have an arbitrarily long sequence and have some sort of more, I don't want to say human-like, but on the grand scale of possibility space of memory, at least one step toward more human-like in the sense that obviously my memory, our memories are not all tokens to all tokens, but rather some sort of lossy compression of everything that came before. And I wonder, what would be your outlook for that now?

• 1:34:06Nathan Labenz?

  We've been almost a year since the original Mamba paper. I would say it's definitely shown itself to be a big deal, but I haven't seen like, outside of maybe Magic, I don't know what they're doing, but I haven't seen like many millions of tokens contributing to a single long-lived state. You talk about like on-device, constant, continued pre-training, like learning about me, baking that into the weights. That sort of suggests a disbelief in the idea that we can compress my full history into some long-lived state.

• 1:34:46Nathan Labenz?

  So what would be your outlook for training, not training necessarily, but running my whole Gmail history through a Zamba 3, let's say, that, you know, that, do we have a line of sight to something that could handle 100 million tokens and have this holistic, fuzzy sense of who I am, what I care about, what I, how I respond to things? Or does that seem just still like too far to envision a path to?

• 1:35:12Speaker 01

  Yes, in terms of engineering, there's some work to be done. For example, we have ring attention and tree attention for attention blocks for these million plus context lengths. We don't really have that for Mamba. And this is just because the memory of training time is not high enough to support million context Mamba, right? You need some sort of sequence parallelism for the Mamba blocks that just doesn't exist right yet, right now. And we're working on it. We hopefully will be the first. In terms of from the modeling angle, one thing I'll say is that I don't really want my models to be brain to model. I forget why I'm walking into a room. Like, I can't reverse a string or anything else. Like, the cross sequence dependencies of attention are more powerful than the way your brain works in a good way. So I still want them to stick around.

• 1:35:57Speaker 01

  In terms of how model or how Mamba behaves as super long context, I don't think anyone knows this. I agree that it's intuitively makes sense, but I, I don't have a specific bet either way whether Mamba will, will, will do super well or super badly. Because there's something also to be said about, about losing cross sequence dependencies for Mamba at, at these really large contexts that are definitely noticeable right now. And whether attention will make up for those sins at even larger scales, I don't know. We might need more attention, for example, if we want a million contexts. So I want it to be there, but whether, well, I think there's a lot of challenges that need to be figured out first.

• 1:36:37Jason Meaux?

  Yeah. Yeah, I'll just pick up on the Okay. On one detail you mentioned, uh, having tried to get longer sequence lengths trained into the original Mamba. You know, you reach a point, you go through all the steps that you normally would to do long sequences and to save memory. And then you memory out and then the next big leap you'd have to do is figure out sequence parallelism. Much less straightforward, it seems like for the original Mamba. Actually beat my head against it. I was like, okay, cool. I'll do a pull request. I think there's for deep speed, there's the Ulysses library. I don't know if you're familiar with it. Okay, cool. Let's get a pull request going. Could not even crack 1% of what it would have taken.

• 1:37:18Jason Meaux?

  With Mamba 2, it seems a little bit more tractable. They even mention it, I think, in the paper. That it's going to be much more transformer like, I mean, any insights to, are we going to get the same kind of sequence parallels that we're used to working with transformers?

• 1:37:36Speaker 01

  It will look similar and I don't want to say exactly how because I want to be the first. I do think this has an interesting implication for hybrids, though. As I was mentioning before, whether attention can make up for it. I see a mixture of depths on the attention block being the best way to move forward where the model decides how much attention it actually needs at inference time and also a training time. So like, maybe if you want a 2 million context of all of Nathan's emails for his entire life, then maybe the model wants attention every single block. Can we also make the model decide like how many LoRs or what like how aggressive the LoRs need to be at every block? I feel like dynamic attention is what I'm getting at. It's gonna be a big determining factor on how long context for Mamba really will work.

• 1:38:21Speaker 01

  But yeah, in terms of SP, it's unknown. It's still still being figured out. Mamba 2 definitely made things easier, I will say that.

• 1:38:27Jason Meaux?

  Excellent, looking forward to that. Yeah, no worries. I understand you wouldn't want to share details at this time.

• 1:38:32Nathan Labenz?

  For anybody who might be a little confused about the whole sequence parallelism thing, can you maybe just give a little bit more intuition for the problem and how this is even conceptually possible to solve? Because again, a naive understanding would be like the Mamba block is sequential. So how are we ever going to get around that? Isn't it generating one token at a time? And like, how do we change the perspective or change the frame on the computation to go from the naive understanding to at least opening up the space of possibility for parallelism in that respect?

• 1:39:08Speaker 01

  Sure, so like sequence parallelism is purely a necessary evil to get around memory constraints. So even though Mamba is generating one token at a time and needs to store somewhere the activations and just the input sequence of a million tokens. And that is, even if you only require a constant hidden state, yeah, activations plus the actual input sequence are just too large to fit on even an H100 at training time. Training time is even worse because you have gradients and optimizer states. Those have to be in higher precision. For example, you can't fit it. And if you don't train on a longer sequence length, then your model doesn't know how to generalize to a longer sequence length. Then you have accuracy problems, so. At training time, there will definitely be less sequence parallelism needed just because

• 1:39:53Speaker 01

  there's less pressure on your device memory. But at training time, those are the reasons why we need it. How it actually looks is the thing that like we're still cracking, but that's the reason we require memory. That's all.

• 1:40:05Nathan Labenz?

  Yeah, okay. We will stay tuned for Zamba 3.

• 1:40:10Speaker 01

  Absolutely. Yeah.

• 1:40:12Jason Meaux?

  Yeah, I guess just one last breakdown. If you could go step by step, the Zamba 2 architecture, it's not just, hey, let's swap in the SSD algorithm rather than the original SSM and let's call it a day. You're making these other changes. It makes sense that if you're going to make the new architecture, hey, let's update a few other things. If you could talk about exactly what you've updated, the LoRa modules on the MLPs and then what you've done with attention and collectively all of that together. I think the highlight result from the Zamba 2 release is it's sort of at this pareto frontier of efficiency

• 1:40:53Jason Meaux?

  and performance in certain respects against the other open models. So with your intuition, you already said it. I think you had three trillion training tokens rather than just one. So a lot goes into this performance booth. But yeah, as well as architecturally, what do you feel is giving the uplift, the most significant uplift for Zamba 2?

• 1:41:14Speaker 01

  Yeah. Okay, so it's different per model. So let's start with the 1.2B, let's say. So the 1.2B says there is still this global transformer block of attention plus MLP. But if you only have one block, there is correlations across depth, but maybe we want a little bit of more representation. So that's why we have LoRa's on both the MLP and the attention block. And these do different things. So remember, like, attention is correlated across depth, but MLPs are not. So we wanted that MLP to do a bit more heavy lifting. We thought there was still room for it to specialize across depth, which is why we introduced that LoRa. I think that one does a bit more heavy lifting than the attention LoRa, which is

• 1:41:54Speaker 01

  why we only have the MLP LoRa on the 2.7B and the 7B Zamba 2 models.

• 1:42:02Speaker 01

  We stick with a single global block instead of the A B A B for the 1.2B just because we were parameter constrained. We thought those parameters were better used on Mamba blocks. That's the 1.2B. So that's the 1.2B is really good if you have one specific thing you want the model to do. 1.2B is enough for summarization on its own and you want it to just to four bit quantize and you want it to be like very low overhead on any Raspberry Pi or Jetson Nano in the world. The 2.7B has this A B A B, so there's two global blocks. Only the lore is on the MLP for the reason I mentioned, but it specializes more across depth and we wanted to kind of like for a very small overhead and flops and params, we could get a bit more expressivity.

• 1:42:47Speaker 01

  That was also when we put in the rotary, the positional encodings. That was just because we got like the slight boost, like I was mentioning along context. Whether it was a mistake or not in terms of context length extension is still to be determined. And that was the main change though for the 2.7B and then the 7B is the exact same architecture as the 2.7B. So we just scaled that one up. But the main differences from Zamba 1 to Zamba 2, if I were to summarize it, is that we were trying to really investigate the depth relationship with the shared block and try and find out how much do you want the model to be able to be flexible across depth? What is actually being flexible as the MLP or the attention?

• 1:43:28Speaker 01

  And then the other two pillars of our training stack, right? Like our data, we improved a lot and then scaled it way up. So we doubled the size of our annealing set. We tripled the size of our phase one. This helps the model like be more, be more broad in what it's good at. So first model was a bit brittle in terms of what like the trunk I was talking about was not super wide. But on three trillion tokens, like we can anneal and fine tune the model to be much more flexible across a lot of different tasks that our customers and partners wanted. Cause just because it's seen more data, like if you could look at this actually in the weights themselves, if you look at the Lama 3 weights, you can see they're much more evenly distributed across the representation. You filled out that weight more, you can, and you can see it.

• 1:44:09Speaker 01

  And going from one trillion to three trillion, we were able to get most of the way there as well. The, the model's much more flexible. But yeah, just improving end to end the stack of our model production, but with a focus on depth was Zamba 1 to Zamba 2.

• 1:44:25Jason Meaux?

  Yeah, and I think you mentioned on the first read of the blog post, I didn't see the state size of the Mamba blocks in Zamba 2.

• 1:44:33Speaker 01

  64. So yeah, all of those details are in our hug and face.

• 1:44:37Jason Meaux?

  Excellent. Okay. Yeah, I didn't check out the model card yet. So 64, that's interesting. So you could have trained larger state sizes. Of course, inference efficiency being governed by state size. Could you talk through experimentally how your team landed at 64? And is it actually a little bit lower than my expectation because the attention blocks are doing so much heavy lifting? There's less pressure for the state size to be larger for the sake of memory.

• 1:45:06Speaker 01

  This is pretty much right. So if I recall, larger states by larger state sizes is mostly for improving your long context ability. For like you can more accurately like store longer range states and we have attention. So that's it. It's a less clear benefit for us when we try to increase it. We've tested this, but it's a very clear damage at inference time. So 64, we found was enough for our Zamba based models. Yeah, just empirically. And this intuitively makes sense for the reasons I just mentioned.

• 1:45:37Jason Meaux?

  Yeah, fascinating. That makes total sense. I guess if we could just go deep on, I haven't seen that many papers or architectures where we're applying LoRa projectors to MLP blocks. It like intuitively makes sense that wow, okay. A really certainly memory efficient way because you're not storing a whole nother set of weights to add some specialization into the MLP blocks. It's sort of like is in alignment with that sort of performance per parameter as you know, idea. I guess if you could just talk about. That idea, what you've seen like in the appliations or the experiments you've run, it's a very interesting feature.

• 1:46:19Speaker 01

  Yeah, it's going to boil down to what I was mentioning earlier of like, the MLPs specifically definitely are not correlated across depth very strongly. So this even boils down to what we've learned with CNNs where like, the first few layers learn very general representations. And then the deeper layers learn more specific representations in the CNN world. This is a face, this face is a child, this child is crying. You get more and more abstract and specific as you go deeper. And in the language world, this is all manifested in the MLPs. So it's important that the MLPs learn the same notes, textualize themselves across depth. So that's why attention is learning that anyway.

• 1:47:00Speaker 01

  Oh, remind the MLPs that these are all the notes that we have across all tokens of the sequence. But those MLPs are like learning different deep relationships across depth. But all of this to say, that is the intuitive reason why we have MLPs or why the MLP that we were carrying around anyway in this shared block should probably have some expressivity across depth. And a Laura is a good way to get some of that expression, but with lower memory and compute costs, which is really what we want. If I was not bound by memory and compute, then I would just probably have independent MLPs and a shared attention. But this is a good middle ground. And then empirically it improved accuracy across depth. So we stuck with it.

• 1:47:42Jason Meaux?

  And is it the kind of thing where you've you also ran the experiment with independent MLPs? Or you just know that's not feasible for the constraints you want for the Zamba model? So you don't even bother running that appliation?

• 1:47:56Speaker 01

  Yeah. Unfortunately, that's just not on the table. Like that model is not in the cards for us. So I would hope someone does it. But yeah, ideally, I'm sure we would get some pretty big boosts from doing that. And if I had more parameter budget, I would go to full MLPs. Yeah.

• 1:48:13Jason Meaux?

  Yeah, makes sense.

• 1:48:14Speaker 01

  But the specific application of no MLP Laura to MLP Laura is worth the the jump in params and flops for sure in terms of loss.

• 1:48:24Jason Meaux?

  Very good. Interesting. I guess a little bit more of a general question is. If you read the literature, it seems like almost every time someone tries some type of Mamba Attention Hybrid, it's like this unreasonable effectiveness, at least across the metrics that they're measuring. We talked about how loss can be misleading. That's not always a measure of how well that model will actually perform, but it maybe tells you at least something. And then you can go down the list. You can look at other evaluations. It's not surprising to see that a hybrid, Mamba hybrid model would beat Mamba by itself. If you think about the limitations of a fixed state size, the way information would propagate

• 1:49:09Jason Meaux?

  through that network. But it's always surprising to me when I see all other factors held equal, that it's beating something that's like a full transform or full attention blocks. Do you have any insight around that? How these dynamics are playing together? Almost like, maybe it's greater than the sum of its parts?

• 1:49:28Speaker 01

  I really think it's powerful when you're doing hybrids to boil down, to almost apply qualitative, like track qualitative differences in how well your model is encoding cross-sequence dependencies and projecting them to mix heads and to find deeper relationships from those notes of the cross-sequence dependencies. Like, Mamba does both. MLPs just does taking the notes and finding deeper relationships, like they're projecting onto some abstract feature space. And then attention is just finding the dependencies themselves and not doing much else with them. At least in terms of how we approach things at Scifra, it's really just boiling down to this critical intuition of things. And for example, that's

• 1:50:09Speaker 01

  how we can explain why Black Mamba had struggled with the long sequence, because Mamba notes on themselves about sequence dependencies are not rich enough. So you need some attention. The unreasonable effectiveness comes from reaching the critical point of where everything is balanced, where you are not spending all of your flops in memory and everything on taking notes and then doing nothing with them, like the tiny MLPs or something and big attentions. It's not with like a weaker cross-sequence relationship from Mamba and then big MLP projections where you're like trying to find really deep things about very minimal notes. I really just, if you find a balance of these two, that's when you get really strong model architectures. And eventually, I think, we'll probably converge to like

• 1:50:53Speaker 01

  a better way to do both. Where maybe we boil down attention to like, maybe you can have a retrieval head on a Mamba block, whatever that may look like, to help you patch the problem with long sequence dependencies. Maybe there's some unified block that finds the exact mixture of projections and like note taking across cross sequence dependencies, and you just scale it up and you're done. That seems to be the case, right? Based on this intuition that you just need a balance. But yeah, I would still tie it back to you need balance between these two effects, and these two effects are what you need for a language modeling. It differs per modality.

• 1:51:28Jason Meaux?

  Yeah, very interesting. Do you get the sense, Nathan alluded to this, we see several iterations of Mamba hybrids out there, but at the same time,

• 1:51:41Jason Meaux?

  it seems like people are, like there's some lock-in maybe of even some of the private labs. Are there lock-in effects once you commit to a certain type of architecture where even though something begins to be promising, you almost ignore it purposefully to some degree, just because you've invested so many resources in a particular path? And do those lock-in effects make it more unlikely that a lab that's been committed to a completely different tech tree for years is going to just jump ship and start? Are those lock-in effects there?

• 1:52:17Speaker 01

  The lock-in effects are stronger here than in most places. One thing is that as your lab continues, you have more and more ablations that you can look back on. And when you jump ship to a new architecture, all of those ablations, you have to retake a bunch of stuff, right? Another thing is that we were talking about evaluations. It's really hard to evaluate specific tasks on very few tokens, meaning that, for example, if you have two architectures, how are you going to test how well they do on MMLU, on in-context learning? You have to train for hundreds of billions of tokens, remember? And I can't just take a pre-trained checkpoint and then test both architectures because you have to train these architectures from scratch. So to predict how well a model is going to

• 1:52:59Speaker 01

  do on those emergent evaluations is a real leap of faith on new architectures that a lot of people in labs just don't want to do. Like I mentioned, Llama CPP and the Ollama, like the whole hacker space and all of the, every serving framework is built around transformers. It's slowly changing to support hybrids and SSMs and such, but do I really want to relearn fine tuning? Do I really want to relearn efficient CPU inference? Or like, what efficient set of fused kernels like will work at inference time? Do I really want to get it in VLM and TensorRT? Like, we're going through that pain, but I can definitely see why a big company with even more technical debt and more ablations.

• 1:53:42Speaker 01

  And interpretability is another one. Like, big labs have really finely honed and interpreted, like especially at somewhere like Anthropic or OpenAI, they can interpret really well what these models are doing out of the hood, way beyond my intuitions that I'm, that I have right now. And if you jump into a new model architecture, you start from scratch with all of that, in some cases. I'm sure there's some, an attention block might act similarly when it's in a hybrid versus when it's in a dense transformer, but there's still like, a lot of resistance there to the new idea. But yeah, totally in agreement with your point. These are all the examples of like, why I would probably myself not want to jump into hybrids right away if I was big tech, or if I had a big existing presence in dense transformers. But yeah, there's a reason that they're moving slowly.

• 1:54:27Jason Meaux?

  Yeah, that makes sense. In that way, do you see, so the lock-in effects are a way in which you commit, but then, obviously, the longer you commit, you're beginning to build out, not just code, but best practice, where if you could imagine the possibility of Zyphra developing this architecture over the next five plus years or longer, what kinds of, so I guess the positive side of lock-in effects is that you get really good at managing this architecture that maybe not that many others are managing. Do you, any insights as to what that could look like for, for you and Zyphra and the hybrid world?

• 1:55:06Speaker 01

  There's some first mover advantage here where we can kind of steer the story like we have like internal inferencing frameworks that we've already ported that if you want a Zamba model, we can run it for you and with you the fastest compared to everyone else. If you want like community engagement, this is a negative. I think we're nimble enough to avoid the pitfalls of being locked in. So we have some very non-Zamba architectures that are cooking right now as long as they are efficient on device and punch way above their weight. We're going to look at whatever that architecture is and we're going to try and find creative ways to make it train well. In terms of avoiding lock-in with that, it's really just comes down to effort of you

• 1:55:49Speaker 01

  need to figure out fine tuning again. You need to figure out sequence parallelism for every new architecture you make and this kind of thing. And putting the effort in is our way forward. Like, we don't have any super crazy intuition behind. You just have to grind it out.

• 1:56:04Jason Meaux?

  Nice. Yeah, that makes sense. I guess related to that is, and this might be part speculation, so I apologize for that, but it seems like you do have some real world data, given that Zephyr is developing not just one flagship Mamba hybrid model, but like various parameter counts, various sizes. Any sense in which you're getting sense of like the scaling properties of this kind of model architecture, how those scaling properties are, you know, you have Chinchilla scaling laws, like we know what it looks like for just base transformers. Do you have any sense of what the scaling laws could look like for the Mamba hybrid models? Are they even that much different than what we've seen in transformers?

• 1:56:48Speaker 01

  I think they're pretty similar. We just have a slightly better slope for a lot of these. I would say I am curious to see how shared blocks affect Chinchilla, because we reach Chinchilla optimality way faster on the shared blocks. So if we were, if we train for 15 trillion tokens, at what point does the attention block saturate? It's kind of a similar thing to like, no one, I don't think is really correct scaling laws for MOEs. A similar, you can think of MOEs as a similar case where some blocks are getting trained more than others. I really don't see any reason why a Zamba architecture wouldn't scale well to 20B, 40B, 70B. And I think we're, we at some point

• 1:57:31Speaker 01

  will try this ourselves to get distillation in-house.

• 1:57:36Speaker 01

  Those are my only hot takes. I think they'll look pretty similar. It's a little bit better.

• 1:57:41Jason Meaux?

  That's awesome. This is maybe a big picture question. What does it take to do foundational ML research? And in particular, what is the kind of approach you and your team have? Because Xifer seems fairly interesting. It's a company, but it also seems to be a lab doing what I would consider pretty clearly research. So it's not necessarily product focused solely, but there's some in-commercial goals. Like, what is the process? Is it a lot of open-ended search in the team? Or there's always specific objectives with the commercial interests in mind?

• 1:58:18Speaker 01

  We have research guided by specific objectives. So those objectives are efficient inference is one. So you get a big first movers advantage if you come up with a novel architecture that actually runs efficiently on device. Because now everyone else either has to use your model or figure out all of the arcane knobs that you had to turn to get your model to work. We're sort of an applied research lab. We're not like making Mamba 3, right? We see all of the available blocks and training processes and laws. Like, every element of this Zamba model is built by someone else. But in order to recognize what all of those knobs do, and then to turn them in a way that you get the same modeling performance as a dense transformer,

• 1:59:02Speaker 01

  but with a much more amenable on-device performance and even cloud performance, is the research that we're doing. So how do you take everyone else's work? And how do you put it together in a way that's perfect at inference time? Yeah. And then this ties into the product story of if you have that inference efficient and training efficient model, it's way easier. It is now achievable to make it personalized to everybody. It's achievable to put it on a phone. It's achievable to personalize it to enterprise applications and these sort of things. It's definitely guided by products.

• 1:59:35Jason Meaux?

  Yeah, that makes sense. And this is just a side question. The, you know, when I think about the, the properties that Mamba has, but in particular, these hybrid models and I've had an increasing interest in robotics. I feel like that's could be the next big market. It's not really a market right now.

• 1:59:53Speaker 01

  Yeah.

• 1:59:53Jason Meaux?

  And the, what you need to run models locally, that just ties in almost directly with

• 2:00:01Jason Meaux?

  What your lab is focusing on, in a sense, obviously robotics usually involves completely different stack of data. So it's multimodal. You have things like the ACT Transformer. It'd be interesting to know what that would look like for Mamba. If you were to apply your work, which is primarily focused on NLP, how big of a jump is it to now be looking at like the robotics technical use cases? Or is that a bigger jump than what I'm thinking in my head? Do y'all ever think about things like that?

• 2:00:30Speaker 01

  Oh, definitely. So robotics, gaming, there's all kinds of markets where it's really nice to have flexible models on device. For robotics specifically to make, if it's just chat with humans and maybe ability to do visual question answering about like how to recognize and respond what the robot is seeing and speaking, that's pretty easy. That's pretty achievable. And I think we probably have the best offering for that right now. Once we get voice to voice done, then yeah, a robot walking around that can talk to you and is totally on devices is like nothing for us. What else would be required? It would all be very low power as well. So you could get it running on a Raspberry Pi or a Jetson.

• 2:01:13Speaker 01

  I think there's a lot of cases with like smart cities with you could put this on a Tesla and you can put in a self-driving car and just have a chat interface, run the windshield, speed them up, run the windshield at the speed that I like. All these sort of things become possible when you have a model architecture like Zamba that's multimodal and voice to voice. And that's why we're pushing for it so hard.

• 2:01:36Jason Meaux?

  Yeah, that's awesome. I guess just one more practical question is oftentimes really, if you look at the adoption of models, it often comes down to just like how easy are they to use? And that often involves like how compatible are they with the Transformers library from Hugging Face? But I guess if you could just talk with, talk about that. If I wanted to do some really interesting experiments with Zamba 2 right now, or I wanted to, whether that was fine tunes, whether that was like trying to some extended pre-training stuff, or whether, you know, how easy it is to use, what do I actually do? What libraries do I need to get into to use Zamba?

• 2:02:22Speaker 01

  So right now, you're restricted to Hugging Face Transformers. So we have Zamba 1 has been upstreamed now. It's in Transformer. You can use it with all Accelerate and Path and all of the Hugging Face like class of frameworks. For Zamba 2, we have our own fork of Transformers. It's currently getting upstreamed, but there's always just a lot of compatibility work with this kind of stuff. You can get like basic functionality right away, which is why we have our own fork. But getting it really solidly ingrained in the entire Hugging Face ecosystem is the hard leg of the journey that we're still working on. And we also have a pure PyTorch implementation for both Zamba 1 and Zamba 2. So if people have their own custom pipelines, they want to port it to their own internal inference framework. We tend to find people like that a lot.

• 2:03:07Speaker 01

  It's also more performant than putting it in like the Hugging Face ecosystem to just have pure PyTorch. You can compile it with Torch Compile, you can create custom Triton kernels very easily. Going forward, I think we're about to finally like become more ingrained because we're about ready to crack G G M L and Llama C P P and O Llama and that entire ecosystem. We've pretty much finished porting Zamba 2. I don't know, I struggle to put a date on it all and miss a date, but like in the next couple of weeks, I think people will be actually, will be able to use this seamlessly with V L M, O Llama, Llama C P P, G G M L, the entire ecosystem I think we're right about to jump onto.

• 2:03:44Jason Meaux?

  Very interesting. So the choice to fork Transformers simply, we'd like people to be able to use it immediately and there's still some kinks to work out to get it integrated. That's essentially it, right?

• 2:03:57Speaker 01

  Yeah. Yeah, we understand that like people will balk, balk when they see, oh, it's a fork? Done, I'm not jumping into it. And even like we've produced the best model at the weight class. So there's definitely incentive for those who want to use it, you got to use the fork. But for all of those who are more focused on what is the fastest way to get this running on my laptop or something, you will very soon be able to actually do that, very soon. But currently, Hugging Face fork or pure PyTorch, all your options.

• 2:04:27Jason Meaux?

  And you said like, what would be the most performant, both for inference and for, let's say I wanted to like do fine tuning or training? If I use one of those training frameworks, is it like fully optimized? Or am I, are there some trade offs with having to make sacrifices with getting compatibility? Or should I go straight to PyTorch or could I use those frameworks? It's not going to come with costs.

• 2:04:50Speaker 01

  Yeah, so for small scales, like by small scale, I mean like fine tunes. So if you're at the pre-training scale, you've got to have your own thing. You've got to have your own GPT-NOX or Megatron Deep Speed or Megatron LM. You've got to have a pre-training framework and those are not super Mamba friendly right now. We have, our stuff is in-house unfortunately right now, it's just the name of the game. How people watch your PRs otherwise. In terms of fine tuning, at the fine tuning scale, Hugging Face is totally fine. We even use Hugging Face internally to do some of our fine tunes, some of our context length extensions. When you're working on sub, you know, a couple billion tokens to 10 billion and below, 50 billion below, Hugging Face is totally fine and we recommend people use that. So, yeah.

• 2:05:35Speaker 01

  If I had like a mid-range kind of thing, maybe I would try to wrap the pure PyTorch implementation of Zamba in some like simple data loader and just try and get that going. If I have maybe 50 billion or something, 50 to 100 billion in that range. But I think most people fall into the fine tuning range where Hugging Face suffices.

• 2:05:54Jason Meaux?

  Very good, makes sense. Excellent. And if people today, like right now, wanted to start experimenting, fine tuning, I, I believe there's a GitHub link to the blog. There I can access the fork of the Transformers library. If someone was motivated, they could spin something up today. Is that correct?

• 2:06:15Speaker 01

  Oh, yeah. Definitely, yeah. It's not hard to fork and to just build our fork and load our checkpoint, yeah. And then if you just want to talk to the model, then we had a day zero release with Nvidia as a Nim, as an inference microservice. We self-hosted it on like the Maya endpoint on our website. So if you just want to talk to the model, use one of those. If you want to spin it up, you can do it today with our fork.

• 2:06:37Jason Meaux?

  Excellent. So I guess we've been talking a lot about Mamba and the Zamba models.

• 2:06:45Jason Meaux?

  But Zephyr also published a very interesting paper recently, the Tree Attention paper. And we've barely talked about it. I've not been able to dedicate myself to a complete deep dive, but my high level view is it's very interesting in the sense that it seems to be in some aspects improving upon ring attention, which was the huge thing earlier this year, where it seemed like everyone was training million plus context length llama models. If you could just talk about tree attention, how is it different from ring attention? And then what impact can it have? What is it really pushing on?

• 2:07:20Speaker 01

  Sure. So ring attention is good when you have a network topology that's a bit more flat. So like a mesh topology, like the TPUs and stuff, I think is what it was designed more for. And that's because you have a ring with point to point across each GPU in the ring. If you're in GPU land, you have two level topologies where on a single node, you have NVLink that connects all of your GPUs, it's really high bandwidth. But across nodes, you have a much slower InfiniBand, or maybe you have a RAKI, a converged Ethernet thing. That's an order of magnitude lower bandwidth than the intra-node interconnect. So there's a reason why they only did up to two nodes in the original Ring Attention paper. And it's because you're bottlenecked by this cross-node link. So the boundary

• 2:08:03Speaker 01

  GPUs at the cross node are going to be dictating how quick your infant's speed is. And it's not great. As you scale up to more and more nodes, it becomes worse and worse. If you have smaller PV states, then you're going to be, the compute is not enough to overlap that expensive communication. So it's not scalable. The whole point of Tree Attention is to recognize this fact that two level topologies are not good for rings, and to reformulate the distributed attention operation as an energy function, which allows you to use all reduce operations, communication operations, instead of point to point. All reduce operations. In existing like communication libraries, Nvidia's communication library called Nickel,

• 2:08:44Speaker 01

  Nvidia Collective Communication Library, AMD has something like Rickle, the Rockham Communication Library, AllReduce, AllReduce themselves are topology aware for GPU clusters. So they account for this two level topology and they also add some of the computation is like in-network. So with an AllReduce operation, you can do some of the sum in the network card and you can overlap that some computation, that kernel, be it on the GPU or on the card if you're doing sharp with the communication of those KV states. So basically it can scale. In short, it can scale on two level topologies because you're reformulating to an AllReduce, which is inherently topology aware and adds some of the compute is easier to overlap

• 2:09:27Speaker 01

  versus point-to-point operations, which are restricted and sort of assume a flat topology of network.

• 2:09:34Jason Meaux?

  So I guess if you could help me, maybe let's begin with the end in mind, like what will that enable me to do that Ring Attention would struggle with? What, in what cases will Tree Attention help me?

• 2:09:47Speaker 01

  So this all boils down to parallelism. Whenever you have parallelism, it boils down to memory. GPUs have a limited memory. If you have long context, so let's say if I want to do a million context on a Llama 8b, or if I need a lot of parameters that I need to store in memory at inference time, then I need more GPUs. For both of those cases I mentioned, with super long context of super large models, you need lots of GPUs to fit. More GPUs than are on a node. So you're going to two nodes, four nodes, et cetera. When you're going across nodes, you have a two level, two level topology and AllReduce is work better. So if you wanted a 16-node like context extension training run with like

• 2:10:32Speaker 01

  2 million context, Tree Attention will work much better. If you have an inference run of Llama 405b on older GPUs, you need four nodes or two nodes, Tree Attention will scale much better on that. Anytime you're scaling across nodes for GPU clusters or any imbalance of communication costs, Tree Attention will win.

• 2:10:51Jason Meaux?

  Interesting. Okay, so this is not necessarily something where we had Ring Attention and all of a sudden everyone's just going to implement Tree Attention. It's on a case-by-case basis of which one you might use if you were to try to do multi-node training and, you know, distribute the memory.

• 2:11:08Speaker 01

  Yeah, so there are cases where like if you perfectly have enough compute in your communication costs, like those two perfectly overlap one another and it's effectively free and you're on a single node, ring attention is pretty good, even though all reduces are really closely optimized because they're used in training so much, you can't really do much better than that. And free communication is totally overlapped. Yeah, not everyone will benefit right away, but those cases where you have a smaller model, longer context, which translates like smaller KV states, that means you're gonna be communication bottlenecked. There's not much compute to hide the communication. So those people will benefit from us and then the people with really large models that are inherently across nodes will benefit from us. Yeah. Yeah, you're right. It's not every single case.

• 2:11:53Jason Meaux?

  Not every single case. Trying to, for like a very, very long, training on very, very long sequences, it would help for that. And that was only if it was a multi-node training setup, or it doesn't come into play.

• 2:12:05Speaker 01

  It would help more. It would help more. Even on a single node, all reduces are a little bit more effective than point-to-point, just because tuning libraries or like collective libraries today are really closely tuned for that specific operation. Like, we're kind of exploiting the fact that all reduce is really good than point-to-point because you don't really use point-to-point much in training unless you're doing like pipeline parallelism, which is less common than tensor parallelism or data parallelism, which are both all reduce. So yeah, again, it'll help you a lot more at scale. It might help you a little bit at smaller scales, but I wouldn't promise big speedups.

• 2:12:41Jason Meaux?

  Okay. No, that that's very interesting. I guess just because we hit that topic, I have I thought of another question. The Zamba 1, I believe, was natively trained with, I think, if I'm reading this correctly, 4,096 token context length. Did Zamba 2 train in the same native token context length, or is that higher?

• 2:13:03Speaker 01

  Yeah, yeah. And then we were able to extend to the 16K range for free just by interpolating uh rotary positional embeddings. Getting beyond that is an open question that we're still working on. So I think we have a path now for 64K and 100K, which should be coming out soon, especially for the smaller models, because remember I mentioned that it's a memory overhead thing at the end of the day, right? So like Zamba 2.7B, we can go to way longer context than this. But just the engineering challenges to get us to a million context, still working on that. But yeah, trained on 4K.

• 2:13:37Jason Meaux?

  Yes, yeah. So I guess just because it's something I've experimented with Mamba 1, like really basically the idea of, okay, we're natively trained at even 2,000 token context length, but I want something capable of much more. So I'm just going to do continuous pre-training. And now all of a sudden, my data distribution, if you want to do it smart, maybe there's a curriculum learning approach. We don't just throw that much more context length at the training. At once you want to do it smart, but it seems there are all these papers and my own experimentation is it works pretty well to do continuous pre-training to expand context length. To what extent is that a solution? If somebody wanted to grab, let's say the Zamba model, or we even see this with some of the just a transformer Lamba models.

• 2:14:21Jason Meaux?

  Do you have any thoughts on continuous pre-training to extend context length? Is that sort of like a hack that you're just not going to get the quality that a natively trained model would have?

• 2:14:30Speaker 01

  No, you can get the quality. That's what we're doing ourselves. So yeah, definitely continue to be trained. You definitely need to do a curriculum of where you're warming up the sequence length and batch size across your continual pre-training run. But eventually you're going to hit a memory wall. And then that's really, that's all I'm trying to say here is that you gotta finish up sequence parallelism for the Mamba blocks as well, if you want to get past that memory wall, but it totally works.

• 2:14:55Jason Meaux?

  Yeah, absolutely. So playing around with the 2.8 billion parameter model, I could never figure out sequence parallelism. I use the DeepSpeed library because it just makes it easy to do all the 03 stuff. Of course, you have to pay tremendous costs with a training efficiency when you start offloading everything to CPU memory or yeah, so sequence parallelism is the big unlock. That's going to be exciting. Okay, cool. Does Zyfer have plans to open source that once that gets cracked?

• 2:15:23Speaker 01

  Sequence parallelism for Mamba will definitely be open sourced. The Zyfer specific training stack that we're using to produce Zamba will probably not be open sourced, but the general process and like what we did with Tree Attention, the kernels that we use, the scripts that we use, those will all be open source because we want the community to like move towards Mamba and see that there's something there.

• 2:15:43Jason Meaux?

  Yeah, awesome. Great to hear. Anything else that we didn't cover about Zamba? The timing is in around the Zamba 2 release. Anything else you think we should cover before we wrap?

• 2:15:53Speaker 01

  No, I think we've covered pretty much everything.

• 2:15:55Jason Meaux?

  Appreciate your time, Quinn. It's absolutely great speaking with you.

• 2:16:00Speaker 01

  Likewise.

• 2:16:00Nathan Labenz?

  So let's do this zoomed out one. I think this has been fantastic. I think that this is for the real ones out there who are, you know, very interested in learning more about what it actually takes to make these models work and just how many little nitty gritty details go into that. So I've learned from this. I think a lot of people will appreciate all of the lessons that you've shared from many long days and probably some long nights working on this stuff. I think it's been great. Let's do the zoomed out thing. With all that in mind, where is Zyfer going as a company? What role do you guys want to play in the lives of users? You talked a little bit about how you see yourselves competitively against the big guys, but I'd love to just hear on whatever time scale, you know, what

• 2:16:45Nathan Labenz?

  you think my AI assisted life is going to look like and what role Zyfer and Zyfer models will play in that.

• 2:16:52Speaker 01

  Yeah, definitely. We've focused really hard on producing the best models for one modality and now we're ready to expand out. So this includes other modalities. So we want like visual question answering. I want to be able to live edit my pictures. I want to be able to talk voice to voice. Like I talk to you voice to voice. I want to talk to my AI voice to voice as well. This includes like actually deploying personalization ability to people. This includes like actually launching like a Maya both enterprise and on device for consumers. This includes the broader ecosystem, for example, Olama and Llama CPP and stuff. Those are all transformer based and getting the hybrid architectures into those is one of our next priorities so that everyone can actually deploy the

• 2:17:34Speaker 01

  model really quickly instead of using our fork of hugging face all this sort of thing and having a much more like a better user experience for hackers and for consumers. And then there's also some higher level things. We think like memory, so retrieval is really important for personalization that we're going to like, how does retrieval interplay with long context reinforcement learning so that like we can really tie down? It's not just continual pre-training when you're following. Like, we definitely need some sort of approximation at least of RLHF or RAIF for users so that we can really extract as much as possible from the little data that we're going to get per user. There's a lot of ways to expand, but those are some of the main ones.

• 2:18:13Nathan Labenz?

  Cool. I love it. For now, I'll say Jason Mo and Quentin Anthony from ZIFRA. Thank you both for being part of the Cogit revolution. Thanks, Zizan. Thanks. It is both energizing and enlightening to hear why people listen and learn what they value about the show. So please, don't hesitate to reach out via email at [email protected] or you can DM me on the social media platform of your choice.