今天小编分享的科学经验：Mamba架构第一次做大！混合Transformer，打败Transformer，欢迎阅读。
精彩精彩，第一个把爆火 Mamba 架构真正扩展到足够大的工作来了。
520 亿参数，还是Mamba+Transformer混合架构。
它的名字叫Jamba。
取两种架构之长，模型质量和效率兼得，要吞吐量有吞吐量，要低内存有低内存。
初步跑分显示：
Jamba 性能总体接近Mixtral 8x-7B，处理 128k 长上下文时吞吐量却是其 3 倍。
一共支持 256k 上下文，而单张 A100 GPU 即可处理 140k，直接拿下同等规模模型之最，高效又经济。
这项成果，来自以色列 AI 公司AI21labs。
Mamba 原作者看了之后都激动转发：
绝对的 " 大新闻 "。
Mamba、Transformer，合体
由 CMU 和普林斯顿大学提出的 Mamba，解决了 Transformer 的局限性（随着推理上下文越长，模型内存占用量越大，同时推理速度变慢，由此导致算力消耗巨大）。
但它也有自己的缺点——
在不关注整个上下文的情况下，Mamba 的输出质量很差，尤其是在召回相关的任务上。
本着 " 既要也要 " 的原则，Jamba 站出来提供两全其美之作。
Jamba 由 Transformer、Mamba 和 MoE 层组成，可同时优化内存、吞吐量和性能。
如下图所示，为了集成两种架构，Jamba 采用块层（blocks-and-layers）组合的创新方法。
简单来说，就是每个 Jamba 块包含一个注意力层或一个 Mamba 层，再跟一个多层感知器 MLP，总体比例保证为每八层一个 Transformer 层。
其次，Jamba 利用 MoE 来增加模型参数的总量，同时简化推理中使用的活动参数量。
最终模型容量高了，计算需求也没有相应的增加。
而为了在单张 GPU（80GB）上最大限度地提高模型吞吐量，Jamba 还优化了所用 MoE 层和专家数量，最终为日常推理工作负载留出足够内存。
值得一提的是，在推理时，Jamba 的 MoE 层仅需 520 亿可用参数中的 120 亿，就能同时保证比同等大小的仅 Transformer 模型更高效。
要知道，此前有人光是尝试过扩展 Mamba，就没能做到 30 亿参数之上。
因此，除了成功合体 Mamba 和 Transformer，Jamba 也达成了第二大成就：
同类中第一个达到生产级规模和质量的混合架构（SSM 混 Transformer）（ps. Mamba 就是一种状态空间模型 SSM）。
吞吐量和效率 up
初步评估显示，Jamba 在吞吐量和效率等关键指标上表现出色。
首先，Jamba 可以在长上下文中提供 3 倍吞吐量，比 Mixtral 8x7B 等大小相当的 Transformer 模型都要高效。
如下图所示，当上下文視窗达到 128k 时，Jamba 的每秒 token 数近乎 1500，而此时表现最好的 Mixtral 8x7B 应该才在 500 往上的样子。
其次，在单张 GPU 上，Jamba 最多可以容纳 140k 上下文，经济又高效。
相比之下，Mixtral 8x7B 为 64k，Llama2 70B 则仅为 16k。
第三，Jamba 的输出质量也得到了保证。
在如下一系列推理基准上，4 项中有 3 项它都拿下了 SOTA。同时，在 GSM8K 等基准上，Jamba 即使没有夺魁，也和 SOTA 模型打了个不相上下。
总体来说，Jamba 的性能接近 Mixtral 8x7B。
最后，作者提示，别忘了，这些都还只是初步改造后的结果，后续还有很多优化空间（比如 MoE 并行、更快的 Mamba 实现）。所以到时性能会更强。
好消息：Jamba 现在已经上线 Hugging Face，并且划重点：采用 apache-2.0 许可。
（Jamba 的指令版本则将很快通过 AI21labs 平台上线。）
网友看完都感动哭了。
传送门：
https://huggingface.co/ai21labs/Jamba-v0.1
参考链接：
[ 1 ] https://www.ai21.com/blog/announcing-jamba
[ 2 ] https://www.ai21.com/jamba
[ 3 ] https://twitter.com/AI21Labs/status/1773350888427438424?s=20
[ 4 ] https://twitter.com/tri_dao/status/1773418926518734957?s=20
