文章标题: Jamba: 一种混合Transformer-Mamba语言模型
作者/机构: 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, Avshalom Manevich, Nir Ratner, Noam Rozen, Erez Schwartz, Mor Zusman, Yoav Shoham - Al21labs

A1 主要贡献

本文介绍了Jamba，一种新型的、公开可用的大型语言模型。Jamba基于一种创新的混合架构，该架构融合了Transformer层、Mamba层（一种最新的状态空间模型SSM）以及专家混合（MoE）模块。这种设计结合了两种正交的架构思想，共同提升了模型的性能和吞吐量，同时保持了可管理的内存占用。

核心问题与研究目标:
Transformer架构虽然流行，但存在两大缺陷：
1. 高内存与计算需求：在处理长上下文时，键值（KV）缓存的大小成为限制因素。
2. 推理速度慢、吞吐量低：由于缺乏单一的摘要状态，每个生成词元都需要对整个上下文进行计算。
而传统的循环神经网络（RNN）和现代的状态空间模型（SSM）如Mamba，虽然能有效处理长上下文且推理高效，但在训练并行性和捕捉长距离依赖关系方面不如Transformer，并且在同等规模下性能通常落后于Transformer模型。

Jamba的目标是结合Transformer和Mamba的优点，创造一个既能高效处理长上下文、推理速度快，又能保持顶尖性能的模型。同时，通过引入MoE，进一步在不显著增加计算需求的情况下扩大模型容量。

创新点与主要贡献:
1. 首个生产级的混合注意力-SSM模型: Jamba是首个将Transformer的注意力和SSM（Mamba）成功结合并扩展到大规模的生产级模型。它通过交错排布Transformer和Mamba层，实现了两者的优势互补。
2. 灵活且高效的架构:
- 混合架构: 通过调整Transformer与Mamba层的比例，可以在内存使用、训练效率和长上下文能力之间进行权衡。
- 专家混合（MoE）集成: 在部分层级应用MoE，有效增加了模型总参数量（容量），同时保持了较低的活跃参数量和计算需求。
3. 卓越的性能与效率:
- 高性能: 发布的7B基础模型（12B活跃参数，52B总参数）在标准语言模型基准测试中，性能与同等规模的Mixtral-8x7B和更大规模的Llama-2 70B相媲美。
- 超长上下文支持: 模型支持高达256K词元的上下文长度，在长上下文评测中，多数任务优于Mixtral。
- 高吞吐量与低内存占用: 在长上下文场景下，Jamba的吞吐量是Mixtral-8x7B的3倍。由于Mamba层的引入，其KV缓存极小，使得模型在处理超过128K词元上下文时也能装入单个80GB的GPU。
4. 开源与社区贡献: 团队以Apache 2.0许可证公开发布了Jamba模型（12B活跃/52B总参数）的权重，并计划发布消融实验的检查点，以鼓励社区对这种新颖架构进行进一步的研究、实验和优化。

下图展示了Jamba块的基本结构及其不同类型的层。


图 1: (a) 单个Jamba块。(b) 不同类型的层。这里展示的实现中，l = 8，注意力与Mamba层的比例a : m = 1 : 7，并且每e = 2层应用一次MoE。

A2 方法细节

2. 模型架构

Jamba的混合解码器架构。Jamba是一种混合解码器架构，它融合了Transformer层【索引51，Attention is all you need，2017，Advances in neural information processing systems】、Mamba层【索引17，Mamba: Linear-time sequence modeling with selective state spaces，2023，arXiv】（一种近期的状态空间模型（SSM）【索引18，Efficiently modeling long sequences with structured state spaces，2021，International Conference on Learning Representations；索引19，Combining recurrent, convolutional, and continuous-time models with linear state space layers，2021，Advances in neural information processing systems】），并额外加入了专家混合（MoE）模块【索引14，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，Journal of Machine Learning Research；索引46，Outrageously large neural networks: The sparsely-gated mixture-of-experts layer，2017，International Conference on Learning Representations】。我们将这三个元素的组合称为一个Jamba块，如图1所示。

