Jamba (modèle de langage)

Jamba est une famille de grands modèles de langage (LLM) développée par la société de recherche israélienne AI21 Labs. Jamba introduit la première architecture hybride de son genre, qui combine les éléments clés de deux approches dominantes dans le développement de l'IA : les transformeurs et les modèles à espace d'états (State Space Models, SSM), en particulier l'architecture Mamba.

L'objectif principal de Jamba est de résoudre le compromis fondamental des LLM modernes : la haute qualité et la performance (caractéristiques des transformeurs) face à l'efficacité et à la capacité de traiter des contextes ultra-longs (caractéristiques des SSM). En combinant ces approches et en y ajoutant la sparsité grâce à Mixture-of-Experts (MoE), Jamba propose un modèle à la fois puissant, efficace et capable de traiter d'énormes volumes de texte en une seule requête.

Architecture de Jamba en détail

Jamba n'alterne pas simplement des couches de transformeur et de Mamba. Il utilise une structure en blocs soigneusement conçue, où chaque bloc est composé de huit couches.

Structure d'un bloc Jamba :

• Une couche de Transformeur : Cette couche est responsable de la compréhension « profonde » et du raisonnement complexe. L'architecture Mixture-of-Experts (MoE) y est intégrée.
• Sept couches de Mamba : Ces couches suivent la couche de transformeur et sont responsables du traitement efficace des séquences et de la propagation de l'information à travers de longs contextes.

Cette structure asymétrique permet au modèle de gérer efficacement les ressources de calcul : les opérations lourdes mais puissantes du transformeur sont exécutées moins souvent, tandis que les opérations légères et rapides de Mamba sont exécutées plus fréquemment.

Intégration de Mixture-of-Experts (MoE)

Jamba utilise une architecture MoE pour améliorer davantage son efficacité.

• Le MoE est appliqué uniquement aux blocs feed-forward (FFN) à l'intérieur des couches de transformeur. Les couches Mamba restent denses.
• Dans le premier modèle Jamba, il y avait 16 experts.
• Pour chaque token, un réseau de routage sélectionne les 2 meilleurs experts (Top-2 gating).

Cela signifie que bien que le nombre total de paramètres du modèle soit élevé (52 milliards), seuls 2 des 16 experts sont actifs à chaque étape de traitement d'un token dans la couche de transformeur, ce qui rend les calculs très rapides.

Évolution des modèles Jamba

Jamba-v0.1 (mars 2024)

Le premier modèle présenté dans cette famille possède les caractéristiques suivantes :

Spécifications techniques de Jamba-v0.1

Caractéristique	Valeur
Nombre total de paramètres	52 milliards
Paramètres actifs	~12 milliards
Nombre d'experts (MoE)	16 (2 actifs)
Fenêtre de contexte	256 000 tokens
Licence	Apache 2.0

Grâce à son architecture hybride, Jamba-1 est capable de traiter un contexte de 256 000 tokens, ce qui équivaut à un roman d'environ 400 pages, et peut être déployé sur un seul GPU grand public avec 80 Go de mémoire.

Jamba-1.5 (2024)

En 2024, AI21 Labs a présenté la famille de modèles mise à jour Jamba 1.5, qui inclut deux versions : Jamba 1.5 Mini (12 milliards de paramètres actifs sur 52 milliards au total) et Jamba 1.5 Large (94 milliards de paramètres actifs sur 398 milliards au total). Ces modèles présentent des améliorations significatives en termes de performance :

• Une inférence jusqu'à 2,5 fois plus rapide sur les longs contextes par rapport aux concurrents.
• La prise en charge de neuf langues, dont l'anglais, l'espagnol, le français et l'arabe.

Principaux avantages et performances

• Fenêtre de contexte immense : 256 000 tokens — l'une des plus grandes fenêtres parmi tous les modèles disponibles (y compris propriétaires) au moment de sa sortie. Cela rend Jamba idéal pour les tâches nécessitant l'analyse de grands documents : contrats juridiques, articles scientifiques, bases de code entières ou longs dialogues.
• Haute performance et efficacité : Dans les tests, Jamba affiche des performances comparables ou supérieures aux principaux modèles ouverts de taille similaire, tels que Llama et Mixtral, tout en offrant un débit 3 fois plus élevé sur les longs contextes.
• Ouverture et accessibilité : Jamba est distribué sous la licence permissive Apache 2.0, ce qui permet son utilisation libre à des fins commerciales et de recherche. Les poids du modèle sont disponibles sur la plateforme Hugging Face.

Résultats sur les benchmarks

Jamba 1.5 obtient des résultats compétitifs sur divers benchmarks :

• Jamba 1.5 Mini a obtenu un score de 46.1 sur Arena Hard, ce qui en fait le modèle public de premier plan dans sa catégorie.
• Jamba 1.5 Large a obtenu un score de 65.4 sur Arena Hard, dépassant Llama 3.1 70B et 405B.

Application et disponibilité

Jamba est optimisé pour les applications d'entreprise et prend en charge des fonctionnalités telles que l'appel de fonctions, la sortie structurée en JSON et le traitement de documents. Le modèle est disponible sur de nombreuses plateformes, notamment :

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

Pour prendre en charge une inférence rentable, AI21 Labs a introduit ExpertsInt8, une nouvelle technique de quantification qui permet de déployer Jamba 1.5 Large sur une machine avec 8 GPU de 80 Go chacun, sans perte de qualité lors du traitement d'un contexte de 256K tokens.

Bibliographie

• 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. (2023). 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.