MMLU 在训练后期才脱离随机水平，loss 只是一个警报

本期《The Cognitive Revolution》，Nathan Labenz 请回了 Jason Meaux。Jason 是那位对 Mamba 相关文献追得很深、并做出了 statespace.info 的联合主持人。节目里，两人与 Zyphra 的模型训练负责人 Quentin Anthony 坐下来对谈。这家公司刚发布了 Zamba 2，这是一组混合语言模型，把选择性状态空间模块与少量共享的注意力模块混在一起。整场对话大致分成两半：前半是关于个性化端侧 AI 的产品论点，后半是一次很长、异常具体的复盘，讲的是训练这类模型时真正会发生什么。

Zamba 2 有三种规格。1.2B 面向受严重限制的设备；2.7B 是质量上的主力，适合邮件和会议摘要这类任务；7B 则是一个能在更强的笔记本上运行、开始推理更难问题的分界点。三者用的是同一套混合配方。让这段训练经历值得聊上两个小时的，是大家盯着的指标与用户实际感受到的能力之间的落差，单看某一个醒目的数字，说明不了什么。

这种落差在 MMLU 上表现得最清楚。Quentin 说，这个基准在训练的大部分时间里都是随机水平，等模型 grok 了（即突然领会）一种严格的格式之后，它才会在后期抬升：答案必须是一个字母，而这个字母要代表题目里所有的背景知识。在稠密 Transformer 上，这种起飞要等学习率稳定下来、训练看过几千亿个 Token 之后才会出现。纯状态空间模型更慢。他回忆说，Falcon Mamba 需要大约 3 万亿到 4 万亿个 Token，同样的信号才会出现，因为“只用一个字母作答”这条指令，在纯 SSM 模块里会丢失。

Zyphra 的网站写道，通用人工智能（AGI）的未来是云端与端侧部署的混合，而且这种混合正在向本地推理倾斜。Quentin 给出了三个理由。首要的一个是个性化：为全世界每个人专门定制一个巨型模型是做不到的，而“为每个人单独改一条系统提示词”也远远不够。他希望把偏好——比如你喜欢别人怎么跟你说话，哪家餐厅是你的常去之处——直接写进权重里。第二是隐私：笔记本和手机里有许多人不想上传到云端的数据，企业也不想把专有代码交给第三方 API。第三是钱。前沿模型的推理资本开支，目前主要靠风险投资撑着；如果设备能够运行自己的模型，这笔账就转移到了用户身上，离线使用和更低的延迟则是附带的好处。

Nathan 一直在试着让模型用他的口吻写作。微调 GPT-4o 能较好地学会风格，却难以学会事实。把大约 50000 个 Token 的个人背景塞进系统提示词，仍然会产生幻觉。Quentin 押注的是在小模型上做持续预训练：手机插着电的夜里完成一次权重更新，在消费级硬件上足够便宜，而 LoRA 还能让它更便宜。对于“你太干巴巴了”到“今晚烘焙完成”之间的这段滞后，他指向激活引导（activation steering），也就是 Golden Gate Claude 的把戏，把它当作一种假装学得更快的办法，直到真正的更新落地。“它必须是持续学习，而且必须是按用户来做，”他说，“而唯一能做到这一点的办法，就是把权重放在手机上。”

1.2B 面向电量有限的手机；在 4 位量化下，它的目标是能在树莓派或 Jetson Nano 上运行。2.7B 则是在摘要任务上一直让他感到惊喜的那个版本，团队已经有了早期的内部微调版本，用于角色扮演、会议，甚至音频笔记。7B 是转折点：它足够强，能想得更清楚；在产品设计上，它还能察觉到什么时候应该去问云端的“神谕”，而且切换是无缝的，多久切换一次由用户决定。产品愿景是 Maya，这是一个跨设备的个人云，用来存放照片、对话和笔记，分享多少由用户决定。用这些私人语料训练，不过就是在你手边的任何边缘硬件上做持续预训练。你现在已经可以下载权重开始尝试了，只是还做不到一键完成。