平衡内存、吞吐量与质量的设计。结合Transformer、Mamba和MoE元素使得架构在低内存使用、高吞吐量和高质量这些时常冲突的目标之间取得了灵活的平衡。在内存使用方面，比较模型的总参数量可能产生误导。在MoE模型中，任何给定前向传播步骤中参与计算的活跃参数数量可能远小于总参数数量。另一个重要考量是KV缓存——即存储上下文中注意力键和值所需的内存。当将Transformer模型扩展到长上下文时，KV缓存成为一个限制因素。通过用Mamba层替换注意力层，可以减少KV缓存的总大小。我们的架构旨在不仅提供少量的活跃参数，而且与纯Transformer模型相比，实现8倍小的KV缓存。表1比较了Jamba与近期公开可用模型的参数和长上下文下的KV缓存，显示了其在维持小KV缓存方面的优势。


表 1: Jamba与近期开放模型在总可用参数、活跃参数和长上下文KV缓存内存方面的比较。Jamba在KV缓存内存需求上实现了显著的减少。

吞吐量优势。在吞吐量方面，对于短序列，注意力操作在推理和训练的浮点运算（FLOPS）中只占一小部分【索引7，Palm: Scaling language modeling with pathways，2023，Journal of Machine Learning Research】。然而，对于长序列，注意力机制会占据大部分计算资源。相比之下，Mamba层在计算上更为高效。因此，增加Mamba层的比例可以提高吞吐量，尤其是在处理长序列时。

Jamba块的基本构成。该模型的基本组件是一个Jamba块，可以顺序重复。每个Jamba块是Mamba层或注意力层的组合。每个这样的层包含一个注意力或Mamba模块，其后跟一个多层感知机（MLP）。图1(b)展示了不同类型的层。一个Jamba块包含l个层，这些层以a:m的比例混合，即每m个Mamba层对应a个注意力层。

MoE模块的集成。在Jamba中，一些MLP可以被MoE层取代，这有助于在保持活跃参数数量（从而保持计算量）较小的情况下增加模型容量。MoE模块可以每e层应用于MLP。使用MoE时，每层有n个可能的专家，路由器在每个词元处选择前K个专家。总结来说，Jamba架构中的不同自由度包括：
- l: 层的数量。
- a:m: 注意力层与Mamba层的比例。
- e: 使用MoE替代单个MLP的频率。
- n: 每层的总专家数。
- K: 每个词元使用的顶尖专家数。

设计空间的灵活性与权衡。基于这个设计空间，Jamba在优先考虑某些属性上提供了灵活性。例如，增加m并减少a，即以牺牲注意力层为代价增加Mamba层的比例，会减少存储键值缓存所需的内存。这降低了整体内存占用，对于处理长序列尤其重要。增加Mamba层的比例也提高了吞吐量，特别是在长序列上。然而，减少a可能会降低模型的能力。

MoE参数的平衡。此外，平衡n、K和e会影响活跃参数和总可用参数之间的关系。较大的n以内存占用为代价增加了模型容量，而较大的K则增加了活跃参数的使用和计算需求。相反，较大的e会降低模型容量，同时减少计算（当K>1时）和内存需求，并允许更少的通信依赖（在专家并行训练和推理期间减少内存传输以及GPU间的通信）。

Mamba层的实现细节。Jamba对Mamba层的实现包含了几种归一化方法，有助于在大型模型规模下稳定训练。特别地，我们在Mamba层中应用了RMSNorm【索引53，Root mean square layer normalization，2019，Advances in Neural Information Processing Systems】。

位置信息的处理。我们发现，对于Mamba层，位置嵌入或像RoPE【索引47，Roformer: Enhanced transformer with rotary position embedding，2024，Neurocomputing】这样的机制是不必要的，因此我们没有使用任何显式的位置信息。

