Jamba是AI21 Labs于2024年2月发布的大型语言模型。它以其混合架构而著称,该架构将Mamba状态空间模型(SSM)与传统Transformer相结合,并拥有高达256,000个令牌的大上下文窗口。该模型以开放权重发布,允许研究人员和开发者下载并使用,同时采用了混合专家(MoE)设计,以在性能与计算效率之间取得平衡。
此次发布将Jamba定位为探索标准Transformer架构替代方案的重要一步,该架构支撑着当代大多数大型语言模型。通过将提供长序列线性时间处理的Mamba SSM与Transformer注意力机制的既有优势相结合,AI21 Labs旨在应对处理超长上下文时的计算挑战。
背景与发展
AI21 Labs由Yoav Shoham、Ori Goshen和Amnon Shashua于2017年11月在以色列特拉维夫创立,此前已开发了Jurassic系列语言模型。该公司的早期工作,包括2021年8月推出的Jurassic-1和2023年3月推出的Jurassic-2,奠定了其在自然语言处理领域的专业地位。Jamba代表了对这些纯Transformer模型的突破,融入了由Albert Gu和Tri Dao等研究人员于2023年底引入的Mamba架构。
Jamba的开发是更广泛研究趋势的一部分,该趋势探索状态空间模型作为基于注意力机制的潜在补充或替代方案。Mamba架构以其选择性状态空间而闻名,提供了一种处理序列的方式,其计算复杂度随序列长度线性扩展,而不同于标准注意力的二次方扩展。AI21 Labs旨在将这种效率与Transformer在需要复杂推理和上下文理解的任务中的成熟能力相结合。
架构与技术细节
Jamba的架构是Mamba SSM层与Transformer层的混合体,以交错模式排列。该模型在其前馈层中采用了混合专家(MoE)方法,允许每个令牌仅激活参数子集,从而在不牺牲模型容量的情况下提高推理效率。总参数数量庞大,但由于MoE设计,每个令牌的活跃参数显著较低。
256,000个令牌的上下文窗口是一个关键卖点,使模型能够单次处理超长文档,如整本书或大规模代码库。这一能力得益于Mamba层对长距离依赖的高效处理。该模型以多种尺寸发布,最大版本拥有520亿总参数,其中每个给定令牌激活120亿参数。
发布与可用性
AI21 Labs于2024年2月28日宣布了Jamba,并在开放许可下公开了模型权重。这使研究社区能够在其自有基础设施上微调和部署该模型。发布内容包括模型检查点和文档,模型可在Hugging Face等平台上获取。开放权重的方法与OpenAI和Anthropic等竞争对手提供的封闭模型形成对比,尽管AI21 Labs也通过其AI21 Studio平台提供商业访问。
发布时机将Jamba置于由OpenAI的GPT-4、Anthropic的Claude和Google DeepMind的Gemini主导的竞争格局中。虽然Jamba在整体基准性能上未达到这些前沿模型,但其独特架构和长上下文窗口使其成为该领域的重要贡献,尤其适用于需要广泛文档处理的应用。
性能与基准
在AI21 Labs报告的基准测试中,Jamba在涉及长上下文理解的任务(如长文档摘要和问答)中表现出色。该模型在性能上与类似规模的其他开放权重模型(如Mixtral 8x7B)具有竞争力,同时提供更长的上下文窗口。在MMLU和HellaSwag等标准自然语言处理基准上,Jamba的表现稳健但并非最先进,反映了其混合设计的权衡。
Mamba层带来的效率提升在吞吐量测试中得到突出体现,Jamba在令牌生成速度上高于同等规模的纯Transformer模型。这使其在资源受限环境中的部署具有吸引力,例如边缘设备或需要实时响应的应用。
接受度与影响
Jamba的发布引起了机器学习社区的兴趣,尤其是探索Transformer架构替代方案的研究人员。模型的开放权重促进了混合SSM-Transformer设计的进一步研究,多项后续工作引用了Jamba作为灵感。该模型还证明了状态空间模型可扩展到生产级规模,鼓励了对这一研究方向的投入。
批评者指出,Jamba在某些基准上的性能未达到领先封闭模型,且其架构的复杂性可能对微调和部署构成挑战。然而,模型的长上下文窗口和效率被视为真正的优势,AI21 Labs继续通过后续发布迭代设计。
后续发展
在初始Jamba发布后,AI21 Labs继续开发模型系列。2024年9月,该公司与Amazon Web Services合作,在Amazon Bedrock上提供Jamba 1.5系列,将模型集成到企业应用中。2025年3月,AI21 Labs发布了Jamba 1.6,用于私有企业部署,声称在多个基准上改进了性能。同月,该公司还推出了AI规划系统Maestro。
这些后续版本完善了混合架构并扩展了模型能力,尽管将Mamba SSM与Transformer层相结合的核心创新仍是中心。Jamba系列确立了AI21 Labs在开放权重模型领域的重要地位,与Meta和Mistral AI等组织并列。
在AI研究中的意义
Jamba的发布促进了深度学习架构创新的持续探索。混合方法的成功表明,未来大型语言模型可能不仅依赖注意力机制,还可能纳入其他序列处理技术。该模型还凸显了AI开发中效率的重要性,因为对训练和部署大型模型计算成本的担忧日益增长。
开放权重的发布使学术机构和独立研究人员能够研究模型行为,有助于更广泛地理解不同架构组件如何相互作用。这种透明度在众多领先模型封闭的领域中备受重视,并帮助民主化了先进AI技术的获取。