IBM 开源 Bamba:结合变压器与状态空间模型,突破长序列处理效率瓶颈
IBM 研究人员将变压器架构与状态空间模型(SSM)结合,开发出了一种名为“Bamba”的新模型。这一混合实验旨在解决变器在处理长对话时出现的“二次瓶颈”问题,即随着对话长度增加,生成响应的成本和延迟呈平方增长。 背景与起因 自 2022 年由 ChatGPT 推广以来,基于变压器的大语言模型(LLM)因其能够生成类似人类的文本而受到广泛关注。然而,变压器的一大缺陷在于其记忆机制。每次生成响应时,它需要将整个对话序列存储在 KV(键值)缓存中,这导致了内存需求的急剧增加和计算效率下降。相比之下,SSM 通过维护一个固定大小的隐藏状态来总结过去的对话信息,从而减少记忆开销并提高推理速度。2021 年,斯坦福大学的 Albert Gu 及其团队发布了 S4,这是一种将状态变量应用于语言处理的 SSM。此后,研究人员开始探索将变压器和 SSM 混合的可行性。 经过与结果 2023 年,CMU 的 Gu 教授与普林斯顿大学的 Tri Dao 发布了一个改进的 SSM 版本——Mamba2,进一步激发了混合模型的发展。同年,IBM 研究员安基特·古普塔(Ankit Gupta)提出了一个门控机制,使 SSM 能够更好地过滤无关信息,匹配变压器的序列建模能力。IBM 研究团队注意到 Nvidia 的结果显示这些混合模型在处理长对话方面表现出色,决定基于 Mamba2 架构开发自己的混合模型 Bamba-9B。 Bamba-9B 通过简化和优化 SSM,显著减少了变压器的 KV 缓存需求。据 IBM 研究员拉格胡·甘蒂(Raghu Ganti)介绍,这是 Bamba 获得更高吞吐量、更低延迟和更长上下文处理能力的关键。该模型最初在 2 万亿个词符上进行了训练,后来又增加了 1 万亿个词符,并通过量化技术将其从 18GB 缩小到 9GB。在关键基准测试中,Bamba 的表现与 Meta 的 Llama-3.1 8B 模型相当,尽管后者的训练数据量是 Bamba 的七倍之多。 为了使 Bamba 更容易使用,IBM 团队与 Red Hat 合作,优化了 vLLM(虚拟大语言模型)的运行支持。vLLM 是目前开源领域最流行的大语言模型推理服务器之一,但 SS M 的集成相对复杂,因为需要特殊的状态管理机制。最终,Bamba 在 4,000 个词符的训练序列下能够处理 32,000 个词符的对话,未来还可能扩展到 1 百万个词符以上,运行速度有望提升至变压器的五倍之快。 评价与公司背景 业内人士对 Bamba 的发布给予了高度评价,认为其创新性的设计和开放源代码为解决变压器的二次瓶颈问题提供了新的思路。IBM 一直致力于为企业用户开发更高效、更强大的大语言模型,其最新一代 Granite LLM 系列产品就是这一努力的成果。Bamba 作为 IBM 下一代 Granite 4.0 模型的一部分,不仅展示了其技术实力,也为行业内的其他研究者和开发者提供了一个有价值的参考。IBM 科学家们相信,随着 vLLM 对 SSM 支持的增强,Bamba 将进一步展示其潜力,带来更加流畅的语言处理体验。