Jason 把本地推理的难题归结为 Transformer 撞上的两堵墙：内存与算力。KV 缓存随序列长度增长，因此一篇长文章就可能超出手机内存；而与固定大小的状态相比，注意力的计算代价很高。Mamba 的吸引力正在于这些系统层面的特性：状态恒定，吞吐量更高。这也正是机器人方面的论文早就在关注它的原因，用来做毫秒级的本地决策。Quentin 的训练原则是，宁可多花一些训练成本，换来推理阶段的“快得飞起”。他说，像 Phi 那样的稠密注意力模型，在手机上处理 Zyphra 想要的长输入（比如新闻和邮件摘要）时跑不动。

有些任务仍然需要注意力。他从没玩过真正能在上下文学习场景中说得很好的纯 SSM 或 RWKV 模型；基准测试可能看起来还不错，对话却会散架。混合方案的解法来自几个团队各自独立的发现：注意力模块在各深度之间高度相关。它们不断地告诉 MLP“这些才是重要的 token”，而各层之间大部分的专门化，其实来自 MLP。因此，Zamba 1 只用一个全局注意力模块，每隔六个 Mamba 模块调用一次。Zamba 2 则增加了一点深度上的专门化：1.2B 仍然只有一个共享模块，而 2.7B 和 7B 则在两个共享模块 A 和 B 之间交替，并在每个深度上加入各自独立的 LoRA。

共享权重意味着这些参数会被反复使用，因此在参数预算相同的情况下，总 FLOP 会增加。但在并行硬件上，Mamba 2 的吞吐量足以补偿这一点。Quentin 引用的最终结果是：与一个相近的稠密 Transformer 相比，首 Token 时间和每个输出 Token 的时间都好了大约 20% 到 30%；KV 缓存的调用次数降到大约 10 到 13 次，原本则要 30 多次。注意力模块完全绑定，查询、键和值都是共享的，而且由于它被应用了六次，梯度会变得非常饱和。他说，线性注意力目前还没有做到他们需要的程度；Liquid AI 那种“先用注意力训练，再线性化”的做法，是一个有意思的后训练思路，但不是他们的路线。

他心中的工作直觉是一种分工。注意力负责尽可能多地记下输入序列的笔记：每个 token 与其他每个 token 的关系，多个头，每个头从略有不同的角度去看。MLP 则试图从这些笔记里找出更深层的含义，并在各个头之间进行混合。线性化注意力，赌的是你能扔掉多少笔记。Mamba 模块的投影器小得多，因此，它能从选择性扫描得到的笔记中做的事，远不如 MLP 配合完整注意力能做的多。这就是为什么 Black Mamba 那种“也许我们并不需要那些昂贵的笔记”的押注失败了，也是为什么他们最终选定了这样一种折中方案：一个共享的、完全绑定的注意力模块，加上专门化的 MLP。

Zyphra 第一次公开押注 Mamba，是 Black Mamba：一个混合专家（MoE）加 Mamba 的模型，于 2024 年 2 月发布，是最早几篇 Mamba 论文之一。Nathan 和 Jason 当时还数得过来，如今这类论文已超过 300 篇。原始 Mamba 论文里有一张动机图，比较了纯 Mamba、Mamba 加注意力，以及 Mamba 加 MLP 每 FLOP 的 loss，但没有给出绝对的 loss。他们把这解读为一种可能性：注意力中昂贵的跨序列笔记或许并不需要，而 MoE 的稀疏性能大幅压低推理成本。长序列上的开销确实表现优异。但上下文学习并不行。MMLU 始终没有真正起来，Quentin 把这视作模型无法在上下文中学习的早期警报。他们把它作为一次实验公开发表，并未部署。

那次失败让一个更宽泛的论断变得更加尖锐：“loss 其实一点也不是很有用的指标。”他用它来检测尖峰，判断模型是否已经“死掉”，并借此早期判断训练是在饱和，还是学习率太过保守。loss 最终降到 3 还是降到 2，取决于数据集；数据里的噪声决定了下限。如果曲线因为学习率不够激进而趋于平缓，那才是真正的警报。其他的一切，他说，都来自评测。

