MMLU Stays at Chance Until Late Training. Loss Is Only an Alarm.

On this episode of *The Cognitive Revolution*, Nathan Labenz brings back Jason Meaux — the co-host who followed the Mamba literature far enough to build statespace.info — and sits down with Quentin Anthony, head of model training at Zyphra. The company has just released Zamba 2, a family of hybrid language models that mix selective state-space blocks with a small number of shared attention blocks. The conversation splits roughly in half: a product thesis about personalized on-device AI, and a long, unusually concrete debrief of what actually happens when you try to train those models.

Zamba 2 ships at three sizes. The 1.2B is for severely constrained devices; the 2.7B is the quality workhorse for things like email and meeting summaries; the 7B is the point where a model can still fit on a stronger laptop and start to reason through harder problems. All three use the same hybrid recipe. What makes the training story worth two hours is not a single headline number. It is the gap between the metrics people watch and the capabilities users actually feel.

The gap shows up most clearly on MMLU. Quentin says the benchmark stays random for most of training, then lifts late, once the model has grokked a strict format: answer with a single letter that has to stand in for all the background knowledge in the question. On a dense Transformer that takeoff waits until the learning rate has settled and the run has seen on the order of several hundred billion tokens. Pure state-space models are slower. Falcon Mamba, he recalls, needed something like three or four trillion tokens before the same signal appeared, because the instruction to answer with a letter gets lost in a pure SSM block.

Zyphra’s site says the future of AGI is a mix of cloud and on-device deployment, with the mix shifting toward local inference. Quentin gives three reasons. The main one is personalization: frontier labs serving one giant model cannot specialize to every person on the planet, and “changing a system prompt per person is not enough.” He wants preferences — how you like to be talked to, which restaurant is yours — baked into the weights. The second is privacy: laptops and phones hold data people do not want in the cloud, and enterprises do not want proprietary code sent to a third-party API. The third is money. Frontier inference capex is being fueled by venture capital; if devices can run their own models, that bill moves to the user, with offline use and lower latency as extras.

Nathan has been trying to get a model to write as him. Fine-tuning GPT-4o picks up style reasonably well and learns facts poorly. Stuffing about 50,000 tokens of personal context into a system prompt still hallucinates. Quentin’s bet is continual pre-training on a small model: a weight update overnight while the phone is plugged in, cheap enough on consumer hardware that LoRAs can make it cheaper still. For the lag between “you’re being too dry” and tonight’s bake, he points at activation steering — the Golden Gate Claude trick — as a way to fake faster learning until the real update lands. “It’s got to be continual learning and it’s got to be per user,” he says, “and the only way to do that is with weights on the phone.”

The 1.2B is for battery-constrained phones; at 4-bit it is meant to run on a Raspberry Pi or a Jetson Nano. The 2.7B is the one that keeps surprising him on summarization, and the team already has early in-house fine-tunes for roleplay, meetings, and even audio note-taking. The 7B is the turning point: strong enough to think more clearly, and, in the product design, able to notice when it should ask a cloud oracle — seamlessly, with the user choosing how often that happens. The product vision is Maya, a personal cloud across devices for photos, conversations, and notes, with the user deciding how much to share. Training on that private corpus is just continual pre-training on whatever edge hardware you have. You can already download the weights and start; it is just not one-click yet.

Jason frames the local-inference problem as two walls Transformers hit hard: memory and compute. The KV cache grows with sequence length, so a long article can blow past phone RAM, and attention is expensive compared with a fixed-size state. Mamba’s appeal is exactly those systems properties — constant state, better throughput — which is why robotics papers were already looking at it for split-second local decisions. Quentin’s training principle is to eat some training cost in order to be “blazing fast” at inference. Dense attention in the style of a Phi model, he says, is not going to run on a phone for the long inputs Zyphra wants, such as news and email summarization.

