Jamba : Le nouveau modèle hybride Transformer-Mamba d’AI21 Labs

Les modèles de langage ont connu des progrès rapides, avec des architectures basées sur les Transformers à la tête de la charge dans le traitement du langage naturel. Cependant, à mesure que les modèles s’étendent, les défis liés à la gestion de longs contextes, à l’efficacité de la mémoire et au débit sont devenus plus prononcés.

AI21 Labs a introduit une nouvelle solution avec Jamba, un modèle de langage grand et avancé (LLM) qui combine les forces des architectures Transformer et Mamba dans un cadre hybride. Cet article détaille l’architecture, les performances et les applications potentielles de Jamba.

Présentation de Jamba

Jamba est un modèle de langage grand hybride développé par AI21 Labs, exploitant une combinaison de couches Transformer et de couches Mamba, intégrées avec un module Mixture-of-Experts (MoE). 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 une large gamme de tâches de traitement du langage naturel. Le modèle est conçu pour tenir dans une seule carte graphique de 80 Go, offrant un débit élevé et une petite empreinte mémoire tout en maintenant des performances de pointe sur divers benchmarks.

L’architecture de Jamba

L’architecture de Jamba est la pierre angulaire de ses capacités. Elle est basée sur une conception hybride novatrice qui intercale des couches Transformer avec des couches Mamba, en incorporant des modules MoE pour améliorer la capacité du modèle sans augmenter significativement les exigences de calcul.

1. Couches Transformer

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

2. Couches Mamba

Mamba est un modèle d’espace d’état récent (SSM) conçu pour gérer les relations à longue distance dans les séquences de manière plus efficace 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 de clés-valeurs (KV) dans les Transformers. En intercalant des couches Mamba avec des couches Transformer, Jamba réduit l’utilisation globale de la mémoire tout en maintenant de hautes performances, en particulier dans les tâches nécessitant la gestion de longs contextes.

3. Module Mixture-of-Experts (MoE)

Le module MoE dans Jamba introduit une approche flexible pour mettre à l’échelle 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 des couches MLP, avec le mécanisme de routage sélectionnant les meilleurs experts à activer pour chaque jeton. Cette activation sélective permet à Jamba de maintenir une grande efficacité tout en traitant des tâches complexes.

L’image ci-dessous démontre la fonctionnalité 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 responsable de prédire des étiquettes telles que « Positif » ou « Négatif » en réponse à des tâches d’analyse de sentiments. Les mots mis en évidence illustrent comment l’attention du modèle est fortement concentrée sur les jetons d’étiquette à partir des exemples à few-shot, en particulier au moment critique avant de prédire l’étiquette finale. Ce mécanisme d’attention joue un rôle crucial dans la capacité du modèle à effectuer un apprentissage en contexte, où le modèle doit déduire l’étiquette appropriée en fonction du contexte et des exemples à few-shot donnés.

Les améliorations de performances offertes par l’intégration de Mixture-of-Experts (MoE) avec 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’augmentation significative des performances sur divers benchmarks tels que HellaSwag, WinoGrande et Natural Questions (NQ). Le modèle avec MoE non seulement atteint une précision plus élevée (par exemple, 66,0 % sur WinoGrande par rapport à 62,5 % sans MoE), mais démontre également des log-probabilités améliorées sur différents domaines (par exemple, -0,534 sur C4).

Caractéristiques architecturales clés

• Composition des couches : L’architecture de Jamba se compose de blocs qui combinent des couches Mamba et des 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.
• <strong<Intégration MoE : Les couches MoE sont appliquées toutes les quelques couches, avec 16 experts disponibles et les deux meilleurs experts activés par jeton. Cette configuration permet à Jamba de mettre à l’échelle de manière efficace tout en gérant les compromis entre l’utilisation de la mémoire et l’efficacité de calcul.
• Normalisation et stabilité : Pour assurer la stabilité pendant l’entraînement, Jamba intègre RMSNorm dans les couches Mamba, ce qui aide à atténuer des problèmes tels que les pics d’activation importants qui peuvent survenir à grande échelle.

Performances et évaluation de Jamba

Jamba a été rigoureusement testé sur une large gamme de benchmarks, démontrant des performances compétitives dans l’ensemble. Les sections suivantes mettent en évidence certains des benchmarks clés sur lesquels Jamba a excellé, mettant en valeur ses forces dans les tâches de traitement du langage naturel générales et les scénarios de contexte long.

1. Benchmarks NLP courants

Jamba a été évalué sur plusieurs benchmarks universitaires, notamment :

• HellaSwag (10-shot) : une tâche de raisonnement de bon sens où Jamba a obtenu un score de performance de 87,1 %, dépassant de nombreux modèles concurrents.
• WinoGrande (5-shot) : une autre tâche de raisonnement où Jamba a obtenu un score de 82,5 %, démontrant à nouveau sa capacité à gérer des raisonnements linguistiques complexes.
• ARC-Challenge (25-shot) : Jamba a démontré de fortes performances avec un score de 64,4 %, reflétant sa capacité à gérer des questions de choix multiple complexes.

Dans des benchmarks agrégés comme MMLU (5-shot), Jamba a obtenu un score de 67,4 %, indiquant sa robustesse sur diverses tâches.

2. Évaluations de contexte long

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

3. Débit et efficacité

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

Dans les tests comparant le débit (jetons par seconde) entre différents modèles, Jamba a constamment surpassé ses pairs, en particulier dans les scénarios impliquant de grandes tailles de lots et de longs contextes. Par exemple, avec un contexte de 128 K jetons, Jamba a atteint 3 fois le débit de Mixtral, un modèle comparable.

Utilisation de Jamba : Python

Pour les développeurs et les chercheurs impatients d’expérimenter avec Jamba, AI21 Labs a mis le modèle à disposition sur des plateformes comme Hugging Face, le rendant accessible pour une large gamme d’applications. Le code suivant montre comment charger et générer du texte en utilisant 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("Dans le Super Bowl LVIII récent,", 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()