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

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

AI21 Labs has introduced a brand new solution with Jamba, a state-of-the-art large language model (LLM) that mixes the strengths of each Transformer and Mamba architectures in a hybrid framework. This text delves into the small print 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 mixture 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 strong tool for a wide selection of NLP tasks. The model is designed to suit inside 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’s built on a novel hybrid design that interleaves Transformer layers with Mamba layers, incorporating MoE modules to boost the model’s capability without significantly increasing computational demands.

1. Transformer Layers

The Transformer architecture has turn out to be the usual for contemporary LLMs resulting from its ability to handle parallel processing efficiently and capture long-range dependencies in text. Nonetheless, its performance is commonly limited by high memory and compute requirements, particularly when processing long contexts. Jamba addresses these limitations by integrating Mamba layers, which we’ll 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 and even Transformers. Mamba layers are particularly effective at reducing the memory footprint related to storing key-value (KV) caches in Transformers. By interleaving Mamba layers with Transformer layers, Jamba reduces the general 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 versatile approach to scaling model capability. MoE allows the model to extend the number of obtainable parameters without proportionally increasing the lively parameters during inference. In Jamba, MoE is applied to a few of the MLP layers, with the router mechanism choosing the highest experts to activate for every token. This selective activation enables Jamba to keep up 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 instance, the eye head is liable for predicting labels equivalent to “Positive” or “Negative” in response to sentiment evaluation tasks. The highlighted words illustrate how the model’s attention is strongly focused on label tokens from the few-shot examples, particularly on the critical moment before predicting the ultimate label. This attention mechanism plays an important role within the model’s ability to perform in-context learning, where the model must infer the suitable 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. Through the use of MoE, Jamba increases its capability without proportionally increasing computational costs. This is especially evident in the numerous boost in performance across various benchmarks equivalent to HellaSwag, WinoGrande, and Natural Questions (NQ). The model with MoE not only achieves higher accuracy (e.g., 66.0% on WinoGrande in comparison with 62.5% without MoE) but additionally 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 mix Mamba and Transformer layers in a particular ratio (e.g., 1:7, meaning one Transformer layer for each 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 make sure stability during training, Jamba incorporates RMSNorm within the Mamba layers, which helps mitigate issues like large activation spikes that may occur at scale.

Jamba’s Performance and Benchmarking

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

1. Common NLP Benchmarks

Jamba has been evaluated on several academic benchmarks, including:

• HellaSwag (10-shot): A standard sense reasoning task where Jamba achieved a performance rating of 87.1%, surpassing many competing models.
• WinoGrande (5-shot): One other 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 rating of 64.4%, reflecting its ability to administer difficult multiple-choice questions.

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

2. Long-Context Evaluations

Considered one of Jamba’s standout features is its ability to handle extremely long contexts. The model supports a context length of as much as 256K tokens, the longest amongst publicly available models. This capability was tested using the Needle-in-a-Haystack benchmark, where Jamba showed exceptional retrieval accuracy across various context lengths, including as much as 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 example, with a context of 128K tokens, Jamba achieved 3x the throughput of Mixtral, a comparable model.

Using Jamba: Python

For developers and researchers desirous to experiment with Jamba, AI21 Labs has provided the model on platforms like Hugging Face, making it accessible for a wide selection of applications. The next code snippet demonstrates how you can 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("Within 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()