其他标准架构细节。架构的其他细节是标准的，包括分组查询注意力（GQA）、SwiGLU激活函数【索引7，Palm: Scaling language modeling with pathways，2023，Journal of Machine Learning Research；索引45，Glu variants improve transformer，2020，arXiv；索引50，Llama 2: Open foundation and fine-tuned chat models，2023，arXiv】以及MoE的负载均衡【索引14，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，Journal of Machine Learning Research】。词汇表大小为64K。分词器使用BPE【索引16，A new algorithm for data compression，1994，The C Users Journal；索引33，Between words and characters: A brief history of open-vocabulary modeling and tokenization in NLP，2021，arXiv；索引44，Neural machine translation of rare words with subword units，2016，Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics】进行训练，每个数字都是一个单独的词元【索引7，Palm: Scaling language modeling with pathways，2023，Journal of Machine Learning Research】。我们还移除了Llama和Mistral分词器中使用的虚拟空格，以实现更一致和可逆的分词。

3. 收获益处

3.1 针对单个80GB GPU的Jamba实现

具体配置选择。在我们的实现中，具体的配置旨在适应单个80GB的GPU，同时在质量和吞吐量方面实现最佳性能。我们的实现中有一系列共4个Jamba块。每个Jamba块具有以下配置。

详细参数配置。
- l = 8: 层的数量。
- a:m = 1:7: 注意力层与Mamba层的比例。
- e = 2: 使用MoE替代单个MLP的频率。
- n = 16: 总专家数。
- K = 2: 每个词元使用的顶尖专家数。

配置选择的依据。选择a:m = 1:7的比例是基于初步的消融实验，如第6节所示，因为在性能最佳的变体中，这个比例是计算效率最高的。

MoE配置的考量。专家的配置是为了使模型能够装入单个80GB的GPU（使用int8权重），同时为输入保留足够的内存。特别地，n和e经过平衡，使得每层平均有约8个专家。此外，我们平衡了n、K和e以实现高质量，同时控制计算需求和通信依赖（内存传输）。因此，我们选择每隔一层用MoE替换MLP模块，并总共设置16个专家，每个词元使用其中两个。这些选择受到了先前关于MoE工作的启发【索引8，Unified scaling laws for routed language models，2022，International conference on machine learning；索引54，ST-MoE: Designing stable and transferable sparse expert models，2022，arXiv】并在初步实验中得到了验证。

长上下文处理能力。图2展示了我们的Jamba实现与Mixtral 8x7B和Llama-2-70B相比，在单个80GB GPU上能容纳的最大上下文长度。Jamba提供的上下文长度是Mixtral的2倍，是Llama-2-70B的7倍。


图 2: 在单个A100 80GB GPU上可容纳的最大上下文长度比较。Jamba实现的上下文长度是Mixtral的2倍，是Llama-2-70B的7倍。

训练与发布模型的上下文长度。总的来说，我们的Jamba实现成功地在高达1M词元的上下文长度上进行了训练。发布的模型支持高达256K词元的长度。

3.2 吞吐量分析

具体设置下的吞吐量结果。为了具体说明，我们展示了两种特定设置下的吞吐量结果。

设置一：不同批次大小下的吞吐量。在第一个设置中，我们使用变化的批次大小、单个A100 80GB GPU、int8量化、8K上下文长度，并生成512个词元的输出。如图3a所示，Jamba允许处理大批量数据，与Mixtral相比，吞吐量（词元/秒）提高了3倍，而Mixtral尽管活跃参数数量相近，却无法容纳16的批次大小。

设置二：不同上下文长度下的吞吐量。在第二个设置中，我们使用单个批次、4个A100 GPU、无量化、变化的上下文长度，并生成512个词元的输出。如图3b所示，在短上下文长度下，所有模型的吞吐量相似。Jamba在长上下文方面表现出色；在128K词元下，其吞吐量是Mixtral的3倍。值得注意的是，尽管Jamba尚未享受到社区过去六年来为纯Transformer模型开发的各种优化，但仍取得了这一成绩。随着针对Jamba的优化被开发出来，我们可以预期吞吐量差距会进一步扩大。