评测也有自己的失败模式。如果不惜一切代价去追求它们，你就会用合成教科书来训练，把数字刷得很高，最终得到一个干巴巴、难以微调、在不干净数据上很脆弱的模型，因为它假设世界就是这么干净的。在他看来，HellaSwag 几乎就是一条更平滑的 loss 曲线：一般性的语言建模，呈线性上升。MMLU 既是涌现的指示器，也是污染的检测器：如果它很早就暴涨，你多半是把题目下载进了训练数据。ARC Easy 和 ARC Challenge 追踪的是事实记忆，以及你放进去了多少维基百科和教科书数据。然后，你还需要凭感觉。在他的描述里，一组要求模型扮演海盗的提示词是一项严肃的测试：Phi 风格的模型只会短暂地坚持，然后就放弃角色。这一点基准测试不会显示，但用户会立刻感觉到。

他把 MMLU 后期的抬升，视作顿悟（grokking）加上学习率调度的共同作用。模型必须看到足够多的 token，才能注意到这类指令遵循数据的存在；而且学习率必须先在损失盆地里稳定下来，细粒度的变化才能保留下来。MiniCPM 就是他的例子：先用很长的恒定学习率，再快速退火到零，几乎所有的质量都在最后这一段退火中显现。Nathan 试图为“涌现是海市蜃楼”这一观点做最有力的辩护：对数坐标轴可能会让一个数量级的缓慢积累看起来像一次跳跃。而 Quentin 并不假装自己有超出“token 数量与调度的经验配对”之外的理论。两者加在一起，才是他愿意称之为顿悟的东西。

Zyphra 尝试过 Sophia，但没有换过去。Quentin 怀疑，当你其实只是在调整步长时，更激进的学习率调度会制造出一种优化器更强的假象。他们在看 Shampoo，但目前仍然用 Adam。他宁可第一天就把训练启动起来，多喂一些 token，依然能得到更低的 loss，也不愿意花三周去寻找一个脆弱的超参数组合，而某篇论文说它本该更好。

他认可的地形图景是：一个总体向下的斜坡，沿途散布着一些凹坑。更大的模型会让这片地形更平滑，因此优化器更容易前进。这正是为什么许多传统观点都默认大模型：在大规模下，一步走得大体上就是准的，人们说 Shampoo 不能扩展，是因为你已经身处一个巨大而平滑的盆地里。而在一个 3B 的新型架构上，地形更崎岖，步子准不准就真的很重要了。在那种情况下，二阶方法和细致的批量大小选择，就不再是可有可无的装饰。

他在 EleutherAI 的 Pythia 上，学到了批量大小版本的这一课。更大的批量本应意味着更准确的梯度，但他们却看到了不稳定。回头看，问题在于他们把学习率随批量大小一起调高了。大批量里的几个离群样本，会把梯度拧到错误的方向上；随之而来的是一个很大的步长；连续几个坏批次之后，训练就发散了。降低批量大小，再加上梯度裁剪（希望接下来的几步能扭转糟糕的那一步），是那个并不好看的解决办法。“我注意到，每次训练都有它自己诅咒般的方式，”他说。Adam 非常稳健。Sophia 或 Adafactor 在某些场景里可能胜出，但它们也可能干脆无法收敛。

混合专家模型会让这锅汤变得更糟。注意力模块看到的是完整的批次；专家模块看到的是被切分后的批次，因此它们处在不同的梯度噪声尺度上。如果按专家数量放大批量大小，注意力模块就会拿到一个巨大的批次。这种错配，正是 Black Mamba 很难训练的主要原因之一。混合架构也有一个较温和的同类问题：共享的注意力模块看到的数据更多，得到的梯度比 Mamba 模块更准确。Cerebras 正在帮他们把 μTransfer 移植到这种设置下，使得在一个便宜的窄模型上做宽度方向的超参数搜索，可以再映射到真实的宽度上。在此之前，在超参数的汤里找路，决定了一次训练能不能活下来。

