Jamba: O Novo Modelo de Linguagem Híbrido do AI21 Labs

Os modelos de linguagem têm testemunhado avanços rápidos, com arquiteturas baseadas em Transformers liderando o caminho na processamento de linguagem natural. No entanto, à medida que os modelos escalam, os desafios de lidar com contextos longos, eficiência de memória e throughput tornaram-se mais pronunciados.

AI21 Labs introduziu uma nova solução com Jamba, um modelo de linguagem grande de última geração (LLM) que combina as forças de ambos os arquitetos Transformer e Mamba em um quadro híbrido. Este artigo detalha a arquitetura do Jamba, seu desempenho e aplicações potenciais.

Visão Geral do Jamba

Jamba é um modelo de linguagem grande híbrido desenvolvido pelo AI21 Labs, aproveitando uma combinação de camadas Transformer e Mamba, integradas com um Módulo de Mistura de Especialistas (MoE). Essa arquitetura permite que o Jamba equilibre o uso de memória, throughput e desempenho, tornando-o uma ferramenta poderosa para uma ampla gama de tarefas de NLP. O modelo é projetado para caber dentro de uma única GPU de 80GB, oferecendo alto throughput e uma pequena pegada de memória, mantendo o desempenho de ponta em vários benchmarks.

A Arquitetura do Jamba

A arquitetura do Jamba é a pedra angular de suas capacidades. É construída em um projeto híbrido inovador que intercala camadas Transformer com camadas Mamba, incorporando módulos MoE para aumentar a capacidade do modelo sem aumentar significativamente os requisitos computacionais.

1. Camadas Transformer

A arquitetura Transformer se tornou o padrão para os LLMs modernos devido à sua capacidade de lidar com processamento paralelo de forma eficiente e capturar dependências de longo alcance em texto. No entanto, seu desempenho é frequentemente limitado por altos requisitos de memória e computação, particularmente ao processar contextos longos. O Jamba aborda essas limitações integrando camadas Mamba, que exploraremos a seguir.

2. Camadas Mamba

Mamba é um modelo de espaço de estado recente (SSM) projetado para lidar com relações de longo alcance em sequências de forma mais eficiente do que os RNNs tradicionais ou mesmo os Transformers. As camadas Mamba são particularmente eficazes na redução da pegada de memória associada ao armazenamento de caches de chave-valor (KV) nos Transformers. Ao intercalar camadas Mamba com camadas Transformer, o Jamba reduz o uso geral de memória, mantendo um alto desempenho, especialmente em tarefas que exigem o manejo de contextos longos.

3. Módulo de Mistura de Especialistas (MoE)

O módulo MoE no Jamba introduz uma abordagem flexível para dimensionar a capacidade do modelo. O MoE permite que o modelo aumente o número de parâmetros disponíveis sem aumentar proporcionalmente os parâmetros ativos durante a inferência. No Jamba, o MoE é aplicado a algumas das camadas MLP, com o mecanismo de roteamento selecionando os principais especialistas para ativar para cada token. Essa ativação seletiva permite que o Jamba mantenha uma alta eficiência ao lidar com tarefas complexas.

A imagem abaixo demonstra a funcionalidade de uma cabeça de indução em um modelo híbrido de Atenção-Mamba, uma característica-chave do Jamba. Nesse exemplo, a cabeça de atenção é responsável por prever rótulos como “Positivo” ou “Negativo” em resposta a tarefas de análise de sentimento. As palavras destacadas ilustram como a atenção do modelo está fortemente focada em tokens de rótulo dos exemplos de few-shot, particularmente no momento crítico antes de prever o rótulo final. Esse mecanismo de atenção desempenha um papel crucial na capacidade do modelo de realizar aprendizado em contexto, onde o modelo deve inferir o rótulo apropriado com base no contexto dado e nos exemplos de few-shot.

As melhorias de desempenho oferecidas pela integração do MoE com a arquitetura híbrida de Atenção-Mamba são destacadas na Tabela. Ao usar o MoE, o Jamba aumenta sua capacidade sem aumentar proporcionalmente os custos computacionais. Isso é particularmente evidente no aumento significativo no desempenho em vários benchmarks, como HellaSwag, WinoGrande e Perguntas Naturais (NQ). O modelo com MoE não apenas alcança uma precisão mais alta (por exemplo, 66,0% no WinoGrande em comparação com 62,5% sem MoE), mas também demonstra log-probabilidades melhoradas em diferentes domínios (por exemplo, -0,534 no C4).

