[Paper Review] Jamba: A Hybrid Transformer-Mamba Language Model

Opher Lieber, Barak Lenz+20 more|arXiv (Cornell University)|Mar 28, 2024

Language, Discourse, Communication Strategies40 citations

TL;DR

Jamba introduces a hybrid Transformer-Mamba mixture-of-experts architecture that interleaves Transformer and Mamba layers with MoE to achieve high performance and long-context capability while fitting in a single 80GB GPU.

ABSTRACT

We present Jamba, a new base large language model based on a novel hybrid Transformer-Mamba mixture-of-experts (MoE) architecture. Specifically, Jamba interleaves blocks of Transformer and Mamba layers, enjoying the benefits of both model families. MoE is added in some of these layers to increase model capacity while keeping active parameter usage manageable. This flexible architecture allows resource- and objective-specific configurations. In the particular configuration we have implemented, we end up with a powerful model that fits in a single 80GB GPU. Built at large scale, Jamba provides high throughput and small memory footprint compared to vanilla Transformers, and at the same time state-of-the-art performance on standard language model benchmarks and long-context evaluations. Remarkably, the model presents strong results for up to 256K tokens context length. We study various architectural decisions, such as how to combine Transformer and Mamba layers, and how to mix experts, and show that some of them are crucial in large scale modeling. We also describe several interesting properties of these architectures which the training and evaluation of Jamba have revealed, and plan to release checkpoints from various ablation runs, to encourage further exploration of this novel architecture. We make the weights of our implementation of Jamba publicly available under a permissive license.

Motivation & Objective

• Investigate whether interleaving Transformer and Mamba layers can combine the strengths of attention-based and state-space models.
• Assess how MoE integration affects capacity, throughput, and memory in a hybrid architecture.
• Evaluate performance on standard benchmarks and long-context tasks, including 256K-token contexts.
• Demonstrate training stability and practicality of deploying a 7B-12B parameter-scale hybrid model on commodity hardware.

Proposed method

• Define a Jamba block that combines Transformer or Mamba layers followed by an MLP or MoE module.
• Interleave blocks with a:m attention-to-Mamba ratio and apply MoE every e layers with n total experts and top-K routing per token.
• Use RMSNorm in Mamba layers and omit explicit positional embeddings, relying on the hybrid structure for implicit position information.
• Train on large-scale data with 64K vocabulary and BPE tokenizer, optimizing for throughput and memory efficiency on 80GB GPU setups.
• Evaluate on academic benchmarks, long-context QA datasets, and throughput measurements across different hardware scales.

Experimental results

Research questions

• RQ1 Can a hybrid Attention-Mamba architecture match or surpass pure Transformer models of similar size on standard benchmarks?
• RQ2 Does incorporating MoE into the hybrid architecture improve capacity without prohibitive compute?
• RQ3 How does the ratio of Attention to Mamba layers affect memory usage, throughput, and long-context performance?
• RQ4 Can Jamba effectively handle very long contexts (up to 256K tokens) with reasonable KV cache requirements?
• RQ5 What are the practical training stability considerations for large-scale hybrid models?

Key findings

• Jamba achieves competitive or superior accuracy on standard benchmarks compared to similar-size open models like Mixtral and Llama-2 70B.
• The hybrid Attention-Mamba architecture reduces KV cache requirements to 4GB at 256K context, enabling long-context processing on a single 80GB GPU.
• MoE variants improve performance over non-MoE hybrids at large scale (7B parameters trained on 50B tokens).
• An Attention-Mamba hybrid outperforms pure Mamba on several tasks and supports in-context learning similar to Transformers, suggesting complementary strengths.
• Explicit positional information is not required for Jamba, as the Mamba-first structure provides implicit positional information.