Jamba: AI21 Labs’ New Hybrid Transformer-Mamba Language Model

Language models has witnessed rapid advancements, with Transformer-based architectures leading the charge in natural language processing. However, as models scale, the challenges of handling long contexts, memory efficiency, and throughput have become more pronounced.

AI21 Labs has introduced a new solution with Jamba, a state-of-the-art large language model (LLM) that combines the strengths of both Transformer and Mamba architectures in a hybrid framework. This article delves into the details of Jamba, exploring its architecture, performance, and potential applications.

Overview of Jamba

Jamba is a hybrid large language model developed by AI21 Labs, leveraging a combination of Transformer layers and Mamba layers, integrated with a Mixture-of-Experts (MoE) module. This architecture allows Jamba to balance memory usage, throughput, and performance, making it a powerful tool for a wide range of NLP tasks. The model is designed to fit within a single 80GB GPU, offering high throughput and a small memory footprint while maintaining state-of-the-art performance on various benchmarks.

The Architecture of Jamba

Jamba’s architecture is the cornerstone of its capabilities. It is built on a novel hybrid design that interleaves Transformer layers with Mamba layers, incorporating MoE modules to enhance the model’s capacity without significantly increasing computational demands.

1. Transformer Layers

The Transformer architecture has become the standard for modern LLMs due to its ability to handle parallel processing efficiently and capture long-range dependencies in text. However, its performance is often limited by high memory and compute requirements, particularly when processing long contexts. Jamba addresses these limitations by integrating Mamba layers, which we will explore next.

2. Mamba Layers

Mamba is a recent state-space model (SSM) designed to handle long-distance relationships in sequences more efficiently than traditional RNNs or even Transformers. Mamba layers are particularly effective at reducing the memory footprint associated with storing key-value (KV) caches in Transformers. By interleaving Mamba layers with Transformer layers, Jamba reduces the overall memory usage while maintaining high performance, especially in tasks requiring long context handling.

3. Mixture-of-Experts (MoE) Modules

The MoE module in Jamba introduces a flexible approach to scaling model capacity. MoE allows the model to increase the number of available parameters without proportionally increasing the active parameters during inference. In Jamba, MoE is applied to some of the MLP layers, with the router mechanism selecting the top experts to activate for each token. This selective activation enables Jamba to maintain high efficiency while handling complex tasks.

The below image demonstrates the functionality of an induction head in a hybrid Attention-Mamba model, a key feature of Jamba. In this example, the attention head is responsible for predicting labels such as “Positive” or “Negative” in response to sentiment analysis tasks. The highlighted words illustrate how the model’s attention is strongly focused on label tokens from the few-shot examples, particularly at the critical moment before predicting the final label. This attention mechanism plays a crucial role in the model’s ability to perform in-context learning, where the model must infer the appropriate label based on the given context and few-shot examples.

The performance improvements offered by integrating Mixture-of-Experts (MoE) with the Attention-Mamba hybrid architecture are highlighted in Table. By using MoE, Jamba increases its capacity without proportionally increasing computational costs. This is particularly evident in the significant boost in performance across various benchmarks such as HellaSwag, WinoGrande, and Natural Questions (NQ). The model with MoE not only achieves higher accuracy (e.g., 66.0% on WinoGrande compared to 62.5% without MoE) but also demonstrates improved log-probabilities across different domains (e.g., -0.534 on C4).

Key Architectural Features

• Layer Composition: Jamba’s architecture consists of blocks that combine Mamba and Transformer layers in a specific ratio (e.g., 1:7, meaning one Transformer layer for every seven Mamba layers). This ratio is tuned for optimal performance and efficiency.
• MoE Integration: The MoE layers are applied every few layers, with 16 experts available and the top-2 experts activated per token. This configuration allows Jamba to scale effectively while managing the trade-offs between memory usage and computational efficiency.
• Normalization and Stability: To ensure stability during training, Jamba incorporates RMSNorm in the Mamba layers, which helps mitigate issues like large activation spikes that can occur at scale.

Jamba’s Performance and Benchmarking

Jamba has been rigorously tested against a wide range of benchmarks, demonstrating competitive performance across the board. The following sections highlight some of the key benchmarks where Jamba has excelled, showcasing its strengths in both general NLP tasks and long-context scenarios.

1. Common NLP Benchmarks

Jamba has been evaluated on several academic benchmarks, including:

• HellaSwag (10-shot): A common sense reasoning task where Jamba achieved a performance score of 87.1%, surpassing many competing models.
• WinoGrande (5-shot): Another reasoning task where Jamba scored 82.5%, again showcasing its ability to handle complex linguistic reasoning.
• ARC-Challenge (25-shot): Jamba demonstrated strong performance with a score of 64.4%, reflecting its ability to manage challenging multiple-choice questions.

In aggregate benchmarks like MMLU (5-shot), Jamba achieved a score of 67.4%, indicating its robustness across diverse tasks.

2. Long-Context Evaluations

One of Jamba’s standout features is its ability to handle extremely long contexts. The model supports a context length of up to 256K tokens, the longest among publicly available models. This capability was tested using the Needle-in-a-Haystack benchmark, where Jamba showed exceptional retrieval accuracy across varying context lengths, including up to 256K tokens.

3. Throughput and Efficiency

Jamba’s hybrid architecture significantly improves throughput, particularly with long sequences.

In tests comparing throughput (tokens per second) across different models, Jamba consistently outperformed its peers, especially in scenarios involving large batch sizes and long contexts. For instance, with a context of 128K tokens, Jamba achieved 3x the throughput of Mixtral, a comparable model.

Using Jamba: Python

For developers and researchers eager to experiment with Jamba, AI21 Labs has provided the model on platforms like Hugging Face, making it accessible for a wide range of applications. The following code snippet demonstrates how to load and generate text using Jamba:

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]
outputs = model.generate(input_ids, max_new_tokens=216)
print(tokenizer.batch_decode(outputs))

import torch
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
model = AutoModelForCausalLM.from_pretrained(
"ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)
lora_config = LoraConfig(r=8,
target_modules=[
"embed_tokens","x_proj", "in_proj", "out_proj", # mamba
"gate_proj", "up_proj", "down_proj", # mlp
"q_proj", "k_proj", "v_proj" 
# attention],
task_type="CAUSAL_LM", bias="none")
dataset = load_dataset("Abirate/english_quotes", split="train")
training_args = SFTConfig(output_dir="./results",
num_train_epochs=2,
per_device_train_batch_size=4,
logging_dir='./logs',
logging_steps=10, learning_rate=1e-5, dataset_text_field="quote")
trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args,
peft_config=lora_config, train_dataset=dataset,
)
trainer.train()