Attention is still required for some tasks. He has never played with a pure SSM or RWKV model that actually speaks well in an in-context-learning regime; benchmarks can look fine and the conversation falls apart. The hybrid fix came from a finding several groups hit independently: attention blocks are highly correlated across depth. They keep telling the MLPs “these are the important tokens,” and most of the specialization across layers comes from the MLPs. Zamba 1 therefore uses a single global attention block, invoked every six Mamba blocks. Zamba 2 adds a little depth-wise specialization: the 1.2B still has one shared block, while the 2.7B and 7B alternate two shared blocks, A and B, plus independent LoRAs at each depth.

Sharing weights means those parameters are applied over and over, so total FLOPs go up for a given parameter budget. Mamba 2’s throughput on parallel hardware more than pays for that sin. The net result Quentin quotes is about 20 to 30 percent better time to first token and time per output token versus a similar dense Transformer, with the KV cache invoked something like 10 or 13 times instead of 30-plus. The attention block is fully tied — queries, keys, and values — and because it is applied six times it gets very saturated gradients. Linear attention, he says, has not yet done what they need; Liquid AI-style “train with attention, then linearize” is an interesting post-training play, not their path.

His working intuition is a division of labor. Attention is there to take as many notes as possible about the incoming sequence — every token to every other token, multiple heads, each from a slightly different angle. MLPs try to find deeper meaning in those notes and mix across heads. Linearizing attention is a bet about how many of those notes you can throw away. A Mamba block has much smaller projectors, so it can do less with the notes from a selective scan than an MLP can do with full attention. That is why Black Mamba’s “maybe we do not need the expensive notes” bet failed, and why a shared, fully tied attention block plus specialized MLPs is the compromise they landed on.

Zyphra’s first public Mamba bet was Black Mamba, a mixture-of-experts plus Mamba model released in February 2024, among the first handful of Mamba papers — a count Nathan and Jason could still keep, and which has since passed 300. The motivating figure in the original Mamba paper compared pure Mamba, Mamba plus attention, and Mamba plus MLP on loss per FLOP, without showing absolute loss. They read that as a chance that the expensive cross-sequence notes from attention might not be needed, and that MoE sparsity would crush inference cost. Long-sequence overhead was indeed excellent. In-context learning was not. MMLU never really got up, which Quentin treats as an early alarm that the model cannot learn in context. They published it as an experiment and did not deploy it.

That failure sharpened a broader claim: “loss is not a very useful indicator like at all.” He uses it to detect spikes — to tell whether the model has died — and as an early read on whether training is saturating or the learning rate is too timid. Whether loss goes to three or to two is dataset-dependent; noise in the data sets the floor. If the curve plateaus because the learning rate was not aggressive enough, that is a real alarm. Everything else, he says, comes from evaluations.

Evals have their own failure modes. Maximize them at all costs and you would train on synthetic textbooks, blow up the numbers, and get a model that is dry, hard to fine-tune, and brittle on unclean data because it assumed the world looks that clean. HellaSwag, in his reading, is almost a smoother loss: general language modeling that rises linearly. MMLU is both an emergence indicator and a contamination detector — if it explodes early, you probably downloaded the questions. ARC Easy and Challenge track factual recall and how much Wikipedia and textbook data you put in. Then you need vibes. A prompt set that asks the model to role-play as a pirate is, in his account, a serious test: Phi-style models stick with it briefly and then drop it, which benchmarks will not show and users will feel immediately.

He treats the late MMLU lift as grokking plus schedule. The model has to see enough tokens to notice that instruction-following data of this type exists, and the learning rate has to settle into the loss basin before fine-grained changes stick. MiniCPM is his exhibit: a long constant learning rate, then a fast anneal to zero, with almost all of the quality showing up in that final anneal. Nathan tries to steel-man “emergence is a mirage” — logarithmic axes can make an order-of-magnitude grind look like a spike — and Quentin does not pretend to have a theory beyond the empirical pairing of tokens and schedule. Both, together, are what he is willing to call grokking.

Zyphra tried Sophia and did not switch. Quentin’s suspicion is that more aggressive learning-rate schedules can create a mirage of optimizer quality when you are really just playing with step size. They are looking at Shampoo; they are staying on Adam for now. He would rather kick a run off on day one, feed it more tokens, and still land at a lower loss than spend three weeks hunting a brittle hyperparameter combination that a paper said should be better.

