突破内存限制：Jamba模型的高效文本处理能力

人工智能咨询培训老师叶梓 转载标明出处

在当今信息爆炸的时代，处理和理解海量文本数据的需求日益增长。自然语言处理（NLP）领域的研究者们一直在探索如何构建更高效、更强大且更灵活的语言模型来应对这一挑战。然而，现有的大型语言模型，尤其是基于Transformer架构的模型，虽然在多个任务上取得了显著的成就，但它们在处理长文本时仍面临着内存和计算资源的巨大需求。这些需求限制了模型在资源受限的环境中的应用，并可能导致推理速度变慢，影响用户体验。

为了解决这些问题，一种新型的混合Transformer-Mamba语言模型——Jamba应运而生。Jamba模型通过结合Transformer的注意力机制和Mamba的状态空间建模能力，实现了对长文本更有效的处理，降低了内存占用，并提高了推理速度和模型的吞吐量。这一突破性进展不仅为NLP领域带来了新的技术革新，也为构建更加高效和实用的语言模型铺平了道路。

模型架构

Jamba架构图

Figure 1展示了Jamba模型的基本构成单元和不同类型的层。Jamba模型的结构组成如下：

(a) 单一Jamba块

一个Jamba块是构成Jamba模型的基本单元。在这个块中，包含了多层的网络结构，这些层可以是Transformer层，也可以是Mamba层，或者是混合了专家（MoE）的层。这些层按照特定的比例和顺序排列，以实现模型的最佳性能。

(b) 不同类型的层

Jamba模型中包含几种不同类型的层：

1. Transformer层：这是传统Transformer模型中使用的层，主要负责处理输入数据的自注意力（Self-Attention）机制，允许模型在生成输出时考虑到序列中的所有位置。

2. Mamba层：Mamba层是一种状态空间模型（State-Space Model, SSM），相比于Transformer层，它在处理长序列数据时更为高效，并且能够更好地捕捉长距离依赖关系。

3. MoE层：即混合专家层，是一种引入模型稀疏性的技术，通过多个“专家”网络来处理不同的输入部分，然后通过一个门控机制（Gating Mechanism）来选择最合适的专家输出，从而提高模型的容量和灵活性。

Jamba块具有以下特征：

• l = 8：表示每个Jamba块包含8层。
• a : m = 1 : 7：表示在这8层中，有1层是Transformer的注意力层，而有7层是Mamba层。这个比例旨在平衡模型的计算效率和对长距离依赖的处理能力。
• MoE应用频率 e = 2：表示每2层就会有一个MoE层，这意味着在8层的Jamba块中，有4个MoE层被应用。这样的设置允许模型在保持参数数量可控的同时，增加模型的容量和灵活性。

这种结构设计使得Jamba模型能够结合Transformer和Mamba层的优势，同时通过MoE层提高模型的扩展性和适应性。通过这种方式，Jamba能够在保持较小的内存占用和高效的计算性能的同时，实现对长上下文的理解和处理。

Jamba模型的核心创新在于其混合架构，该架构巧妙地融合了Transformer和Mamba两种不同的神经网络层。Transformer层以其高效的注意力机制而闻名，而Mamba层则是一种新型的状态空间模型，擅长捕捉序列数据中的长距离依赖关系。通过这种混合设计，Jamba能够同时利用两种模型的优势，实现对长文本的高效处理。

在Jamba模型中，Transformer层和Mamba层以特定的比例交替出现。这种设计允许模型在保持Transformer层强大的并行处理能力的同时，通过Mamba层引入对长距离依赖的敏感性。这种协同作用不仅提高了模型对上下文的理解能力，还显著降低了处理长序列时所需的内存和计算资源。

为了进一步提升模型的容量和灵活性，Jamba模型引入了混合专家（MoE）技术。MoE允许模型在不同的专家网络中分配计算任务，每个专家可以专注于解决特定类型的问题。这种设计不仅增加了模型的总参数数量，而且通过智能路由机制，确保了在每次前向传播中只激活一部分参数，从而有效控制了计算成本。

