Bamba:IBM 混合架构开源大模型
融合 Transformer 和 SSM 的新型语言模型,性能与闭源大模型相当。新颖的架构设计值得关注。
融合 Transformer 和 SSM 的新型语言模型,性能与闭源大模型相当。新颖的架构设计值得关注。
当今大语言模型背后的 transformer 架构,展现出了生成类人文本的惊人能力。它之所以如此有效,部分归功于 self-attention 机制:模型在生成回答时,可以衡量输入序列中所有词的重要性。
随着对话越来越长,问题也随之而来。由于模型在回答过程中需要将不断增长的序列保存在内存里,生成的累计成本会呈二次方增长。如果上下文窗口扩大一倍,处理上下文并生成回答的成本并不只是翻倍,而是变成四倍。
这种“二次方瓶颈”,往往正是用户提出问题后,需要恼人地等待一段时间才能得到回答的原因。它还会造成大量重复计算。到 2022 年 ChatGPT 让 transformer 广为人知时,研究人员其实早已开始寻找替代架构。
状态空间模型(State-Space Model,SSM),以及将 transformer 与 SSM 层交错组合的架构,已经成为两种可能的解决方案。IBM Research 刚刚开源了它的第一个混合架构实验:Bamba。这个模型既能像 SSM 一样快速运行,也能像 transformer 一样出色地处理长序列。Bamba 的许多创新也将被纳入 IBM 计划于数月后推出的下一代 Granite 4.0 模型。
通过大幅降低 transformer 的 KV(key value)cache 内存需求,Bamba-9B 已经证明,它的运行速度至少可以达到同等规模 transformer 的两倍,同时保持相当的准确率。“一切最终都归结为 KV cache 的缩减,”该项目负责人、IBM 研究员 Raghu Ganti 表示,“更高的吞吐量、更低的延迟,以及更长的上下文长度。”
在 YouTube 上观看这段视频——Bamba:IBM 最新 LLM 介绍。
状态空间模型的知名度远不及 transformer,但几十年来,人们一直在使用它对动态系统进行建模。
“它们是电气工程的基本工具,广泛用于信号处理、机器人和控制理论,”IBM 研究员 Ankit Gupta 表示。他在将 SSM 引入深度学习领域的过程中发挥了关键作用。“任何使用时间序列数据的领域,都会使用状态空间模型进行分析。”
SSM 所基于的数学方程,可以应用于大脑中的电活动、天气,甚至股票市场。SSM 会根据一系列观测结果,计算出一个大小固定的“隐藏状态”,用来捕捉系统的关键属性。你可以把这个状态理解为对过去的总结。当新数据到来时,隐藏状态会随之更新,并预测接下来将发生什么,但其大小不会增加。
2021 年,Albert Gu 及其在 Stanford 的合作者发布了 S4,将状态变量应用于语言,SSM 由此进入神经网络领域。和 transformer 以及更早之前的循环神经网络(RNN)一样,SSM 擅长处理词语序列。但相比 RNN,它更善于处理长序列;相比 transformer,它的速度则快得多。
transformer 在输出回答时,会关注上下文窗口中的所有词;SSM 则维护一个经过压缩的隐藏状态,用于总结过去的信息。这种选择性保留信息的方式只需要更少的额外内存,也能带来更快的推理速度。
S4 突然出现在 Long Range Arena 上时引发了轰动。Long Range Arena 是一个通过处理长序列的能力来比较语言模型的 benchmark,但 S4 很难实现。当时还是 IBM AI resident 的 Gupta,帮助 Gu 及其团队利用对角状态空间简化了模型。他们的“对角”SSM 将 S4 的 1,000 行代码缩减到了 10 行。后来,Gupta 又协助引入了一种用于过滤无关信息的 gating 机制,使 SSM 首次能够在“表达能力”,也就是序列建模能力上,与 transformer 相匹敌。
该团队还发布了可能是首个混合 transformer 的模型。“探索混合架构是顺理成章的,”Gupta 表示。他目前参与 IBM Granite Vision 模型的研发。“我们可以使用标准 attention block 处理具有局部依赖关系的文本,同时利用 SSM 完成更长距离的上下文化。”
2023 年,当时任 CMU 教授的 Gu 与 Princeton 的 Tri Dao 发布了带有 gating 机制的 SSM 变体 Mamba2。它推动了一波混合架构浪潮,其中包括 Samba 和 MambaFormer 等模型。去年,Nvidia 证实,这些新的混合架构不仅可以超越任意一种单独架构,还能显著提升推理速度,并最终发布了 Nemotron-H。
从一开始,IBM Research 就将效率作为其企业级 Granite LLM 的核心。随着 IBM Granite 变得更小、能力更强,研究人员开始把目标对准二次方瓶颈。Nvidia 的研究结果发布后,IBM 研究人员在内部对其进行了验证,随后开始构建自己的混合架构 Bamba-9B。
他们找到了 Mamba 的创造者 Gu 和 Dao,以及 University of Illinois at Urbana-Champaign 的教授 Minjia Zhang。团队共同选择了 Nvidia 的 Mamba2 架构,并决定将几乎所有与 Bamba 相关的内容开源,包括训练方案、数据、IBM 为大规模分布式训练设计的 data loader,以及一个旨在降低存储和推理成本的量化框架。
他们最初使用 2 万亿个 token(词与词的组成部分)训练 Bamba。受到结果鼓舞后,他们又增加了 1 万亿个 token,并通过量化将模型从 18 GB 缩减到 9 GB,把位宽从 Mamba2 的 16-bit 浮点精度降至 8-bit。在关键 benchmark 上,Bamba 的表现与 Meta 的 Llama-3.1 8B 模型相当,而后者使用的训练数据是 Bamba 的七倍。Ganti 将这一成果归功于 Bamba 的设计与高质量训练数据。
他们面临的下一个障碍,是优化 vLLM 以运行 SSM。“Virtual” LLM 已经成为 LLM 首选的开源 inference server,Bamba 背后的团队与 Red Hat 密切合作,将该模型集成到这个平台中。“SSM 很难支持,因为你需要专门定制的状态管理,”Red Hat 技术人员、vLLM committer Tyler Smith 表示。
Bamba 于去年年底发布时,Ganti 邀请开源社区帮助改进它。“让我们一起突破 KV-cache 瓶颈!”他在 Hugging Face 上的 Bamba 介绍中写道。
Bamba 使用长度为 4,000 个 token 的序列进行训练,却能够出色地处理包含 32,000 个 token 的对话。但 Ganti 认为,它可以扩展到 100 万个 token 甚至更多;随着 vLLM 加入更多对 SSM 的支持,其运行速度最高可以达到 transformer 的五倍。
墨西哥民歌《La Bamba》由 Ritchie Valens 唱红,其中的副歌是:Para bailar La Bamba/Se necesita una poca de Gracia。想踩着节拍起舞,只需要一点优雅。
要突破 transformer 的二次方瓶颈,也可以这么说。
新闻
Kim Martineau
2026 年 7 月 23 日
AI
AI 透明度
公平、问责与透明
生成式 AI
公平、问责与透明
新闻
2026 年 7 月 22 日
AI
量子计算
发布
Kim Martineau
2026 年 7 月 16 日
AI
AI for Code
生成式 AI
自然语言处理
自然语言处理
研究
Peter Hess
2026 年 7 月 9 日
AI
生成式 AI