Pour son premier LLM codeur ouvert, Mistral AI choisit une architecture alternative

Pour développer une version 7B de son modèle Codestral, Mistral AI n’a pas utilisé de transformeur, mais l’architecture Mamba2.

À quand le support de Codestral Mamba dans llama.cpp ? Mistral AI invite à garder un œil sur l’évolution de cette demande.

L’usage en local est effectivement l’une des promesses de ce LLM « spécial code ». À environ 7 milliards de paramètres, il est plus petit que le modèle Codestral 22B disponible depuis quelques semaines. Il s’en distingue aussi par sa licence (Apache 2.0)… et par son architecture, dont il tire son nom : Mamba.

Mamba a émergé il y a quelques mois pour pallier les limites des transformeurs lors du traitement de longues séquences. Dans ces scénarios, le mécanisme d’attention constitue un goulet d’étranglement, du fait qu’il utilise une forme de cache clé-valeur permettant à chaque token d’accéder aux précédents lors de la prédiction. Plus la taille de contexte augmente, plus l’empreinte mémoire et la latence augmentent, de façon quadratique.

Des méthodes telles que la fenêtre glissante et l’attention flash peuvent atténuer cet effet. Mamba va plus loin en remplaçant le composant d’attention par un mécanisme inspiré de la théorie du contrôle : les SSM (State Space Models). Avec eux, la montée en charge est linéaire. On permet aux paramètres SSM d’être fonction de l’input, de sorte qu’une sélection des informations à conserver s’opère au moment la mémorisation – et non au moment de la remémoration, comme c’est le cas pour les transformeurs.

architecture Mamba
Mamba combine le bloc H3 à la base de la plupart des architectures SSM avec le bloc MLP des réseaux neuronaux modernes. Au lieu d’interlacer les deux blocs, il répète le bloc Mamba de manière homogène. Au niveau du bloc H3, la première porte multiplicative est remplacée par une fonction d’activation. Dans le bloc MLP, on ajoute un SSM sur la branche principale.

Une version Mamba pour Codestral

L’architecture Mamba se prête à des cas d’usage impliquant un grand contexte et une mémoire de long terme. Ses deux créateurs – de Princeton et de Carnegie Mellon – l’ont notamment mise à l’épreuve sur la production d’audio et la modélisation d’ADN.

Pour ce qui est de générer du code, Mistral AI a choisi de comparer Codestral Mamba aux modèles suivants :

benchmark Codestral Mamba
Plus ou moins au niveau de Codestral sur ces benchmarks, DeepSeek utilise une toute autre architecture (fenêtre 4k glissante).

La fenêtre d’expérimentation s’est limitée à 256k tokens en dépit de la capacité théorique de Mamba à en gérer des millions. En ligne de mire, une fois encore, l’usage en local.

Llama.cpp prend déjà en charge Mamba… mais pas Mamba2, version sur laquelle repose Codestral 7B. En attendant, on peut déployer le modèle avec le SDK mistral-inference ou via TensorRT-LLM. On peut aussi le tester sur La Plateforme, le « bac à sable » de Mistral AI.

Illustration principale générée par IA