Jamba: A Hybrid Transformer-Mamba Language Model
By Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, S. Meirom, Yonatan Belinkov, Shai Shalev-Shwartz and 12 more. arXiv.org, 2024
467 citations recorded by the source corpus.
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.
Subfields: Language and Linguistics
View Paper: https://www.semanticscholar.org/paper/cbaf689fd9ea9bc939510019d90535d6249b3367
How researchers see this paper
The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.
- An Yang
- Jinze Bai
- Yunfei Chu
- Shuai Bai
- Baosong Yang
- Junyang Lin
- Rui Men
- Jin Xu
Citations in the corpus
- Parallelizing Linear Transformers with the Delta Rule over Sequence Length (2024, unknown): “Following recent work on combining subquadratic token-mixing layers with existing neural network primitives [6, 20, 53], we also experiment with hybridizing DeltaNet models.”
- OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework (2024, methodology): “OpenRLHF seamlessly integrates with Hugging Face Transformers [29] and supports popular technologies such as Mixture of Experts (MoE) [13], Jamba [17] and QLoRA [4].”
- Qwen2 Technical Report (2024, methodology): “These baselines include other MoE models, such as Mixtral-8x7B (Jiang et al., 2024) and Jamba (Lieber et al., 2024), and dense models, such as Yi-1.”
- A Survey on Mixture of Experts in Large Language Models (2024, unknown): “Switch Transformers [34] and many other subsequent studies [26], [28], [33], [66] have embraced this L aux design, and enhancements [30], [65], [67] have been made to cater to diverse requirements.”
- MambaOut: Do We Really Need Mamba for Vision?* (2024, unknown): “Notably, models like RWKV [65] and Mamba [28] are proven to be effective as the backbone for large language models (LLMs) [52, 65].”