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

Apresentamos Jamba, um novo modelo base de linguagem grande baseado numa arquitetura híbrida inovadora de mistura de especialistas (MoE) Transformer-Mamba. Especificamente, Jamba intercala blocos de camadas Transformer e Mamba, aproveitando os benefícios de ambas as famílias de modelos. MoE é adicionado em algumas dessas camadas para aumentar a capacidade do modelo, mantendo o uso ativo de parâmetros gerenciável. Essa arquitetura flexível permite configurações específicas para recursos e objetivos. Na configuração particular que implementamos, resultamos em um modelo poderoso que cabe em uma única GPU de 80GB. Construído em grande escala, Jamba oferece alta taxa de processamento e pequena pegada de memória comparado aos Transformers tradicionais, e ao mesmo tempo desempenho de ponta em benchmarks padrão de modelos de linguagem e avaliações de contexto longo. Notavelmente, o modelo apresenta resultados fortes para comprimentos de contexto de até 256K tokens. Estudamos várias decisões arquiteturais, como combinar camadas Transformer e Mamba, e como misturar especialistas, mostrando que algumas são cruciais em modelagem em grande escala. Também descrevemos várias propriedades interessantes dessas arquiteturas que o treinamento e avaliação do Jamba revelaram, e planejamos liberar checkpoints de várias execuções de ablação para incentivar exploração adicional dessa arquitetura inovadora. Tornamos os pesos de nossa implementação do Jamba publicamente disponíveis sob uma licença permissiva.