(a) 不同批次大小下的吞吐量（单个A100 GPU，8K上下文长度）。Jamba允许处理大批量数据，吞吐量比Mixtral高3倍。

(b) 不同上下文长度下的吞吐量（单个批次，4个A100 GPU）。在128K词元的上下文中，Jamba的吞吐量是Mixtral的3倍，而Llama-2-70B无法适应如此长的上下文。
图 3: Jamba与近期开放模型在吞吐量（词元/秒）上的比较。

A3 消融实验与见解

本节讨论了我们为Jamba架构实现中的不同设计选择所进行的消融实验。我们首先展示了结合注意力和Mamba层的好处，它们应该以何种比例组合，以及如何交错排列。我们研究了纯Mamba模型失败的情况，这表明它在发展上下文学习能力方面存在困难，而注意力-Mamba混合模型则表现出与纯Transformer相似的上下文学习能力。然后我们展示了在混合注意力-Mamba模型之上添加MoE的好处。最后，我们分享了两个我们发现有用的额外学习点：Jamba中不需要显式的位置信息，以及Mamba层需要特殊的归一化来稳定大规模训练。

消融实验的评估指标。对于这些消融实验，我们报告了以下指标，即使在小数据或小模型规模下，这些指标也表现出有信息量的性能。
- 学术基准测试: HellaSwag (10-shot) 【索引52，HellaSwag: Can a machine really finish your sentence?，2019，Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics】, WinoGrande (5-shot) 【索引42，WinoGrande: An adversarial winograd schema challenge at scale，2020，Proceedings of the AAAI Conference on Artificial Intelligence】, Natural Questions closed-book (NQ; 5-shot) 【索引27，Natural questions: a benchmark for question answering research，2019，Transactions of the Association for Computational Linguistics】。
- HuggingFace OpenLLM排行榜 (OLLM) 【索引12，Open LLM leaderboard，2024，Hugging Face】: 多个数据集的汇总统计数据。我们报告的是我们复现的结果。
- 困惑度评估: 我们报告了在三个领域文本上的对数概率（每字节）：C4、书籍和代码。

6.1 结合注意力与Mamba的益处

注意力与Mamba层的比例探索。我们首先研究了注意力和Mamba层的比例（a:m），使用了1.3B参数的模型，训练了250B个词元。如表4所示，混合的Jamba模型优于纯注意力或纯Mamba模型。注意力与Mamba层的比例可以是1:3或1:7，性能上几乎没有差异。图6显示了这些模型的训练损失，其中Jamba在训练过程中展现出更低的损失。考虑到1:7的比例计算效率更高且性能相似，我们在更大规模的实验中选择了它。


表 4: 学术基准测试和对数概率评估结果显示，与纯注意力和Mamba模型相比，注意力-Mamba混合模型（无MoE）性能更优。注意力与Mamba层的1:3和1:7比例之间没有显著差异。模型均为1.3B参数，训练了250B个词元。

图 6: 纯注意力、纯Mamba以及注意力-Mamba混合模型（无MoE，比例a:m为1:3和1:7）的训练损失曲线。所有模型均为1.3B参数。两种混合模型在整个训练过程中都取得了更好的损失，且不同注意力/Mamba比例之间没有明显差异。

7B规模下的架构比较。接下来，我们比较了在7B模型规模下，训练了50B个词元后的纯Transformer、纯Mamba和注意力-Mamba混合模型的性能。如表5所示，纯Mamba模型相当有竞争力，但略微落后于纯注意力模型。混合的注意力-Mamba模型（无MoE）在性能上优于这两种纯模型，同时获得了比纯Transformer模型更高的吞吐量（见3.2节）。


表 5: 学术基准测试和对数概率评估结果，比较了纯注意力、纯Mamba和注意力-Mamba混合模型（无MoE）。模型均为7B参数，训练了50B个词元。

7B规模下的训练损失。图7显示了这三种架构的训练损失。虽然纯Transformer和纯Mamba模型的收敛情况相似，但混合的Jamba模型（无MoE）在整个训练运行中的损失更低。


