Jamba: Un Modello Linguistico Ibrido Transformer-Mamba

Jamba: A Hybrid Transformer-Mamba Language Model

March 28, 2024

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

Abstract

Presentiamo Jamba, un nuovo modello linguistico di base di grandi dimensioni basato su una innovativa architettura ibrida Transformer-Mamba con mixture-of-experts (MoE). Nello specifico, Jamba intercala blocchi di layer Transformer e Mamba, sfruttando i vantaggi di entrambe le famiglie di modelli. Il MoE viene aggiunto in alcuni di questi layer per aumentare la capacità del modello mantenendo gestibile l'utilizzo dei parametri attivi. Questa architettura flessibile consente configurazioni specifiche per risorse e obiettivi. Nella configurazione particolare che abbiamo implementato, otteniamo un modello potente che può essere contenuto in una singola GPU da 80GB. Costruito su larga scala, Jamba offre un'elevata produttività e una ridotta impronta di memoria rispetto ai Transformer tradizionali, mantenendo al contempo prestazioni all'avanguardia nei benchmark standard per modelli linguistici e nelle valutazioni su contesti lunghi. In modo notevole, il modello presenta risultati solidi per contesti fino a 256K token. Studiamo varie decisioni architetturali, come combinare i layer Transformer e Mamba e come miscelare gli esperti, dimostrando che alcune di esse sono cruciali nella modellazione su larga scala. Descriviamo inoltre diverse proprietà interessanti di queste architetture emerse durante l'addestramento e la valutazione di Jamba, e prevediamo di rilasciare checkpoint da varie esecuzioni di ablazione, per incoraggiare ulteriori esplorazioni di questa innovativa architettura. Rendiamo pubblicamente disponibili i pesi della nostra implementazione di Jamba sotto una licenza permissiva.

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.