Jamba: Um Modelo de Linguagem Híbrido Transformer-Mamba

Jamba: A Hybrid Transformer-Mamba Language Model

March 28, 2024

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

Resumo

Apresentamos o Jamba, um novo modelo de linguagem de grande escala baseado em uma arquitetura híbrida inovadora que combina Transformer e Mamba com uma abordagem de mistura de especialistas (MoE). Especificamente, o Jamba intercala blocos de camadas Transformer e Mamba, aproveitando os benefícios de ambas as famílias de modelos. O MoE é adicionado em algumas dessas camadas para aumentar a capacidade do modelo, mantendo o uso de parâmetros ativos em um nível gerenciável. Essa arquitetura flexível permite configurações específicas para diferentes recursos e objetivos. Na configuração particular que implementamos, obtemos um modelo poderoso que cabe em uma única GPU de 80GB. Construído em grande escala, o Jamba oferece alta taxa de transferência e uma pequena pegada de memória em comparação com Transformers convencionais, ao mesmo tempo em que alcança desempenho de ponta em benchmarks padrão de modelos de linguagem e avaliações de contexto longo. Notavelmente, o modelo apresenta resultados robustos para comprimentos de contexto de até 256 mil tokens. Estudamos várias decisões arquiteturais, como combinar camadas Transformer e Mamba, e como misturar especialistas, e mostramos que algumas delas são cruciais na modelagem em grande escala. Também descrevemos várias propriedades interessantes dessas arquiteturas que o treinamento e a avaliação do Jamba revelaram, e planejamos liberar checkpoints de várias execuções de ablação, para incentivar uma exploração mais aprofundada dessa nova arquitetura. Disponibilizamos publicamente os pesos de nossa implementação do Jamba sob uma licença 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.