图 7: 纯注意力、纯Mamba和注意力-Mamba混合模型（无MoE）的训练损失曲线。所有模型均为7B参数。混合模型在整个训练过程中取得了更好的损失。

6.2 为什么这种组合有效？

纯Mamba模型的局限性。纯Mamba模型在早期的大多数任务中表现相当不错，包括在一般的困惑度评估中。然而，在三个常见的基准任务中，它的表现明显差于纯注意力模型：IMDB【索引32，Learning word vectors for sentiment analysis，2011，Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies】、QuAC【索引6，QuAC: Question answering in context，2018，Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing】和NarrativeQA【索引26，The NarrativeQA reading comprehension challenge，2018，Transactions of the Association for Computational Linguistics】。相比之下，混合的注意力-Mamba模型在这些数据集上的表现与注意力模型相似。表6显示了1.3B模型在训练250B个词元后的结果。


表 6: Mamba在某些数据集上表现不佳，而注意力-Mamba混合模型表现与注意力模型相当。

对Mamba局限性的深入分析。进一步研究这些结果后，我们发现纯Mamba模型经常不遵循正确的格式。例如，在IMDB数据集中，答案选项是“Positive”或“Negative”。虽然注意力模型遵守这种格式，但纯Mamba模型经常产生其他答案，如“Very Good”、“Very Positive”、“Funny”、“Bad”、“Poor”和“3/10”。虽然这些可能被认为是正确的答案，但Mamba难以遵守格式表明了潜在的问题。确实，要成功进行上下文学习，模型捕捉输入-输出格式非常重要【索引34，Rethinking the role of demonstrations: What makes in-context learning work?，2022，Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing】。混合的注意力-Mamba模型像纯注意力模型一样，成功地遵循了格式。

关于SSM上下文学习能力的假说。我们假设这种现象指向了SSM的一个局限性——在上下文学习（ICL）方面可能存在困难。确实，执行ICL的能力与Transformer语言模型在训练期间所谓的“归纳头”（induction heads）的出现有关，这些头执行近似的复制操作，支持ICL【索引35，In-context learning and induction heads，2022，arXiv】。我们推测，纯Mamba模型中缺乏注意力机制使其难以进行上下文学习。虽然当明确训练Mamba进行复制和执行简单的ICL时，它可能学会这些【索引17，Mamba: Linear-time sequence modeling with selective state spaces，2023，arXiv；索引36，Can mamba learn how to learn? a comparative study on in-context learning tasks，2024，arXiv】，但ICL是否是SSM中像Transformer模型那样典型的涌现能力尚不清楚。相比之下，混合的注意力-Mamba模型确实能成功执行ICL，即使8层中只有1层是注意力层。

归纳机制的实例证据。作为涌现的归纳机制的一个实例证据，我们在图8中可视化了一个来自1.3B注意力-Mamba混合模型（无MoE）的注意力头的注意力分布，该示例是纯Mamba失败而混合模型成功的IMDB例子。显然，来自最后一个词元“:”的注意力集中在少样本示例的标签上。我们在我们的混合模型中发现了12个这样的头，分布在所有三个注意力层中（对应模型中的第4、12、20层）。

未来研究方向。未来的工作可以进一步研究ICL在混合模型大规模涌现的情况。我们发布的检查点希望能促进此类研究。最后，最近的工作试图从像Mamba这样的状态空间模型中提取类似注意力的分数【索引1，The hidden attention of mamba models，2024，arXiv】，这为在状态空间模型中寻找归纳能力开辟了另一个方向。

|Bos| _Passage :_To_call_this_film_a _disaster _will _be _an _understatement ._I_don 't _even _know _where _to_begin !_I_have _questions _though，_and _lots_of_them ._I _would […]

