Jamba

What it is

Jamba is an open-weights language model released by AI21 Labs on 28 March 2024. Its name stands for Joint Attention and Mamba ^[1][2][4]. It is a hybrid SSM-Transformer model, meaning it mixes two kinds of layers: Mamba state-space layers and Transformer attention layers ^[1]. It was described as the first production-scale implementation of the Mamba architecture ^[2].

How it works

Transformer models use attention, which lets each part of the input look at the other parts when producing output. Attention works well but its cost grows quickly as the input gets longer. Mamba is a state-space design that processes a sequence step by step while keeping a compact running summary of what it has seen, so its cost grows more gently with length.

Jamba uses both. Of every eight layers, one is a Transformer attention layer and the rest are Mamba layers, and each layer is paired with a multilayer perceptron (MLP) ^[1]. A multilayer perceptron is a basic neural-network block that transforms each position's values. The mix lets the model keep the attention layers for tasks that need them while the Mamba layers handle most of the sequence.

Key details

Jamba has 52 billion total parameters, of which 12 billion are active for each token, because it uses mixture-of-experts layers ^[1][2]. In a mixture-of-experts model, a router sends each token to a few specialised sub-networks, so only a part of the total parameters does the work for any one token.

It supports a 256k-token context window. Up to 140k tokens fit on a single 80 GB GPU, which AI21 said was unique in its size class ^[1]. AI21 also claimed 3x the throughput of Mixtral 8x7B on long contexts ^[1]. Throughput is the amount of text a system can process in a given time, and these figures are AI21's own claims.

In its base-model benchmarks, Jamba scored 67.4% on MMLU, 87.1% on HellaSwag, 64.4% on the ARC Challenge, 82.5% on WinoGrande and 59.9% on GSM8K with chain-of-thought prompting ^[2].

The model was released under the Apache 2.0 licence on Hugging Face. It was a pretrained base model, with no instruction tuning or safety moderation, and it was listed in the NVIDIA API catalog ^[1][2].

The design is described in the paper "Jamba: A Hybrid Transformer-Mamba Language Model", posted to arXiv in March 2024 ^[3].

Place in development

Jamba succeeded Jurassic-2 as AI21's model family ^[2]. It was followed in August 2024 by the instruction-tuned Jamba 1.5 Mini and Large models ^[2].