Jamba

AI21 Labs · 2024-03-28 · analysed 2026-08-08 · skill v0.3.1

hybrid hybrid text-only open-weights large frontier-generalist

Design overview

Jamba interleaves Transformer and Mamba layers into repeated 'Jamba blocks' — the report's term for l=8 layers mixed at a 1:7 attention-to-Mamba ratio — and adds mixture-of-experts to some MLPs. The released configuration stacks 4 blocks: 32 layers total, of which only 4 are GQA attention layers (positions 4, 12, 20, 28) and the rest are Mamba state-space layers; every other layer's MLP is replaced by an MoE with 16 experts and top-2 routing, giving 52B total but only 12B active parameters. The design deliberately trades attention for Mamba to shrink the KV cache (reported 8x smaller than a vanilla Transformer; 4GB at 256K context) and raise long-context throughput (3x Mixtral at 128K), while keeping enough attention for in-context learning — the report shows pure-Mamba models fail to follow few-shot formats, and the hybrid restores Transformer-like behaviour with as little as one attention layer in eight. Notable choices: no explicit positional encoding (Mamba layers are said to provide implicit position), RMSNorm inside Mamba layers for training stability, SwiGLU activations, a 64K-token BPE tokenizer with per-digit tokens, and GQA. Trained on contexts up to 1M tokens on H100s with FSDP/tensor/sequence/expert parallelism, the released base model serves 256K-token contexts from a single 80GB GPU with int8 weights and reaches parity with Mixtral-8x7B and Llama-2 70B on standard benchmarks. The report covers a single configuration and discloses few per-dimension numbers; most ablation evidence comes from 1.3B and 7B runs.

What the report claims is novel

Atlas assessment: Every ingredient is pre-existing — the report itself cites H3, Hyena, StripedHyena, Block-State Transformer, and MoE-Mamba as earlier Attention-SSM hybrids — so Jamba's novelty is scale and integration: the first production-grade hybrid at 52B total / 12B active parameters, a 1:7 attention-to-Mamba recipe with MoE on every other layer, and open weights. The empirical findings (no positional encoding needed; RMSNorm on Mamba internals to prevent loss spikes) are genuine, reusable contributions, and the 256K-token context with a 4GB KV cache stood out against contemporaneous open models. Relative to atlas peers, Jamba is the anchor for hybrid sequence mixing (taxonomy Axis 2).

Relation to the atlas

Lineage in

Influence out

Notable omissions

Closest relatives in the atlas

Spec sheet

Identity

FamilyJamba[1]
OrganisationAI21 Labs
Release date2024-03-28
Report URLhttps://arxiv.org/abs/2403.19887[2]
LicenseApache 2.0[3]
Open weightsyes[4]

Scale

Total parameters52B[5]
Active parameters / token12B[5]
Layers32
Hidden dimensionnot disclosed
FFN inner dimensionnot disclosed
FFN rationot disclosed
Query headsnot disclosed
KV headsnot disclosed
Head dimensionnot disclosed
Vocabulary size65536
Embedding tyingnot disclosed

Core block

Block typehybrid
Attention variantGQA[6]
Attention layer pattern4 blocks of 8 layers each (l=8); a:m = 1:7 — one GQA attention layer per block, seven Mamba layers; attention layers at positions 4, 12, 20, 28 (32 layers total, 4 attention layers)[7]
Depth mixingsequential-residual
ActivationSwiGLU[6]
Expert count16[8]
Experts per token2[9]
Shared experts0
Routingtop-K routing[10]
Load balancingload balancing (mechanism not specified)[6]
Expert granularitynot disclosed
MethodNoPE[6]
RoPE base frequencyn/a[12]
Partial RoPEn/a[12]
TypeRMSNorm[13]
Placementnot disclosed
QK-normno
Attention sinksno
Softcappingno
OtherRMSNorm applied to internal activations of Mamba layers to prevent loss spikes when scaling to the 7B-based model (§6.4)[14]

Context

Trained context length1M[15]
Deployed context length256K[15]
Extension methodnone

Tokenizer

AlgorithmBPE[6]
NotesEach digit is a separate token; the dummy space used in the Llama and Mistral tokenizers is removed for more consistent and reversible tokenization. Vocabulary size 64K (recorded in scale.vocab_size).[6]

Training

Training tokensnot disclosed
Data compositionIn-house dataset containing text from the Web, books, and code, last updated March 2024; processing pipeline includes quality filters and deduplication; mixture proportions not disclosed.[16]
Curriculum / stagingnone disclosed
Optimizernot disclosed
LR schedulenot disclosed
Batch schedulenot disclosed
Precisionnot disclosed
ParallelismIn-house proprietary training framework with FSDP, tensor parallelism, sequence parallelism, and expert parallelism[16]
HardwareNVIDIA H100 GPUs[16]
Disclosed computenot disclosed

Post-training

SFTno[3]
Preference optimisationnone disclosed[3]
Reasoning trainingnone disclosed
Distillationnone disclosed

Modality

Typetext-only
Attachmentn/a[17]

Inference efficiency

KV-cache designGQA attention on only 4 of 32 layers: report states an 8x smaller KV cache than a vanilla Transformer and 4GB at 256K context (16-bit) vs 32GB for Mixtral and 128GB for Llama-2 (Table 1)[18]
Quantisation shippedINT8 (inference): model fits in a single 80GB GPU with int8/8bit weights, including at contexts over 128K tokens[19]
Speculative decodingnone disclosed
Serving optimisationsnone disclosed

Evaluation

Benchmarks (report's own numbers only)

BenchmarkValueRef
HellaSwag (10-shot)87.1Table 2
WinoGrande (5-shot)82.5Table 2
ARC-E (0-shot)73.5Table 2
ARC-C (25-shot)64.4Table 2
PIQA (zero-shot)83.2Table 2
NQ (5-shot)45.9Table 2
TruthfulQA (zero-shot)46.4Table 2
BoolQ (10-shot)88.2Table 2
QuAC (zero-shot)40.9Table 2
GSM8K (3-shot CoT)59.9Table 2
HumanEval (pass@1)29.3Table 2
MMLU (5-shot)67.4Table 2
BBH (3-shot)45.4Table 2
NQ (long-context, 3-shot F1)0.6Table 3
Avg (L-Eval long-context F1)0.44Table 3

Provenance references

  1. stated · Title / §1
  2. stated · pipeline
  3. stated · §1
  4. stated · Abstract / §1
  5. stated · Abstract / §1 / Table 1
  6. stated · §2
  7. stated · §3.1 / Figure 1 / §5.2.1
  8. stated · §3.1 (n=16) / §6.3
  9. stated · §3.1 (K=2) / §6.3
  10. stated · §2 / §3.1
  11. n/a · §3.1
  12. n/a · §2
  13. stated · §2 / §6.4
  14. stated · §6.4
  15. stated · §3.1 / §5.2
  16. stated · §4
  17. n/a · §4
  18. stated · §2 / Table 1
  19. stated · §1 / §3.1 / §3.2