退火是 Quentin 举的一个秘密外泄的例子。第一阶段是在噪声更大的网络规模数据上做标准的余弦衰减，最终停在相对较高的学习率上，因此模型仍然知道杂乱的互联网文本和海盗式的说话是什么样。第二阶段，把学习率在一个小得多、经过精心过滤的子集上重新升温，这个子集有几百亿个你真正可以检查的 token；然后激进地退火到零，并重放第一阶段的数据，以免模型遗忘。干净的数据让你可以承受较高的学习率，而不会出现 Pythia 那样的离群尖峰。Zyphra 把这套做法写进了一份 GitHub 上的 cookbook；他们说，这套两阶段配方是在 Zamba 1 上独立发现的。在前沿实验室那一边，它长期是个隐藏技巧，直到 Llama 2 出现，然后是 Mistral 那边的拆分做法，行业其余的公司这才陆续跟上。Zamba 2 把退火数据集扩大了一倍，并把第一阶段的数据增加到大约 3 万亿个 token；Zamba 1 只用了 1 万亿个 token，就已经追平了用 5 万亿到 15 万亿个 token 训练的模型。

蒸馏是那个秘密的新版本。来自大模型的合成数据是穷人版的形式：你是在模型 logits 的 argmax 上训练，那只是大模型关注的东西，甚至可以从闭源模型那里获取。真正的蒸馏需要完整的分布，包括那些接近零的值，它们说明了老师*没有*在看什么；并且理想情况下，每一步训练都要有一个在线的老师。他认为，在一个 7B 模型的预训练过程中，每一步都运行 Llama 3 405B，其成本接近于直接训练一个更大的模型。他们从 logits 中看到了总体上正向的转化效果，但还没有攻克预训练规模的蒸馏。他说，如果你是微软在做 Phi，你就直接在自己的系统上调用 GPT-4。“要么他们找到了更便宜的办法，我们得去找到它；要么他们有更多的钱，我们就得去找更便宜的办法。”

数据清洗是另一个滞后环节。它需要时间、耐心和人力，而开源永远会落后于那些拥有干净的自有数据的实验室。英伟达的 NeMo Curator 降低了门槛；DCLM、FineWeb、Dolma 以及 Zyphra 自己的 Zyda 2，正在改善公开的数据组合。μTransfer 是另一种被囤积的工具：在一个便宜的低宽度模型上扫描学习率和批量大小，然后除以宽度缩放因子再迁移，这样激活值就不会爆炸。针对小模型的二阶优化器（Gemma 会说他们用过一种，但不会说明额外状态是如何分片的），加上 Llama 风格的 4D 并行，以及专门的检查点重启节点，都属于同一类代价高昂的经验知识。即便是大家都抄过的 GPT-3 配置，也藏着一个内核层面的技巧：隐藏层尺寸取二的幂，词表凑成 64 的倍数。再深入一层，你会撞上 wave 量化：101 份工作放在 100 个流式多处理器上，要花两个时间步，吞吐量大约减半。像 Zamba 这样的替代架构并没有被囤着不公开，因为大实验室见过纯 SSM 的质量问题，不愿意把一次大型训练押在一个可能略差的混合模型上。

最后这一点，正是 Zyphra 的生态位。OpenAI、xAI 和 Anthropic 的任务是做出质量最高的模型；Claude Sonnet 3.5 只需要比 ChatGPT 稍好一点，就能抢到份额，所以它们不会去赌奇特的混合架构。小模型常被忽视。保持灵活，不断寻找下一个免费的午餐（二阶优化器、廉价的蒸馏），一个 7B 模型就可能在同级别中名列第一。Nathan 关于 Entropix 的问题，引出了几分钟的鸡同鸭讲：Quentin 起初把它归为“无法扩展”，并援引 Tim Dettmers 那些大参数实验，看起来并不比 Adam 好；而 Nathan 说的是一种推理采样器，会在模型处于分布之外时插入澄清性问题。等两人对上了频道，都喜欢*学习*一个采样函数的想法。Quentin 还指出了一种规律：大模型会自行学会数据筛选；小模型则需要一个显式的流程，去更快地模仿同样的行为。