_to_mention _a _few_of_course._This_film_was _a _painful _experience _for _me _and _I _advise _everyone _to_skip _it _by_all _means_necessary _and _possible ._Bollywood _should _be _terribly _ashamed _of _this _kind _of_film - making .<|newline|>_Sent iment : Negative <|newline|><|newline|>_Passage :_I_woke _up _and _it_was_a _beautiful_day ;_the _sun _was _shining ，_the _birds _were _singing _and _i_fanc ied _getting _a _movie，_something

[…]

_if_nothing _else ,_ok _it _is_nothing _else.Enjoy ，_but _a _Little _advice - _before _pressing _the _play _button _on _your_DvD _player，_throw_it _out _of_the _window. <|newline|>_Sent iment : Negative <|newline|> <|newline|>_Passage : _This _is _an _excellent _James _Bond _movie . Although _it _is _not _part _of _the _original _and _more _famous _series ,_and _it_is_a _standalone _film,_it _is _very _well_done ._En tic ing _Sean _Conner y […]

_different _version _of _Thunder ball,_updated _with _newer_technology .Regardless _of _the _repeated _theme , _there _are _sufficient _differences _to _make _it _most _entertaining ._I _will _watch _this _one _frequently .<|newline|>_Sent iment : Positive <|newline|> _Passage : _I _am_a_fan _of_Ed_Harris _work_and _I _really _had _high _expectations _about _this_film._Having _so _good _actors_as _Harris _and_Von _Sy d ow_is

[…]

_movies _for_the _VHs_era_of _the _ 90 's ._Whatever_the _reason _was ,_though ,_this _movie _was _a _very _bad _choice _for _anyone _involved .<|newline|>_Sent iment :

图 8: 来自混合注意力-Mamba模型的一个归纳头示例（H3，第一个注意力层）。高亮词语反映了来自最后一个词元“:”的强注意力，此时模型即将预测标签。我们可以看到注意力集中在少样本示例的标签词元上。

6.3 专家混合（MoE）的效果

MoE对混合架构的增益。最近的研究表明，MoE可以在保持计算量可控的情况下改进Transformer语言模型【索引24，Mixtral of experts，2024，arXiv】。然而，MoE是否能与状态空间模型，特别是我们的大规模混合注意力-Mamba架构良好集成，尚不清楚。事实上，表7显示，MoE在大规模（7B参数，训练50B个词元）下提升了混合注意力-Mamba架构的性能。MoE变体具有n=16个总专家，每个词元使用K=2个专家，并且如3.1节所述，每e=2层应用一次MoE。
表 7: 专家混合提升了注意力-Mamba混合模型的性能。

6.4 在大规模下稳定Mamba

大规模训练中的稳定性问题。在训练高达1.3B参数的Jamba模型时，我们观察到训练过程稳定，没有特殊问题。然而，当扩展到本文发布的最大模型（基于7B，拥有12B/52B的活跃/总参数）时，我们遇到了巨大的损失尖峰。调查发现，Mamba层的内部部分存在巨大的激活值，导致了尖峰。因此，我们对内部激活值增加了RMSNorm【索引53，Root mean square layer normalization，2019，Advances in Neural Information Processing Systems】。如图9所示，这稳定了训练并防止了额外的损失尖峰。


图 9: 在Mamba层中添加RMSNorm可防止损失尖峰。

6.5 Jamba不需要显式的位置信息

位置信息消融实验。表8显示了Jamba架构（带MoE）在没有位置信息和在注意力层应用RoPE【索引47，Roformer: Enhanced transformer with rotary position embedding，2024，Neurocomputing】时的结果（1.3B参数模型，250B个词元）。结果相似，表明混合架构可能不需要显式的位置信息。据推测，放置在注意力层之前的Mamba层提供了隐式的位置信息。


表 8: Jamba在有无显式位置信息时的比较。

A4 实验

实验环境

  • 硬件配置: 模型在NVIDIA H100 GPU上进行训练。
  • 软件配置: 使用了一个内部专有框架，该框架支持高效的大规模训练，包括FSDP、张量并行、序列并行和专家并行。
  • 数据集: Jamba在一个内部数据集上进行训练，该数据集包含来自网络、书籍和代码的文本数据，最后更新于2024年3月。数据处理流程包括质量过滤和去重。