Jamba模型的另一个显著特点是其高度的可配置性。模型的设计者可以根据不同的硬件资源和性能要求，调整模型中的参数，如层数、注意力与Mamba层的比例、MoE的使用频率等。这种灵活性使得Jamba能够适应各种不同的应用场景，从资源受限的移动设备到高性能的服务器环境。

在Jamba模型的设计中，特别关注了内存使用和吞吐量的优化。通过减少关键值（KV）缓存的大小，Jamba显著降低了模型的内存占用。同时，通过优化Mamba层的计算效率，Jamba在处理长序列时展现出了更高的吞吐量，这对于需要快速响应的实时应用场景尤为重要。

表格比较了Jamba与其他几个最近公开的模型（如LLAMA-2、Mistral、Mixtral）在总可用参数、活跃参数和长上下文（256K tokens）下的KV缓存大小

高效部署实现与性能优化

Jamba模型特别针对单个80GB GPU进行了优化配置，以实现在保证质量和吞吐量的同时，最大程度地适应硬件限制。具体来说，Jamba由4个Jamba块组成，每个块包含8层，其中包括1:7的注意力层到Mamba层的比例，以及每两层使用一次MoE（专家混合）代替单个MLP（多层感知器）。这种配置不仅使得模型能够适应单个80GB GPU的内存限制，同时还保持了高效的计算性能。

该图比较了Jamba模型实现与Mixtral 8x7B和Llama-2-70B在单个80GB A100 GPU上支持的最大上下文长度。显示了Jamba可以支持比Mixtral多两倍、比Llama-2-70B多七倍的上下文长度

Jamba模型在处理长序列数据时展现出了卓越的吞吐量。在不同的批处理大小和上下文长度设置下，Jamba的吞吐量表现均优于现有的Mixtral-8x7B和Llama-2 70B模型。特别是在长上下文处理方面，Jamba的吞吐量是Mixtral的三倍，这一优势在处理超过128K个token的上下文时尤为明显。

图(a)展示了在不同批量大小下，使用单个A100 80GB GPU、int8量化、8K上下文长度生成512个token的输出时，Jamba与其他模型的吞吐量比较。图(b)展示了在单个批量、4个A100 GPU、不使用量化、不同上下文长度下，生成512个token的输出时的吞吐量比较。Jamba在长上下文上展现出更高的吞吐量

Jamba模型的训练采用了NVIDIA H100 GPU，并使用了一种专有的内部框架，该框架支持大规模训练，包括FSDP、张量并行、序列并行和专家并行等技术。Jamba的训练数据集包含了来自网络、书籍和代码的文本数据，这些数据经过了质量筛选和去重处理。

Jamba模型的训练得益于NVIDIA H100 GPU的强大计算能力，以及开发团队自研的高效训练框架，该框架集成了全参数数据并行、张量并行、序列并行和专家并行等先进技术，确保了大规模训练的效率。模型所依赖的内部数据集涵盖了来自网络、书籍和代码的文本数据，且数据集经过了最新的更新和严格的质量筛选与去重处理，以保证训练数据的质量和多样性。这些因素共同为Jamba模型的优异性能打下了坚实的基础。

评估

Jamba在多个标准的学术基准测试上进行了评估，这些测试覆盖了常识推理、阅读理解、语言理解等多个方面。例如，在HellaSwag、WinoGrande、ARC等测试中，Jamba展现了出色的推理能力；在BoolQ、QuAC等阅读理解测试中，模型的理解和回答问题的能力得到了验证。Jamba还在MMLU和BBH等综合基准测试中表现优异，这些测试综合考察了模型在多个任务上的语言理解能力。

Jamba与其他公开可用的模型在多个学术基准测试上的性能比较。Jamba在保持相似或更好的性能的同时，具有更高的吞吐量

