What is Jamba
Jamba is the first production-grade large language model based on the Mamba architecture launched by AI21 Labs. Currently, most large models (such as GPT, Gemini, and Llama) are based on the Transformer structure. Jamba combines the Mamba structured state space model (SSM) and the traditional Transformer architecture, offering the advantages of high-quality output, high throughput, and low memory usage. This hybrid architecture model provides a 256K context window, significantly improving throughput and efficiency when processing long texts.
Jamba is released in the form of open weights under the Apache 2.0 open-source license, encouraging the community to conduct further research and optimization. Currently, Jamba is released as a research model, without the necessary guarantees for commercial use. AI21 Labs plans to release a fine-tuned, safer version in the coming weeks.
Jamba's official website entry
- Official project homepage:https://www.ai21.com/jamba
- Official blog introduction:https://www.ai21.com/blog/announcing-jamba
- Hugging Face address:https://huggingface.co/ai21labs/Jamba-v0.1
Jamba's main features
- SSM-Transformer hybrid architecture: Jamba is the first production-grade model that combines the Mamba SSM (structured state space model) with the Transformer architecture. This innovative hybrid architecture is designed to improve the model's performance and efficiency.
- Large-capacity context window: Jamba provides a 256K context window, which enables the model to process longer text sequences and is suitable for more complex natural language processing tasks.
- High throughput: Compared with the Mixtral 8x7B model of the same scale, Jamba achieves a 3x throughput improvement when processing long contexts, allowing it to process large amounts of data more efficiently.
- Large-capacity processing on a single GPU: Jamba can process up to 140K context on a single GPU, significantly improving the model's accessibility and deployment flexibility.
- Open weights license: Jamba's weights are released under the Apache 2.0 license, providing researchers and developers with the freedom to use, modify, and optimize the model, promoting technology sharing and innovation.
- NVIDIA API integration: Jamba will be available as an NVIDIA NIM inference microservice in the NVIDIA API catalog, enabling enterprise developers to easily deploy the Jamba model using the NVIDIA AI Enterprise software platform.
- Optimized MoE layers: Jamba utilizes MoE (mixture of experts) layers in the hybrid architecture, activating only part of the parameters during inference, improving the model's operating efficiency and performance.
Jamba's technical architecture
The Jamba architecture adopts a block and layer approach, enabling Jamba to successfully integrate both the Mamba SSM and Transformer architectures. Each Jamba block contains either an attention layer or a Mamba layer, followed by a multilayer perceptron (MLP), resulting in an overall ratio of one Transformer layer in every eight layers.
In addition, Jamba uses MoE to increase the total number of model parameters while simplifying the number of active parameters used in inference, thereby achieving higher model capacity without a corresponding increase in computational requirements. To maximize model quality and throughput on a single 80GB GPU, AI21 Labs optimized the number of MoE layers and experts used, leaving enough available memory for common inference workloads.
Jamba's performance comparison
According to a report by AI21 Labs, the Jamba model demonstrates excellent results on various benchmarks (such as HellaSwag, ArcChallenge, MLLU, etc.), and in a wide range of task tests (such as language understanding, scientific reasoning, common-sense reasoning, etc.) it is comparable to or even surpasses the most advanced models in the same size category (such as Llama2 13B, Llama2 70B, Gemma 7B, Mixtral 8×7B).
