AI21 Labs releases Jamba, a hybrid Mamba-Transformer model

Background

Many language models of this period were built on the Transformer, an architecture whose attention mechanism lets every part of the input relate to every other part. Attention works well, but its cost rises quickly as inputs grow longer. Mamba was a newer alternative that processes a sequence step by step, keeping a compact summary of what it has read, so that its cost grows more gently with length. AI21 Labs, which had earlier released the Jurassic models, released Jamba as a model that combines the two designs.

What happened

On 28 March 2024 AI21 Labs released Jamba, which combines Mamba state-space layers with Transformer attention layers ^[1]. The company billed it as the first production-scale Mamba-based model ^[2]. The weights were released under the Apache 2.0 licence on Hugging Face ^[1].

Key details

Jamba has 52 billion total parameters, with 12 billion active parameters for each token, through its use of mixture-of-experts layers ^[1]. A mixture-of-experts layer routes each token to a small number of specialised sub-networks, so the model stores many more parameters than it uses for any single token.

The model supports a 256k-token context window. Up to 140k tokens fit on a single 80 GB GPU ^[1]. AI21 claimed 3x the throughput of Mixtral 8x7B on long contexts ^[1]. These are AI21's own figures.

In the layout that AI21 described, the model uses one Transformer attention layer for every eight layers, and the rest are Mamba layers, each paired with an MLP. The attention layers give the model a way to relate distant parts of the text directly, while the Mamba layers do most of the sequential work at lower cost.

Impact

The release introduced a hybrid architecture as an openly available model that others could download and test. Its significance lies in that combination: a production-scale model that mixes two families of layers rather than relying on attention alone. The model was followed in August 2024 by the instruction-tuned Jamba 1.5 Mini and Large models ^[2]. This entry covers the initial release; the later versions have their own entries.