Jamba (language model) — 混合架构模型
Jamba是由以色列人工智能研究公司AI21 Labs开发的一个大型语言模型 (LLM) 家族。Jamba是首个采用混合架构的模型，它结合了AI开发中两种主流方法的关键元素：Transformer和状态空间模型 (State Space Models, SSM)，特别是Mamba架构[1]。
Jamba的主要目标是解决现代LLM面临的一个根本性权衡：一方面是Transformer架构所具备的高质量和高性能，另一方面是SSM架构所具备的高效率和处理超长上下文的能力。通过结合这两种方法，并借助Mixture-of-Experts (MoE)引入稀疏性，Jamba提供了一个既强大又高效，并且能够单次处理海量文本的模型。
Jamba架构详解
Jamba并非简单地交替使用Transformer层和Mamba层，而是采用了一种精心设计的块状结构，每个块由八个层组成。
Jamba单个块的结构：
- 一个Transformer层： 该层负责“深度”理解和复杂推理，并内置了Mixture-of-Experts (MoE)架构。
- 七个Mamba层： 这些层紧随Transformer层之后，负责高效处理序列并通过长上下文“传递”信息[2]。
这种非对称结构使模型能够有效管理计算资源：计算量大但功能强大的Transformer操作执行频率较低，而轻量且快速的Mamba操作执行频率较高。
Mixture-of-Experts (MoE) - 专家混合模型集成
Jamba利用MoE架构进一步提升效率。
- MoE仅应用于Transformer层内的全连接块 (FFN)[3]。Mamba层保持密集结构。
- 在Jamba的第一个模型中，设有16个专家。
- 对于每个token，路由网络会选择2个最佳专家（Top-2 gating）。
这意味着，尽管模型总参数量很大（520亿），但在Transformer层的每个token处理步骤中，16个专家中只有2个处于激活状态，从而使计算速度非常快。
Jamba模型演进
Jamba-v0.1 (2024年3月)
作为该系列的首个模型，其技术规格如下：
| 特性 | 值 |
|---|---|
| 总参数量 | 520亿 |
| 激活参数量 | 约120亿 |
| 专家数量 (MoE) | 16个（2个激活） |
| 上下文窗口 | 256,000个token |
| 许可证 | Apache 2.0[4] |
得益于其混合架构，Jamba能够处理长达256,000个token的上下文，相当于一本约400页的小说，并且可以部署在单块80GB显存的消费级GPU上[5]。
Jamba-1.5 (2024年)
2024年，AI21 Labs推出了更新的Jamba 1.5模型家族，包括两个版本：Jamba 1.5 Mini（总参数520亿，激活参数120亿）和Jamba 1.5 Large（总参数3980亿，激活参数940亿）[6]。这些模型在性能上表现出显著提升：
- 在长上下文推理方面，速度比竞争对手快2.5倍。
- 支持包括英语、西班牙语、法语和阿拉伯语在内的九种语言[7]。
关键优势与性能
- 巨大的上下文窗口： 256,000个token——在发布时是所有可用模型（包括专有模型）中最大的上下文窗口之一。这使得Jamba成为分析大型文档（如法律合同、科研论文、完整代码库或长对话）的理想选择。
- 高性能与高效率： 测试表明，Jamba的性能与Llama和Mixtral等同规模的领先开源模型相当或更优，同时在长上下文处理上表现出高出3倍的吞吐量[8]。
- 开放性与可及性： Jamba采用宽松的Apache 2.0许可证发布，允许用于商业和研究目的。模型权重可在Hugging Face平台上获取。
基准测试结果
Jamba 1.5在多项基准测试中展现了有竞争力的结果[9]：
- Jamba 1.5 Mini在Arena Hard上获得46.1分，成为同类公共模型中的领先者[10]。
- Jamba 1.5 Large在Arena Hard上获得65.4分，超过了Llama 3.1 70B和405B。
应用与可用性
Jamba针对商业应用进行了优化，支持函数调用、JSON格式的结构化输出以及文档处理等功能。该模型已在多个平台上线，包括：
- Hugging Face
- Google Cloud Vertex AI
- Microsoft Azure
- NVIDIA API catalog
- Amazon Bedrock[9]
- AI21 Studio
为了支持经济高效的推理，AI21 Labs推出了ExpertsInt8——一种新的量化技术，它允许在拥有8个80GB GPU的机器上部署Jamba 1.5 Large，并在处理256K token上下文时不会造成质量损失[11]。
参考文献
- Lieber, O.; et al. (2024). Jamba: A Hybrid Transformer‑Mamba Language Model. arXiv:2403.19887.
- Lieber, O.; et al. (2024). Jamba‑1.5 Models and ExpertsInt8 Quantization. OpenReview JFPaD7lpBD.
- Gu, A.; Dao, T. (2023). Mamba: Linear‑Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
- Gu, A.; et al. (2021). S4: Efficiently Modeling Long Sequences with Structured State Spaces. arXiv:2111.00396.
- Fedus, W.; et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.
- Yun, L.; et al. (2024). Toward Inference‑Optimal Mixture‑of‑Expert Large Language Models. arXiv:2404.02852.
- Liu, J.; et al. (2024). A Survey on Mixture of Experts in Large Language Models. arXiv:2407.06204.
- Gupta, V.; et al. (2024). Lynx: Enabling Efficient MoE Inference through Dynamic Batch‑Aware Expert Selection. arXiv:2411.08982.
- Liu, J.; et al. (2024). A Survey on Inference Optimization Techniques for Mixture of Experts Models. arXiv:2412.14219.
- Hsieh, C.‑P.; et al. (2024). RULER: What's the Real Context Size of Your Long‑Context Language Models?. arXiv:2404.06654.
注释
- ↑ “宣布推出Jamba：AI21 Labs开创性的SSM-Transformer模型”。AI21 Labs博客。[1]
- ↑ Lieber, O., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv:2403.19887.
- ↑ “Jamba Documentation”. Hugging Face Transformers. [2]
- ↑ “ai21labs/Jamba-v0.1”. Hugging Face. [3]
- ↑ “AI21 Labs' Jamba: A New Hybrid LLM Architecture”. Gradient Flow. [4]
- ↑ “宣布推出Jamba-1.5模型家族”。AI21 Labs博客。[5]
- ↑ “ai21labs/AI21-Jamba-Large-1.5”. Hugging Face. [6]
- ↑ “AI21 Labs借助Jamba打破新壁垒”。ITinAI。[7]
- ↑ 9.0 9.1 “Jamba 1.5发布：来自AI21 Labs的混合模型”。Дзен。[8]
- ↑ “Jamba-1.5 family of models by AI21 Labs is now available in Amazon Bedrock”. AWS What's New. [9]
- ↑ “ExpertsInt8: A new paradigm for efficient inference of MoE-based LLMs”. OpenReview. [10]