The landscape picture he will endorse is not a flat plain full of pockets. It is a general downward slope with pockets along the way, and larger models make that landscape smoother, so the optimizer has an easier time. That is why a lot of conventional wisdom assumes big models: at large scale a step is broadly accurate anyway, and people say Shampoo does not scale because you are already in a giant smooth basin. On a 3B novel architecture the landscape is bumpier, so step accuracy actually matters. Second-order methods and careful batch-size choices are, in that regime, not optional decoration.

He learned the batch-size version of this on Pythia at EleutherAI. Larger batches should mean more accurate gradients, and yet they were seeing instabilities. Looking back, they had scaled learning rate up with batch size. A few outlier samples in a large batch kick the gradient the wrong way; a large step follows; several bad batches in a row and the run diverges. Lower batch sizes plus gradient clipping — hope the next few steps can reverse a bad one — were the ugly fix. “I’ve noticed that every run has its own cursed way,” he says. Adam is extremely robust. Sophia or Adafactor may win in some settings; they can also simply fail to converge.

Mixture-of-experts models make the soup worse. The attention block sees the full batch; expert blocks see a subdivided batch, so they live at a different gradient-noise scale. Scale batch size by number of experts and the attention block now gets a massive batch. That mismatch is a big part of why Black Mamba was a pain to train. Hybrids have a milder version of the same problem: the shared attention block sees more data and gets more accurate gradients than the Mamba blocks. Cerebras is helping them port μTransfer to this setting so width-wise hyperparameter search on a cheap skinny model can be divided across to the real width. Until that lands, navigating the hyperparameter soup is what determines whether a run lives or dies.

Annealing is the example Quentin uses for a secret that leaked. Phase one is a standard cosine decay on noisier web-scale data, ending at a relatively high learning rate, so the model still knows what messy internet text and pirate-talk look like. Phase two rewarms the learning rate on a much smaller, carefully filtered subset — tens of billions of tokens you can actually inspect — then anneals aggressively to zero, with replay from phase one so the model does not forget. Clean data lets you stomach a high learning rate without the Pythia-style outlier kicks. Zyphra wrote this up in a GitHub cookbook; they say they found the two-phase recipe independently on Zamba 1. On the frontier-lab side it was a hidden trick until Llama 2, then the Mistral split, then the rest of the industry. Zamba 2 doubled the anneal set and tripled phase-one data to about three trillion tokens; Zamba 1, at one trillion, had already matched models trained on five to fifteen trillion.

Distillation is the new version of that secret. Synthetic data from a large model is the poor man’s form: you train on the argmax of the logits, which is only what the big model is paying attention to, and you can even do it from closed models. Real distillation wants the full distribution, including the near-zero values that say what the teacher is *not* looking at, and ideally a live teacher at every training step. Running Llama 3 405B on every step of a 7B pre-train, he argues, is close to the cost of just training a bigger model. They see generally positive conversion from logits and have not cracked pre-training-scale distillation. If you are Microsoft working on Phi, he says, you just call GPT-4 on your own systems. “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.”

Data cleaning is the other lag. It takes time, patience, and people, and open source will always trail labs that own clean first-party data. NVIDIA’s NeMo Curator lowered the bar; DCLM, FineWeb, Dolma, and Zyphra’s own Zyda 2 are improving the public mix. μTransfer is another hoarded tool: sweep learning rate and batch size on a cheap low-width model, then divide by a width scaling factor and transfer, so activations do not explode. Second-order optimizers for small models — Gemma will say they used one and not how they sharded the extra states — plus Llama-style 4D parallelism and dedicated checkpoint-restart nodes, are the same class of expensive empirical knowledge. Even the GPT-3 configs everyone copied hide a kernel trick: powers of two in the hidden size, vocabulary rounded to a multiple of 64. Go deeper and you hit wave quantization, where 101 units of work on 100 streaming multiprocessors take two timesteps and roughly half the throughput. Alternative architectures like Zamba are *not* hoarded, because big labs saw quality issues with pure SSMs and will not bet a huge run on a hybrid that might be slightly worse.

