Paper summary: Jamba: A Hybrid Transformer-Mamba Language Model

Problem

Transformers are strong but expensive for long contexts, with a big key-value cache and quadratic attention. Mamba (the selective state-space model from December 2023) is efficient, but pure SSMs had their own weaknesses. Can a hybrid get the benefits of both?

Method

Jamba interleaves blocks of Transformer layers and Mamba layers, and adds mixture-of-experts in some layers to increase capacity while keeping the number of active parameters manageable. The design is flexible. The released configuration is chosen so that the model fits on a single 80GB GPU. The paper also studies how to combine the two layer types and how to mix experts.

Key results

From the abstract: compared with vanilla Transformers, Jamba has high throughput and a small memory footprint, while reaching state-of-the-art performance on standard benchmarks and long-context evaluations. It shows strong results for contexts up to 256K tokens.

Limitations

There’s no dedicated limitations section in the extracted text. The paper describes one particular configuration, and the authors note that several architectural choices proved crucial at scale, so the design space isn’t fully settled.