Jamba : Un modèle de langage hybride Transformer-Mamba

Jamba: A Hybrid Transformer-Mamba Language Model

March 28, 2024

Auteurs: 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

Résumé

Nous présentons Jamba, un nouveau modèle de langage de base de grande taille, basé sur une architecture hybride novatrice combinant Transformer et Mamba avec un mécanisme de mixture d'experts (MoE). Plus précisément, Jamba entrelace des blocs de couches Transformer et Mamba, bénéficiant ainsi des avantages des deux familles de modèles. Le MoE est ajouté à certaines de ces couches pour augmenter la capacité du modèle tout en gardant l'utilisation des paramètres actifs gérable. Cette architecture flexible permet des configurations adaptées aux ressources et aux objectifs spécifiques. Dans la configuration particulière que nous avons implémentée, nous obtenons un modèle puissant qui tient sur un seul GPU de 80 Go. Conçu à grande échelle, Jamba offre un débit élevé et une empreinte mémoire réduite par rapport aux Transformers classiques, tout en atteignant des performances de pointe sur les benchmarks standards de modèles de langage et les évaluations de contexte long. De manière remarquable, le modèle présente des résultats solides pour des contextes allant jusqu'à 256 000 tokens. Nous étudions diverses décisions architecturales, telles que la manière de combiner les couches Transformer et Mamba, et la façon de mélanger les experts, et montrons que certaines d'entre elles sont cruciales pour la modélisation à grande échelle. Nous décrivons également plusieurs propriétés intéressantes de ces architectures que l'entraînement et l'évaluation de Jamba ont révélées, et prévoyons de publier des points de contrôle issus de diverses expériences d'ablation, afin d'encourager l'exploration approfondie de cette architecture novatrice. Nous rendons les poids de notre implémentation de Jamba disponibles publiquement sous une licence permissive.

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.