Créée à Tel Aviv en 2017, la start-up israélienne AI21 Labs est un des pionniers des LLM. Ses modèles, Jurassic 2 et WordTune, ont longtemps semblé s’inscrire parmi les meilleurs avant que ChatGPT ne balaie tout sur son passage. L’an dernier, elle a levé 208 millions de dollars (196 millions d’euros), ce qui fait d’elle la start-up non américaine la mieux financée de l’IA générative, hormis la canadienne Cohere et, bien sûr, les géants chinois (lire Qant du 13 mars 2023, du 4 septembre 2023 et du 23 novembre 2023).
Elle vient de présenter Jamba, sans doute le modèle d'IA le plus innovant depuis GPT-4. Jamba est le fruit de la fusion entre l'architecture des transformers et Mamba, un modèle en espace d'état (space state model, ou SSM). L’application des SSM dans l’IA dérive largement des travaux de recherche d’Albert Gu à Stanford et Carnegie Mellon, en particulier à partir de 2021. Mais c’est le 22 décembre dernier qu’une nouvelle étude du chercheur, co-écrite et pré-publiée avec Tri Dao, a commencé à mobiliser le monde de la recherche en IA.
En à peine un trimestre, Mamba: Linear-Time Sequence Modeling with Selective State Spaces a déjà généré une bonne vingtaine d’émules, qui couvrent le traitement du langage naturel, la segmentation d’images et la vision par ordinateur, les jeux, la météo, la génomique, la finance et même l’entomologie…
Puissance du serpent
Jusqu’à présent, les modèles de fondation sont tous basés sur l’architecture Transformer, présentée pour la première fois dans une étude devenue illustre, Attention is all you need, par 8 chercheurs de Google. Mais elle garde des faiblesses, notamment la difficulté à maintenir le contexte sur de longues séquences de données. Dans de longues “conversations”, ChatGPT et ses rivaux tendent à oublier les informations précédentes et halluciner. Albert Gu résout le problème en intégrant des SSM sélectifs dans un réseau neuronal sans module d’attention : Mamba.
La version de Mamba à 3 milliards de paramètres s'est avérée plus efficace que des LLM deux fois plus importants, avec une vitesse d’inférence 5 fois supérieure. La nouvelle architecture excelle notamment dans trois modalités : le langage, l’audio et la génomique. Dans des modèles de diffusion, Mamba permet de générer des séquences vidéo plus longues, là où les approches conventionnelles échouent en raison de contraintes de mémoire.
Le héros des deux mondes
Jamba ("Joint Attention and Mamba") est doté d’une architecture hybride, par mélange d’experts – comme Mixtral et, semble-t-il, GPT-4 –, qui combine les deux mondes : c’est aussi bien un LLM qu’un SSM.
L'architecture hybride de Jamba • Source : Opher Lieber et al.
Cette combinaison permet à Jamba d'ingérer de plus longues séquences de données, lui conférant ainsi la capacité de gérer un contexte plus large que les grands modèles de langage traditionnels. Pour être précis, sa fenêtre de contexte atteint les 256 000 tokens, soit environ 210 pages de texte. Il peut gérer jusqu'à 140 000 jetons sur une seule GPU de 80 Go. On est certes loin du vertigineux million de tokens annoncé par Google pour Gemini 1.5 Ultra, mais au-dessus du reste du marché. Et à la différence d’un LLM, Jamba ne ralentit pas quand sa fenêtre de contexte s'alourdit. Avec 52 milliards de paramètres disponibles, dont 12 milliards actifs, il se révèle souvent meilleur que Llama-2 70B et les autres grands LLM open source.
Jamba obtient des résultats similaires aux principaux LLM open-source, avec une bien meilleure inférence • Source : AI21 Labs
Jamba est disponible sous licence Apache 2.0, avec peu de restrictions pour l'utilisation commerciale, mais AI21 Labs prévient qu'il n'est pas encore prêt pour une utilisation commerciale en raison de l'absence de garde-fous contre les contenus toxiques et les biais. La start-up prévoit cependant de créer ces mécanismes de sécurité et proposera alors une version de Jamba prête pour l'entreprise sur sa plateforme.
Pour en savoir plus :
- Opher Lieber et al., Jamba: A Hybrid Transformer-Mamba Language Model, Arxiv, 2024
- Albert Gu et al., Efficiently Modeling Long Sequences with Structured State Spaces, Arxiv 2021
- Albert Gu et al., Mamba: Linear-Time Sequence Modeling with Selective State Spaces, Arxiv 2023
- Yuta Oshima et al., SSM Meets Video Diffusion Models: Efficient Video Generation with Structured State Spaces, Arxiv, 2024
- Carmen Amo Alonso et al., State Space Models as Foundation Models: A Control Theoretic Overview, Arxiv, 2024
- Zhuangwei Shi, MambaStock: Selective state space model for stock prediction, Arxiv, 2024
- Qianning Wang et al., InsectMamba: Insect Pest Classification with State Space Model, Arxiv, 2024
- Jinhong Wang et al., Large Window-based Mamba UNet for Medical Image Segmentation: Beyond Convolution and Self-attention, Arxiv, 2024
- Mark Tech Post
- Techcrunch
- Silicon Angle