Características Arquiteturais-Chave

• Composição de Camadas: A arquitetura do Jamba consiste em blocos que combinam Mamba e camadas Transformer em uma proporção específica (por exemplo, 1:7, significando uma camada Transformer para cada sete camadas Mamba). Essa proporção é ajustada para o desempenho e eficiência ótimos.
• Integração do MoE: As camadas MoE são aplicadas a cada poucas camadas, com 16 especialistas disponíveis e os dois principais especialistas ativados por token. Essa configuração permite que o Jamba dimensione efetivamente, gerenciando os trade-offs entre o uso de memória e eficiência computacional.
• Normalização e Estabilidade: Para garantir a estabilidade durante o treinamento, o Jamba incorpora RMSNorm nas camadas Mamba, o que ajuda a mitigar problemas como grandes picos de ativação que podem ocorrer em escala.

Desempenho e Benchmarking do Jamba

O Jamba foi rigorosamente testado contra uma ampla gama de benchmarks, demonstrando um desempenho competitivo em toda a linha. As seguintes seções destacam alguns dos principais benchmarks onde o Jamba se destacou, mostrando suas forças em tarefas de NLP gerais e cenários de contexto longo.

1. Benchmarks de NLP Comuns

O Jamba foi avaliado em vários benchmarks acadêmicos, incluindo:

• HellaSwag (10-shot): Uma tarefa de raciocínio de senso comum onde o Jamba alcançou uma pontuação de desempenho de 87,1%, superando muitos modelos concorrentes.
• WinoGrande (5-shot): Outra tarefa de raciocínio onde o Jamba marcou 82,5%, novamente mostrando sua capacidade de lidar com raciocínio linguístico complexo.
• ARC-Desafio (25-shot): O Jamba demonstrou um desempenho forte com uma pontuação de 64,4%, refletindo sua capacidade de gerenciar perguntas de múltipla escolha desafiadoras.

Em benchmarks agregados como MMLU (5-shot), o Jamba alcançou uma pontuação de 67,4%, indicando sua robustez em tarefas diversas.

2. Avaliações de Contexto Longo

Uma das características destacadas do Jamba é sua capacidade de lidar com contextos extremamente longos. O modelo suporta uma longitude de contexto de até 256K tokens, a mais longa entre os modelos disponíveis publicamente. Essa capacidade foi testada usando o benchmark Needle-in-a-Haystack, onde o Jamba mostrou uma precisão de recuperação excepcional em diferentes longitudes de contexto, incluindo até 256K tokens.

3. Throughput e Eficiência

A arquitetura híbrida do Jamba melhora significativamente o throughput, particularmente com sequências longas.

Em testes comparando o throughput (tokens por segundo) entre diferentes modelos, o Jamba consistentemente superou seus pares, especialmente em cenários que envolvem grandes tamanhos de lote e contextos longos. Por exemplo, com um contexto de 128K tokens, o Jamba alcançou 3x o throughput do Mixtral, um modelo comparável.

Usando o Jamba: Python

Para desenvolvedores e pesquisadores ansiosos para experimentar o Jamba, o AI21 Labs forneceu o modelo em plataformas como Hugging Face, tornando-o acessível para uma ampla gama de aplicações. O seguinte trecho de código demonstra como carregar e gerar texto usando o Jamba:

<p>from transformers import AutoModelForCausalLM, AutoTokenizer</p>

<p>model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")</p>

<p>input_ids = tokenizer("No recente Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]</p>

<p>outputs = model.generate(input_ids, max_new_tokens=216)</p>

print(tokenizer.batch_decode(outputs))

import torch
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments

<p>tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
model = AutoModelForCausalLM.from_pretrained(
"ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)</p>

<p>lora_config = LoraConfig(r=8,
target_modules=[
"embed_tokens","x_proj", "in_proj", "out_proj", # mamba
"gate_proj", "up_proj", "down_proj", # mlp
"q_proj", "k_proj", "v_proj"
# atenção],
task_type="CAUSAL_LM", bias="none")</p>

<p>dataset = load_dataset("Abirate/english_quotes", split="train")
training_args = SFTConfig(output_dir="./results",
num_train_epochs=2,
per_device_train_batch_size=4,
logging_dir='./logs',
logging_steps=10, learning_rate=1e-5, dataset_text_field="quote")
trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args,
peft_config=lora_config, train_dataset=dataset,
)
trainer.train()