Jamba: Ein Hybrid-Transformer-Mamba-Sprachmodell

Jamba: A Hybrid Transformer-Mamba Language Model

March 28, 2024

Autoren: Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, Yoav Shoham

cs.AI

Zusammenfassung

Wir präsentieren Jamba, ein neues Basis-Großsprachmodell, das auf einer neuartigen Hybrid-Transformer-Mamba-Mischung-von-Experten (MoE)-Architektur basiert. Speziell verwebt Jamba Blöcke von Transformer- und Mamba-Schichten, um die Vorteile beider Modellfamilien zu nutzen. MoE wird in einigen dieser Schichten hinzugefügt, um die Modellkapazität zu erhöhen, während die aktive Parameterverwendung überschaubar bleibt. Diese flexible Architektur ermöglicht ressourcen- und zielabhängige Konfigurationen. In der spezifischen Konfiguration, die wir implementiert haben, erhalten wir ein leistungsstarkes Modell, das auf einer einzigen 80-GB-GPU Platz findet. Im großen Maßstab gebaut, bietet Jamba im Vergleich zu reinen Transformern eine hohe Durchsatzleistung und einen geringen Speicherbedarf sowie gleichzeitig eine erstklassige Leistung bei Standard-Sprachmodell-Benchmarks und Langkontext-Evaluationen. Bemerkenswerterweise liefert das Modell starke Ergebnisse für Kontextlängen von bis zu 256.000 Tokens. Wir untersuchen verschiedene architektonische Entscheidungen, wie die Kombination von Transformer- und Mamba-Schichten und die Vermischung von Experten, und zeigen, dass einige davon für das Modellieren im großen Maßstab entscheidend sind. Wir beschreiben auch mehrere interessante Eigenschaften dieser Architekturen, die das Training und die Evaluation von Jamba aufgedeckt haben, und planen, Checkpoints aus verschiedenen Ablaufläufen zu veröffentlichen, um die weitere Erforschung dieser neuartigen Architektur zu fördern. Die Gewichte unserer Implementierung von Jamba stellen wir öffentlich unter einer freizügigen Lizenz zur Verfügung.

English

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.