Jamba: Een Hybride Transformer-Mamba Taalmodel

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

Samenvatting

We presenteren Jamba, een nieuw basis groot taal model gebaseerd op een innovatieve hybride Transformer-Mamba mixture-of-experts (MoE) architectuur. Specifiek interleeft Jamba blokken van Transformer- en Mamba-lagen, waardoor het de voordelen van beide modelfamilies benut. MoE wordt toegevoegd in sommige van deze lagen om de modelcapaciteit te vergroten terwijl het actieve parametergebruik beheersbaar blijft. Deze flexibele architectuur maakt configuraties mogelijk die afgestemd zijn op specifieke middelen en doelen. In de specifieke configuratie die we hebben geïmplementeerd, resulteert dit in een krachtig model dat past op een enkele 80GB GPU. Op grote schaal gebouwd, biedt Jamba een hoge doorvoersnelheid en een klein geheugenverbruik in vergelijking met standaard Transformers, en tegelijkertijd state-of-the-art prestaties op standaard taal model benchmarks en lange-context evaluaties. Opmerkelijk is dat het model sterke resultaten laat zien voor contextlengtes tot 256K tokens. We bestuderen verschillende architectuurkeuzes, zoals hoe Transformer- en Mamba-lagen te combineren, en hoe experts te mixen, en laten zien dat sommige ervan cruciaal zijn bij modellering op grote schaal. We beschrijven ook verschillende interessante eigenschappen van deze architecturen die de training en evaluatie van Jamba hebben onthuld, en zijn van plan checkpoints van verschillende ablatie runs vrij te geven, om verder onderzoek naar deze nieuwe architectuur aan te moedigen. We maken de gewichten van onze implementatie van Jamba publiekelijk beschikbaar onder een permissieve licentie.

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.