除了标准的学术基准测试，Jamba在处理长上下文数据方面的能力也经过了严格的评估。通过"针堆中找针"的测试，Jamba证明了其在长文本中检索信息的能力。此外，Jamba还在L-Eval中的长上下文问答数据集上进行了评估，这些数据集包括NarrativeQA、LongFQA、Natural Questions等，Jamba在这些测试中的表现进一步证实了其在处理长文本方面的卓越性能。

Jamba在针堆草垛评估中的表现，该评估要求在长上下文窗口中检索一个简单语句。Jamba在这项评估中表现出色，尤其是在只有4个注意力层的实现中

在与其他公开可用的模型进行对比时，Jamba显示出了其在相似参数规模下的性能优势。与Llama-2、Mixtral等模型相比，Jamba在保持较小的总参数量的同时，实现了更高的活跃参数使用效率和更好的吞吐量，同时在多个基准测试中取得了相似或更好的成绩。

Jamba的高效率和低内存占用是其显著特点之一。在处理长达256K个token的上下文时，Jamba的内存占用仅为4GB，远低于其他同类模型。这一优势使得Jamba即使在资源受限的环境中也能高效运行。

Jamba在长上下文问答基准测试中的表现，特别是在3-shot格式下。Jamba在大多数数据集上的表现优于Mixtral，并且平均表现更好

在消融实验中，研究者首先探讨了Transformer注意力层和Mamba层的结合比例对模型性能的影响。实验结果表明，混合模型在1:3和1:7的注意力到Mamba层的比例下表现相似，但1:7的比例在计算效率上更胜一筹。因此，这一比例被选为后续实验的配置。

不同注意力到Mamba层比例（a:m）的模型在学术基准测试和对数概率评估上的结果。1:3和1:7比例的模型表现几乎没有差异

进一步的实验分析了纯Mamba模型在某些任务上表现不佳的原因，尤其是在需要上下文学习能力的任务中。相比之下，混合模型能够成功执行上下文学习，即使只有少数几层是注意力层。这表明注意力机制可能有助于模型更好地学习上下文信息。

1.3B参数的纯Mamba模型在IMDB、QuAC和NarrativeQA数据集上的表现不佳，而Attention-Mamba混合模型的表现与纯Transformer模型相当

研究者还研究了MoE在大规模模型中与混合注意力-Mamba架构结合的效果。实验结果表明，MoE能够显著提升模型的性能，同时保持了计算的可行性。

在7B参数模型上，添加MoE对Attention-Mamba混合模型性能的提升。MoE有助于提高模型在多个基准测试上的表现

在训练大规模模型时，遇到了Mamba层内部激活值过大导致损失激增的问题。为了解决这个问题，研究者引入了RMSNorm来稳定训练过程，有效地防止了损失的剧烈波动。

一个1.3B参数的Attention-Mamba混合模型（无MoE）中的一个示例头在IMDB数据集上的注意力可视化。该图显示了最后一个token（“:”）的注意力集中在少数示例的标签token上

研究者还探讨了Jamba模型是否需要显式的位置信息。实验结果表明，即使没有显式的位置编码，Jamba模型也能取得良好的性能，这表明Mamba层可能已经提供了足够的位置信息。

 Table 8比较了有无显式位置信息（使用RoPE）的Jamba模型在多个基准测试和对数概率评估上的结果。结果表明，对于混合架构而言，可能不需要显式的位置信息

这些实验结果和开发过程中的洞见将为未来混合注意力-状态空间模型的研究提供指导。为了推动这一领域的研究，研究者计划公开小规模训练的模型检查点。发布的最大型模型具备12亿活跃参数和52亿总参数，支持处理长达256K个token的上下文，并且能够在单个80GB GPU上处理140K-token的文本。Jamba模型的成功不仅表现在其卓越的性能上，还在于它为自然语言处理技术的未来研究和应用开辟了新的可能性。

论文链接：https://arxiv.org/abs/2403.19887

GitHub 地址：https://www.ai21.com/jamba