问到这个领域处于什么位置，Quentin 说，扩展的极限和架构的极限都还没有找到；人们仍在重新发明一些训练技巧，而有些实验室早已掌握，只是从不写下来。Zamba 1 的模块是六个 Mamba-1 层，然后是一个共享的注意力加 MLP，如此重复，没有位置编码，因为 Mamba 已经在状态中编码了顺序。在注意力之前，他们把 Mamba 的残差流与原始输入嵌入拼接在一起；为了省算力而先投影回去，结果损失了质量，所以他们选择接受更宽的拼接。6:1 的比例纯属经验：7 和 5 都略差，而且共享模块可能需要比独立模块更频繁地调用。Zamba 2 换上了 Mamba 2 的 SSD 算法（也就是 Albert Gu 和 Tri Dao 在 2024 年 5 月对 A 矩阵施加的限制，使扫描变成矩阵乘法，并填满张量核心），目的是提升吞吐量，质量上并没有跃升。纯 Mamba 2 仍然在上下文学习上失败。1.2B 保留一个全局模块，并在注意力和 MLP 上都加了 LoRA；2.7B 和 7B 使用 A/B 交替，只在 MLP 上加 LoRA，因为深度上的专门化恰恰就存在于那里。他们加入了 RoPE，获得了一点长程能力的提升，代价是上下文扩展变得僵硬；Quentin 不确定这笔交易是否划算。Mamba 的状态大小是 64：更大的状态有助于长上下文，但注意力已经覆盖了这一点，而推理的代价却很明显。

长上下文仍然是一座工程大山。注意力方面有环形注意力（ring attention），以及 Zyphra 自己的树状注意力。后者把分布式注意力重新表述为一个能量函数，因此可以使用拓扑感知的 all-reduce，代替环形的点对点通信，这更适合 GPU 集群，因为节点内的 NVLink 比跨节点链路快一个数量级。Mamba 模块仍然缺少序列并行。激活值加上百万 token 的输入，再加上高精度梯度和优化器状态，放不进一块 H100；如果不做长序列训练，模型就无法泛化。Jason 试着通过 DeepSpeed 的 Ulysses 库，把序列并行打进 Mamba 1，结果连 1% 的工作都没能攻克。Mamba 2 看起来更像 Transformer，也更容易处理。Quentin 并不希望模型像大脑一样（“我忘了我为什么要走进这个房间”），他希望保留注意力那种跨序列的笔记。原生训练上下文为 4K，通过插值 RoPE 可以较为轻松地扩展到 16K，他们认为自己也有通往 64K 和 100K 的路径。配合序列长度课程的持续预训练确实有效；瓶颈在于内存，而 Mamba 的序列并行一旦落地，就会开源。

今天你就可以动手尝试。Zamba 1 已经合并进了 Hugging Face Transformers；Zamba 2 仍然需要 Zyphra 的 Transformers 分支，外加一条纯 PyTorch 路径，供想用 torch.compile 和自定义 Triton 内核的人使用。在几百亿 token 以内的微调和上下文扩展训练，在 Hugging Face 中就可以完成，团队内部也是这样用的。预训练规模仍然需要 Megatron 或 DeepSpeed，而它们对 Mamba 还不太友好。录制时，移植到 llama.cpp、Ollama、vLLM 和 GGML 的工作被说成还需要几周。如果你只想跟模型聊聊，发布当天就有英伟达 NIM 微服务，网站上也有 Maya 的端点。公司接下来想增加视觉问答、实时图像编辑和语音对语音的功能，然后真正推出按用户个性化的功能，再加上用于记忆的检索，以及对偏好学习的一些端侧近似，因为仅靠持续预训练，无法从每个用户产生的少量数据里榨出足够多的东西。