Jamba: A Hybrid Transformer-Mamba Language Model

📄 arXiv: 2403.19887v2 📥 PDF

作者: Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, Yoav Shoham

分类: cs.CL, cs.LG

发布日期: 2024-03-28 (更新: 2024-07-03)

备注: Webpage: https://www.ai21.com/jamba

💡 一句话要点

提出Jamba模型以提升长文本处理能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合模型 长文本处理 专家混合 Transformer Mamba 自然语言处理 模型优化

📋 核心要点

  1. 现有语言模型在处理长文本时面临性能瓶颈，尤其是在上下文长度超过传统限制时。
  2. Jamba模型通过混合Transformer和Mamba层，并引入专家混合机制，提升了模型的容量和效率。
  3. 实验结果显示，Jamba在256K令牌的上下文长度下表现优异，相较于传统Transformer模型具有更高的吞吐量和更小的内存占用。

📝 摘要（中文）

我们提出了Jamba，一种基于新型混合Transformer-Mamba专家混合（MoE）架构的大型语言模型。Jamba交替使用Transformer和Mamba层，结合了两种模型的优点。在某些层中引入MoE以增加模型容量，同时保持可管理的参数使用。该灵活架构支持资源和目标特定的配置。我们实现的特定配置使得模型在单个80GB GPU上运行，具有高吞吐量和小内存占用，同时在标准语言模型基准和长上下文评估中表现出色。模型在256K令牌上下文长度下展现出强大的结果。我们研究了多种架构决策，展示了它们在大规模建模中的重要性，并计划发布各种消融实验的检查点，以鼓励对这一新型架构的进一步探索。

🔬 方法详解

问题定义：现有语言模型在处理长文本时，尤其是上下文长度超过传统限制时，性能显著下降，导致信息丢失和理解能力不足。

核心思路：Jamba模型通过混合Transformer和Mamba层，结合两者的优点，并在部分层中引入专家混合（MoE）机制，以提升模型的容量和效率，从而更好地处理长文本。

技术框架：Jamba的整体架构由交替的Transformer层和Mamba层组成，部分层中集成了MoE机制。该架构允许根据资源和目标进行灵活配置，确保在单个80GB GPU上实现高效运行。

关键创新：Jamba的主要创新在于其混合架构设计，通过有效结合不同类型的层和专家机制，显著提升了模型在长上下文处理上的能力，与传统的Transformer模型相比，提供了更高的吞吐量和更小的内存占用。

关键设计：在模型设计中，关键参数设置包括层的交替方式、专家的混合策略以及损失函数的选择，这些设计决定了模型的性能和效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Jamba模型在标准语言模型基准测试中表现出色，尤其是在长上下文评估中，能够处理高达256K令牌的上下文长度。与传统Transformer模型相比，Jamba在吞吐量和内存占用方面均有显著提升，展示了其在大规模建模中的优势。

🎯 应用场景

Jamba模型在自然语言处理领域具有广泛的应用潜力，尤其是在需要处理长文本的任务中，如文档摘要、长篇对话生成和复杂问答系统。其高效的架构设计使得在资源受限的环境中也能实现高性能的语言理解与生成，未来可能推动更多智能应用的发展。

📄 摘要（原文）

We present Jamba, a new base large language model based on a novel hybrid Transformer-Mamba mixture-of-experts (MoE) architecture. Specifically, Jamba interleaves blocks of Transformer and Mamba layers, enjoying the benefits of both model families. MoE is added in some of these layers to increase model capacity while keeping active parameter usage manageable. This flexible architecture allows resource- and objective-specific configurations. In the particular configuration we have implemented, we end up with a powerful model that fits in a single 80GB GPU. Built at large scale, Jamba provides high throughput and small memory footprint compared to vanilla Transformers, and at the same time state-of-the-art performance on standard language model benchmarks and long-context evaluations. Remarkably, the model presents strong results for up to 256K tokens context length. We study various architectural decisions, such as how to combine Transformer and Mamba layers, and how to mix experts, and show that some of them are crucial in large scale modeling. We also describe several interesting properties of these architectures which the training and evaluation of Jamba have revealed, and plan to release checkpoints from various ablation runs, to encourage further exploration of this novel architecture. We make the weights of our implementation of Jamba publicly available under a permissive license.