实验结果

5.1 学术基准测试

测试任务:
- 常识推理: HellaSwag (10-shot), WinoGrande (5-shot), ARC-E (0-shot), ARC-Challenge (25-shot), PIQA (0-shot)。
- 阅读理解: BoolQ (10-shot), QuAC (0-shot)。
- 其他: GSM8K (3-shot CoT), HumanEval (pass@1), Natural Questions (NQ, 5-shot), TruthfulQA (0-shot)。
- 综合基准: MMLU (5-shot), BBH (3-shot)。

结果与分析:
如表2所示，Jamba与多个公开可用模型进行了比较，包括Llama-2 13B（活跃参数相近）、Llama-2 70B（规模更大）、Gemma 7B和Mixtral（活跃和总参数相近）。

  • 性能相当: 在大多数任务中，Jamba的表现与同等或更大规模的领先模型（包括Llama-2 70B和Mixtral）相当。
  • 效率优势:
    • 与Llama-2 70B相比，Jamba的总参数更少（52B vs 70B）。
    • 作为一个稀疏模型，Jamba只有12B活跃参数，与Mixtral的12.9B相似。然而，由于混合架构，Jamba的KV缓存在256K上下文时仅需4GB，而Mixtral需要32GB。
    • Jamba在实现强劲性能的同时，吞吐量远超Llama-2 70B和Mixtral，最高可达3倍提升。

结论: Jamba证明了混合架构能够在达到同等规模最先进Transformer模型性能的同时，享有SSM带来的效率优势。


表 2: Jamba与其他公开可用模型的比较。Jamba在获得相似性能的同时，具有更好的吞吐量。

5.2 长上下文评估

Needle-in-a-haystack (大海捞针) 测试:
如图4所示，Jamba在该测试中表现出色，能够从长达256K词元的上下文中准确检索到植入的简单陈述。这一结果尤其值得注意，因为Jamba的实现仅使用了4个注意力层。


图 4: “大海捞针”评估显示了Jamba在长达256K词元的上下文中回忆植入语句的能力。

自然长上下文评估:
1. 多样本分类: 在四个标签空间较大的数据集（Trec-Fine, NLU Intent, Banking77, CLINC150）上，通过增加少样本示例数量来测试模型的上下文学习能力，上下文长度最高达128K词元。
- 结果 (图5): 在Trec-Fine和Banking77上，Jamba优于Mixtral，尤其是在少样本示例数量较多时。在NLU Intent和CLINC150上，两者表现相当。


图 5: Jamba与Mixtral在大量少样本示例的分类任务上的比较。结果为贪婪解码下的精确匹配率。

  1. 长上下文问答: 使用来自L-Eval的五个长上下文数据集（NarrativeQA, LongFQA, NQ, CUAD, SFiction），采用3-shot格式进行评估。这些数据集的平均输入长度从6K到62K词元不等。
    • 结果 (表3): 在F1分数上，Jamba在大多数数据集以及平均分上都优于Mixtral。
    • 效率: 在这些计算密集型的长上下文任务中，Jamba的高吞吐量优势尤为突出。

      表 3: 在长上下文QA基准测试上的结果（F1），采用3-shot格式。

A5 结论

我们提出了Jamba，一种结合了注意力层和Mamba层，并带有MoE模块的新型架构，并开源了其实现。该模型达到了最先进的性能，并支持长上下文。我们展示了Jamba如何在性能和内存需求之间提供灵活性，同时保持高吞吐量。我们通过实验探讨了多种设计选择，如注意力与Mamba层的比例，并讨论了开发过程中的一些发现，这些发现将为未来混合注意力-状态空间模型的研究提供信息。为了促进此类研究，我们计划发布小规模训练运行的模型检查点。本次发布的最大的模型拥有12B活跃参数和52B总可用参数，支持高达256K词元的上下文长度，并且在处理140K词元文本时也能装入单个80GB的GPU中。