Jamba : le nouveau modèle de langage hybride Transformer-Mamba d'AI21 Labs

Les modèles linguistiques ont connu des progrès rapides, les architectures basées sur Transformer étant en tête du traitement du langage naturel. Cependant, à mesure que les modèles évoluent, les défis liés à la gestion des contextes longs, à l'efficacité de la mémoire et au débit sont devenus plus prononcés.

Laboratoires AI21 a introduit une nouvelle solution avec Jamba, un modèle de langage étendu (LLM) de pointe qui combine les atouts de Transformer et de Architectures Mamba dans un environnement hybride. Cet article détaille l'architecture, les performances et les applications potentielles de Jamba.

Présentation de Jamba

Jamba est un grand modèle de langage hybride développé par AI21 Labs, exploitant une combinaison de couches Transformer et de couches Mamba, intégrées à un Mélange d'experts (MoE) module. Cette architecture permet à Jamba d'équilibrer l'utilisation de la mémoire, le débit et les performances, ce qui en fait un outil puissant pour un large éventail de tâches NLP. Le modèle est conçu pour s'intégrer dans un seul GPU de 80 Go, offrant un débit élevé et une faible empreinte mémoire tout en conservant des performances de pointe sur divers benchmarks.

L'architecture de Jamba

L'architecture de Jamba est la pierre angulaire de ses capacités. Elle repose sur une conception hybride innovante qui entrelace les couches Transformer et Mamba, intégrant des modules MoE pour améliorer les capacités du modèle sans augmenter significativement les besoins de calcul.

1. Couches de transformateur

L'architecture Transformer est devenue la norme pour les LLM modernes en raison de sa capacité à gérer efficacement le traitement parallèle et à capturer les dépendances à longue portée dans le texte. Cependant, ses performances sont souvent limitées par des exigences élevées en matière de mémoire et de calcul, en particulier lors du traitement de contextes longs. Jamba résout ces limitations en intégrant les couches Mamba, que nous explorerons ensuite.

2. Couches Mamba

Mamba est un modèle espace-état (SSM) récent conçu pour gérer les relations longue distance dans des séquences plus efficacement que les RNN traditionnels ou même les Transformers. Les couches Mamba sont particulièrement efficaces pour réduire l'empreinte mémoire associée au stockage des caches clé-valeur (KV) dans Transformers. En entrelaçant les couches Mamba avec les couches Transformer, Jamba réduit l'utilisation globale de la mémoire tout en conservant des performances élevées, en particulier dans les tâches nécessitant une longue gestion du contexte.

3. Mélange d'experts (MoE) Modules

MoE Le module Jamba introduit une approche flexible pour faire évoluer la capacité du modèle. MoE permet au modèle d'augmenter le nombre de paramètres disponibles sans augmenter proportionnellement les paramètres actifs pendant l'inférence. Dans Jamba, MoE est appliqué à certaines couches MLP, le mécanisme du routeur sélectionnant les meilleurs experts à activer pour chaque jeton. Cette activation sélective permet à Jamba de maintenir une efficacité élevée tout en gérant des tâches complexes.

L'image ci-dessous illustre le fonctionnement d'une tête d'induction dans un modèle hybride Attention-Mamba, une fonctionnalité clé de Jamba. Dans cet exemple, la tête d'attention est chargée de prédire des étiquettes telles que « Positif » ou « Négatif » en réponse à des tâches d'analyse des sentiments. Les mots surlignés illustrent la forte concentration de l'attention du modèle sur les jetons d'étiquettes issus des exemples à faible échantillonnage, notamment au moment critique précédant la prédiction de l'étiquette finale. Ce mécanisme d'attention joue un rôle crucial dans la capacité du modèle à effectuer un apprentissage contextuel, où il doit déduire l'étiquette appropriée en fonction du contexte donné et des exemples à faible échantillonnage.

