IBM 将 Transformer 与 SSM 结合，得到了 'Bamba'

IBM Research 与 CMU、Princeton 和 University of Illinois 合作，构建了一个开源的 LLM，它结合了 Transformer 的表达能力和 state-space model (SSM) 的运行时速度。关键特性将很快添加到 IBM Granite 4.0 中。

通过这篇技术博客了解更多关于 Bamba 的信息，并查看基准测试结果。

欢迎参加 5 月 7 日在曼哈顿 1 Madison Avenue 的 IBM Innovation Studio 举行的首届 vLLM meetup in NYC！本次活动由 IBM 和 Red Hat 主办，将包含技术讲座和讨论，探讨如何优化 LLM 推理以提高性能和效率。

如今大型语言模型背后的 Transformer 架构已经展示了生成类似人类文本的惊人能力。其有效性的一部分来自其自注意力机制，该机制允许模型在生成响应时权衡输入序列中的所有单词。

但问题在于，随着对话变得越来越长，由于模型在响应时会将正在运行的序列保存在内存中，因此生成的累计成本呈二次方增长。如果上下文窗口的大小加倍，则处理上下文和生成响应的成本不仅仅是加倍，而是变为四倍。

这种“二次瓶颈”通常是提问模型和获得答案之间令人沮丧的延迟的原因。它还会产生大量冗余计算。在 ChatGPT 于 2022 年普及 Transformer 之后，研究人员已经在寻找替代架构。

state-space models (SSMs) 以及与 SSM 层交错的 Transformer 已经成为两种可能的解决方案。IBM Research 刚刚开源了其首个混合实验：Bamba，该模型可以像 SSM 一样快速运行，并像 Transformer 一样熟练地处理长序列。Bamba 的许多创新都是 IBM 即将在几个月后推出的下一代 Granite 4.0 模型的一部分。

通过显着降低 Transformer 的 KV (key value) 缓存内存的内存需求，Bamba-9B 表明它可以至少以与类似大小的 Transformer 两倍的速度运行，同时保持其准确性。“一切都回到了 KV 缓存减少，”领导该项目的 IBM 研究员 Raghu Ganti 说道。“更高的吞吐量、更低的延迟、更长的上下文长度。”

你从未听说过的最重要的模型

state-space models 在名称认知度方面远不及 Transformer，但它们已被使用了数十年，用于对动态系统进行建模。

“它们是电气工程的面包和黄油——信号处理、机器人技术和控制理论，”IBM 研究员 Ankit Gupta 说道，他一直在将 SSM 适应于深度学习中发挥了关键作用。“任何使用时间序列数据的领域都会使用 state-space models 来分析它。”

SSM 构建的数学方程式可以应用于大脑中的电活动、天气，甚至股市。SSM 从一系列观察结果中计算出一个固定大小的“隐藏状态”，捕获系统的基本属性。可以将状态视为过去的摘要。当新数据进入时，隐藏状态会更新，而不会增加其大小，同时预测接下来会发生什么。

SSM 于 2021 年跨界到神经网络，当时 Albert Gu 及其在 Stanford 的合作者发布了 S4，这是一个将状态变量应用于语言的 SSM。与 Transformer 和之前的 recurrent neural network (RNNs) 一样，SSM 擅长处理单词序列。但它可以比 RNN 更熟练地处理长序列，并且比 Transformer 快得多。

当 Transformer 在输出响应时关注上下文窗口中的所有单词时，SSM 维护一个压缩的隐藏状态，该状态总结了过去的信息。这种对信息的选择性保留需要更少的内存开销，并导致更快的推理速度。

S4 在突然出现在 Long Range Arena 时引起了轰动，Long Range Arena 是一个通过处理长序列的能力来比较语言模型的基准，但它很难实现。然后，一位 IBM 的 AI 研究员 Gupta 帮助 Gu 团队使用 diagonal state spaces 简化了模型。他们的“对角线”SSM 将 S4 的 1,000 行代码缩减到 10 行。Gupta 后来帮助引入了 a gating mechanism 用于过滤掉不相关的信息，从而首次允许 SSM 匹配 Transformer 的“表达能力”或序列建模技能。

该团队还推出了可能是首个混合 Transformer。“探索混合模型是有意义的，”Gupta 说道，他现在在 IBM 的 Granite Vision models 上工作。“我们可以使用标准的注意力块来处理具有局部依赖关系的文本，同时利用 SSM 来进行更长范围的上下文关联。”

2023 年，当时是 CMU 教授的 Gu 和 Princeton 的 Tri Dao 推出了一个 gated SSM 变体——Mamba2，这有助于激发了混合模型的热潮，其名称如 Samba 和 MambaFormer。去年，Nvidia 证实，这些新的混合模型可以胜过它们自己的架构，同时显着 加速推理，最终发布了 Nemotron-H。

克服 KV 缓存瓶颈

从一开始，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 为大规模分布式训练而设计的数据加载器，以及旨在降低存储和推理成本的量化框架。

他们最初在 2 万亿个 token（单词和部分单词）上训练了 Bamba。在结果的鼓舞下，他们又添加了 1 万亿个 token，并通过量化将模型从 18 GB 缩小到 9 GB，将其位宽从 Mamba2 的 16 位浮点精度降低到 8 位。在关键基准测试中，Bamba 的表现与 Meta 的 Llama-3.1 8B 模型相当，后者在七倍的数据上进行了训练——Ganti 将这一成就归功于 Bamba 的设计和高质量的训练数据。

他们的下一个障碍是优化 vLLM 以运行 SSM。“Virtual” LLM 已经成为 LLM 的首选开源推理服务器，而 Bamba 背后的团队与 Red Hat 密切合作，将该模型集成到平台中。“SSM 很难支持，因为你需要定制的状态管理，”Red Hat 的技术人员和 vLLM 的提交者 Tyler Smith 说道。

当 Bamba 在去年年底发布时，Ganti 邀请开源社区帮助改进它。“让我们一起克服 KV 缓存瓶颈！”他在 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 的二次瓶颈也可以这样说。

订阅我们的 Future Forward 新闻通讯，并及时了解最新的研究新闻

日期

2025 年 4 月 29 日

作者

主题