That last fact is Zyphra’s niche. OpenAI, xAI, and Anthropic are paid to produce the highest-quality model; Claude Sonnet 3.5 only had to be a little better than ChatGPT to take share, so they will not gamble on exotic hybrids. Small models are overlooked. Stay nimble, keep hunting the next free lunch — second-order optimizers, cheap distillation — and a 7B can be best in its class. Nathan’s question about Entropix produces a few minutes of crossed wires: Quentin first files it under “doesn’t scale,” citing Tim Dettmers’ large-parameter trials that looked no better than Adam, while Nathan is talking about an inference sampler that injects clarifying questions when the model is out of distribution. Once they align, both like the idea of *learning* a sampling function, and Quentin notes the pattern: large models grok data selection on their own; small models need an explicit process that imitates the same behavior faster.

Asked where the field is, Quentin says the limits of scaling and of architecture have not been found, and people are still reinventing training tricks that some labs already know and never write down. The Zamba 1 block is six Mamba-1 layers, then a shared attention-plus-MLP, repeated, with no positional encodings — Mamba already encodes order in the state. Before attention, they concatenate the Mamba residual stream with the original input embedding; projecting back down to save compute lost quality, so they eat the wider concat. The 6:1 ratio is purely empirical: seven and five were slightly worse, and a shared block may need to be applied more often than independent ones. Zamba 2 swaps in Mamba 2’s SSD algorithm — Albert Gu and Tri Dao’s May 2024 restriction on the A matrix that turns the scan into matrix multiplies and fills tensor cores — for throughput, not for a quality jump. Pure Mamba 2 still failed in-context learning. The 1.2B keeps one global block with LoRAs on both attention and MLP; the 2.7B and 7B use the A/B alternation and LoRAs only on the MLP, because that is where depth specialization actually lives. They added RoPE and got a small long-range boost at the cost of rigid context extension; Quentin is not sure that trade was right. Mamba state size is 64: larger states help long context, attention already covers that, and the inference penalty is obvious.

Long context is still an engineering mountain. Attention has ring attention and Zyphra’s own tree attention, which reformulates distributed attention as an energy function so you can use topology-aware all-reduces instead of ring point-to-point — a better fit for GPU clusters whose in-node NVLink is an order of magnitude faster than the cross-node link. Mamba blocks still lack sequence parallelism. Activations plus a million-token input, plus high-precision gradients and optimizer states, will not fit on an H100; without long-sequence training the model will not generalize. Jason tried to patch sequence parallel into Mamba 1 via DeepSpeed’s Ulysses library and could not crack one percent of the work. Mamba 2 looks more Transformer-like and more tractable. Quentin does not want the model to be brain-like — “I forget why I’m walking into a room” — and wants attention’s cross-sequence notes to stay. Native training context is 4K, with a free-ish extension to 16K by interpolating RoPE and a path they think they have to 64K and 100K. Continual pre-training with a sequence-length curriculum does work; the wall is memory, and sequence parallel for Mamba will be open-sourced when it lands.

Today you can already experiment. Zamba 1 is upstream in Hugging Face Transformers; Zamba 2 still needs Zyphra’s Transformers fork, plus a pure PyTorch path for people who want `torch.compile` and custom Triton kernels. Fine-tunes and context-extension runs under a few tens of billions of tokens are fine in Hugging Face — the team uses it internally for those. Pre-training scale still wants Megatron or DeepSpeed, which are not very Mamba-friendly yet. Ports to llama.cpp, Ollama, vLLM, and GGML were described as weeks away at recording. For anyone who just wants to talk to the model, there was a day-zero NVIDIA NIM microservice and a Maya endpoint on the site. The company wants to add visual question answering, live image edits, and voice-to-voice, then actually ship per-user personalization — plus retrieval for memory and some on-device approximation of preference learning, because continual pre-training alone will not squeeze enough out of the little data each user produces.