Les améliorations de performances offertes par l'intégration du mélange d'experts (MoE) à l'architecture hybride Attention-Mamba sont mises en évidence dans le tableau. En utilisant MoE, Jamba augmente sa capacité sans augmenter proportionnellement les coûts de calcul. Cela est particulièrement évident dans l'amélioration significative des performances sur divers benchmarks tels que HellaSwag, WinoGrande et Natural Questions (NQ). Le modèle avec MoE atteint non seulement une précision plus élevée (par exemple, 66.0 % sur WinoGrande contre 62.5 % sans MoE), mais démontre également des probabilités log améliorées dans différents domaines (par exemple, -0.534 sur C4).

Principales caractéristiques architecturales

• Composition des couches : L'architecture de Jamba se compose de blocs qui combinent Mamba et les couches Transformer dans un rapport spécifique (par exemple, 1:7, ce qui signifie une couche Transformer pour sept couches Mamba). Ce rapport est réglé pour des performances et une efficacité optimales.
• Intégration du ministère de l'Environnement : Les couches MoE sont appliquées toutes les quelques couches, avec 16 experts disponibles et les 2 meilleurs experts activés par jeton. Cette configuration permet à Jamba d'évoluer efficacement tout en gérant les compromis entre l'utilisation de la mémoire et l'efficacité de calcul.
• Normalisation et stabilité : Pour garantir la stabilité pendant la formation, Jamba intègre RMSNorm dans les couches Mamba, ce qui permet d'atténuer les problèmes tels que les pics d'activation importants qui peuvent survenir à grande échelle.

Performances et benchmarking de Jamba

Jamba a été rigoureusement testé par rapport à un large éventail de références, démontrant des performances compétitives à tous les niveaux. Les sections suivantes mettent en évidence certains des points clés dans lesquels Jamba a excellé, mettant en valeur ses atouts à la fois dans les tâches générales de PNL et dans les scénarios à contexte long.

1. Benchmarks courants en PNL

Jamba a été évalué sur plusieurs critères académiques, notamment :

• HellaSwag (10-coup):Une tâche de raisonnement de bon sens où Jamba a obtenu un score de performance de 87.1 %, surpassant de nombreux modèles concurrents.
• WinoGrande (5 coups): Une autre tâche de raisonnement où Jamba a obtenu un score de 82.5 %, démontrant une fois de plus sa capacité à gérer un raisonnement linguistique complexe.
• ARC-Challenge (25 coups): Jamba a démontré de solides performances avec un score de 64.4 %, reflétant sa capacité à gérer des questions à choix multiples difficiles.

Dans des benchmarks globaux comme MMLU (5-shot), Jamba a obtenu un score de 67.4 %, indiquant sa robustesse dans diverses tâches.

2. Évaluations à contexte long

L'une des fonctionnalités phares de Jamba est sa capacité à gérer des contextes extrêmement longs. Le modèle prend en charge des contextes d'une longueur maximale de 256 256 jetons, soit la plus longue parmi les modèles disponibles publiquement. Cette capacité a été testée à l'aide du benchmark Needle-in-a-Haystack, où Jamba a démontré une précision de récupération exceptionnelle sur différentes longueurs de contexte, jusqu'à XNUMX XNUMX jetons.

3. Débit et efficacité

L'architecture hybride de Jamba améliore considérablement le débit, en particulier avec les séquences longues.

Lors de tests comparant le débit (jetons par seconde) entre différents modèles, Jamba a systématiquement surpassé ses pairs, en particulier dans les scénarios impliquant des lots de grande taille et des contextes longs. Par exemple, avec un contexte de 128 3 jetons, Jamba a atteint un débit XNUMX fois supérieur à celui de Mixtral, un modèle comparable.

Utiliser Jamba : Python

Pour les développeurs et les chercheurs désireux d’expérimenter Jamba, AI21 Labs a fourni le modèle sur des plateformes comme Hugging Face, le rendant ainsi accessible à un large éventail d’applications. L'extrait de code suivant montre comment charger et générer du texte à l'aide de Jamba :

from transformers import AutoModelForCausalLM, AutoTokenizer

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

input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]

outputs = model.generate(input_ids, max_new_tokens=216)

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

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

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" 
# attention],
task_type="CAUSAL_LM", bias="none")

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()