Jamba (modelo de linguagem)

Jamba é uma família de grandes modelos de linguagem (LLM) desenvolvida pela empresa de pesquisa israelense AI21 Labs. Jamba apresenta a primeira arquitetura híbrida do seu tipo, que combina elementos-chave de duas abordagens dominantes no desenvolvimento de IA: Transformers e Modelos de Espaço de Estados (State Space Models, SSM), em particular, a arquitetura Mamba.

O principal objetivo do Jamba é resolver o compromisso fundamental dos LLMs modernos: alta qualidade e desempenho (característicos dos Transformers) versus eficiência e a capacidade de processar contextos ultralongos (característicos dos SSMs). Ao combinar essas abordagens e adicionar esparsidade por meio de Mixture-of-Experts (MoE), o Jamba oferece um modelo que é simultaneamente poderoso, eficiente e capaz de trabalhar com enormes volumes de texto em uma única requisição.

Arquitetura do Jamba em detalhes

O Jamba não apenas alterna camadas de Transformer e Mamba. Ele utiliza uma estrutura em blocos cuidadosamente projetada, onde cada bloco consiste em oito camadas.

Estrutura de um bloco do Jamba:

• Uma camada de Transformer: Esta camada é responsável pela compreensão "profunda" e raciocínios complexos. A arquitetura Mixture-of-Experts (MoE) está embutida nesta camada.
• Sete camadas de Mamba: Essas camadas seguem a camada de Transformer e são responsáveis pelo processamento eficiente da sequência e pela "propagação" da informação através de um longo contexto.

Essa estrutura assimétrica permite que o modelo gerencie eficientemente os recursos computacionais: as operações pesadas, mas poderosas, do Transformer são executadas com menos frequência, enquanto as operações leves e rápidas do Mamba são executadas com mais frequência.

Integração de Mixture-of-Experts (MoE)

No Jamba, a arquitetura MoE é utilizada para aumentar ainda mais a eficiência.

• O MoE é aplicado apenas aos blocos totalmente conectados (FFN) dentro das camadas de Transformer. As camadas Mamba permanecem densas.
• No primeiro modelo Jamba, havia 16 especialistas (experts).
• Para cada token, uma rede de roteamento seleciona os 2 melhores especialistas (gating Top-2).

Isso significa que, embora o número total de parâmetros do modelo seja grande (52 bilhões), em cada etapa do processamento de um token na camada de Transformer, apenas 2 dos 16 especialistas estão ativos, o que torna os cálculos muito rápidos.

Evolução dos modelos Jamba

Jamba-v0.1 (março de 2024)

O primeiro modelo apresentado nesta família tem as seguintes características:

Especificações técnicas do Jamba-v0.1

Característica	Valor
Número total de parâmetros	52 bilhões
Parâmetros ativos	~12 bilhões
Número de especialistas (MoE)	16 (2 ativos)
Janela de contexto	256.000 tokens
Licença	Apache 2.0

Graças à sua arquitetura híbrida, o Jamba-1 é capaz de processar um contexto de 256.000 tokens, o que equivale a um romance de aproximadamente 400 páginas, e pode ser implantado em uma única GPU de consumo com 80 GB de memória.

Jamba-1.5 (2024)

Em 2024, a AI21 Labs apresentou a família de modelos Jamba 1.5 atualizada, incluindo duas versões: Jamba 1.5 Mini (12B de parâmetros ativos de um total de 52B) e Jamba 1.5 Large (94B de parâmetros ativos de um total de 398B). Esses modelos demonstram melhorias significativas de desempenho:

• Inferência até 2,5 vezes mais rápida em contextos longos em comparação com os concorrentes.
• Suporte para nove idiomas, incluindo inglês, espanhol, francês e árabe.

Principais vantagens e desempenho

• Janela de contexto enorme: 256.000 tokens — uma das maiores janelas entre todos os modelos disponíveis (incluindo proprietários) no momento de seu lançamento. Isso torna o Jamba ideal para tarefas que exigem a análise de grandes documentos, como contratos jurídicos, artigos científicos, bases de código inteiras ou longos diálogos.
• Alto desempenho e eficiência: Nos testes, o Jamba demonstra um desempenho comparável ou superior aos principais modelos de código aberto de tamanho semelhante, como Llama e Mixtral, ao mesmo tempo em que apresenta uma taxa de transferência 3 vezes maior em contextos longos.
• Abertura e acessibilidade: O Jamba é distribuído sob a licença permissiva Apache 2.0, o que permite seu uso livre para fins comerciais e de pesquisa. Os pesos do modelo estão disponíveis na plataforma Hugging Face.

Resultados em benchmarks

O Jamba 1.5 apresenta resultados competitivos em vários benchmarks:

• O Jamba 1.5 Mini alcançou 46.1 pontos no Arena Hard, tornando-se o modelo público líder em sua categoria.
• O Jamba 1.5 Large alcançou 65.4 pontos no Arena Hard, superando o Llama 3.1 70B e 405B.

Aplicação e disponibilidade

O Jamba é otimizado para aplicações de negócios e suporta recursos como chamada de função, saída estruturada em JSON e processamento de documentos. O modelo está disponível em várias plataformas, incluindo:

• Hugging Face
• Google Cloud Vertex AI
• Microsoft Azure
• NVIDIA API catalog
• Amazon Bedrock
• AI21 Studio

Para suportar uma inferência economicamente eficiente, a AI21 Labs introduziu o ExpertsInt8 — uma nova técnica de quantização que permite hospedar o Jamba 1.5 Large em uma máquina com 8 GPUs de 80GB cada, sem perda de qualidade ao processar um contexto de 256K tokens.

Literatura

• Lieber, O.; et al. (2024). Jamba: A Hybrid Transformer‑Mamba Language Model. arXiv:2403.19887.
• Lieber, O.; et al. (2024). Jamba‑1.5 Models and ExpertsInt8 Quantization. OpenReview JFPaD7lpBD.
• Gu, A.; Dao, T. (203). Mamba: Linear‑Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
• Gu, A.; et al. (2021). S4: Efficiently Modeling Long Sequences with Structured State Spaces. arXiv:2111.00396.
• Fedus, W.; et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.
• Yun, L.; et al. (2024). Toward Inference‑Optimal Mixture‑of‑Expert Large Language Models. arXiv:2404.02852.
• Liu, J.; et al. (2024). A Survey on Mixture of Experts in Large Language Models. arXiv:2407.06204.
• Gupta, V.; et al. (2024). Lynx: Enabling Efficient MoE Inference through Dynamic Batch‑Aware Expert Selection. arXiv:2411.08982.
• Liu, J.; et al. (2024). A Survey on Inference Optimization Techniques for Mixture of Experts Models. arXiv:2412.14219.
• Hsieh, C.‑P.; et al. (2024). RULER: What's the Real Context Size of Your Long‑Context Language Models?. arXiv:2404.06654.