Mathstral, Codestral Mamba : pour Mistral AI, « small is beautiful »

Mathstral : prouver les bienfaits du fine-tuning de Mistral 7B

Mathstral est une déclinaison de Mistral-7B « instruite » pour résoudre des problèmes de mathématiques avancées.

 Il a été concocté dans le cadre de la collaboration de Mistral AI avec le projet scientifique Numina. Celui-ci a été fondé par Jia Li, Yann Fleureau, Hélène Evain (trois anciens de la startup Cardiologs), Guillaume Lample (cofondateur et directeur scientifique chez Mistral AI) et Stanislas Polu (cofondateur de Dust et ex-ingénieur de recherche chez OpenAI).

L’ONG Numina entend rassembler une grande base de données de problèmes de mathématiques et leurs résolutions, les utiliser pour entraîner des modèles d’IA accessibles librement et fournir des outils pour ce faire. Outre Mistral AI, Numina reçoit le soutien d’Hugging Face, de General Catalyst, d’Answer AI et du centre international de recherche en mathématiques de Beijing (Peking University). Pour l’instant, Numina a décliné deux modèles basés sur les LLM de la startup chinoise DeepSeek.

Avec 7 milliards de paramètres et sa fenêtre de contexte de 32 000 tokens, Mathstral décroche un score de 56,6 % au benchmark générique MATH, au coude à coude avec Qwen2 7B d’Alibaba et devant Gemma 2 9B (48,3 %) ou Llama 3 8B (28,4 %).

Mistral 7B n’atteignait que 13,1 % sur cette évaluation et Mathstral le dépasse de peu sur le fameux MMLU (60,1 % contre 63,47 %). Selon Mistral AI, en allongeant la durée d’inférence (sans dire de combien), il obtient un score de 68,37 % au même benchmark. 

Sur le test GSM8K, le LLM instruit surpasse légèrement Llama 3 8B (77,1 % contre 75,4 %), mais demeure derrière DeepSeek Math 7B. Il affiche la note de 56,9 % au test GRE Math, uniquement doublé par QWen 2 7B (58,2 %) dans le tableau présenté par la startup. Sur les parangonnages Odyssey Math et AMC 2023, Mathstral est devant ses concurrents du moment (37,2 % et 42,30 %). Les LLM semblent toutefois lutter pour résoudre les 30 problèmes de la compétition AIME 2024. Mathstral en a réussi deux sur les 30, tout comme QWen 2 7B, quand les autres modèles de langage échouent totalement ou en complètent un seul.

Ici, Mistral AI veut prouver les bienfaits du fine-tuning, qui plus est avec de « petits » LLM. « Mathstral est un autre exemple des excellents compromis performance/vitesse obtenus lors de l’entraînement de modèles à des fins spécifiques – une philosophie de développement que nous encourageons activement dans La Plateforme, en particulier avec ses nouvelles capacités de fine-tuning », assure l’équipe de Mistral AI.

Codestral Mamba : une nouvelle architecture dans le portfolio de Mistral

La startup ne cible pas le même objectif avec Codestral Mamba. Après l’adoption des architectures denses (Mistral 7B) et SMoE (Sparse Mixture of Experts), Mistral Ai s’essaye à Mamba 2.

Mamba est une architecture mise sur pied par les chercheurs Albert Gu, du pôle machine learning de l’Université Carnegie Mellon et Tri Dao, directeur scientifique de Together AI et professeur assistant au département de sciences computationnelles de l’université de Princeton.

L’architecture Mamba est pensée pour résoudre l’un des plus gros défauts des modèles Transformer et de leurs mécanismes d’attention.

« L’efficacité de l’auto-attention est attribuée à sa faculté d’acheminer l’information de manière dense dans une fenêtre contextuelle, ce qui lui permet de modéliser des données complexes », rappellent les deux chercheurs. « Toutefois, cette propriété présente des inconvénients fondamentaux : l’incapacité de modéliser quoi que ce soit en dehors d’une fenêtre finie, et une échelle quadratique par rapport à la longueur de la fenêtre ».

Pour « la faire courte », la fenêtre de contexte des LLM est justement… trop courte. Plusieurs techniques ont émergé pour tenter de résoudre ce problème, mais sans résultat probant, considèrent-ils.

En réponse à cela, Albert Gu est l’un des artisans des « modèles séquentiels espace-état structurés » (structured state space sequence models ou SSM). « Ces modèles peuvent être interprétés comme une combinaison de réseaux neuronaux récurrents (RNN) et convolutifs (CNN) avec une inspiration des modèles classiques espace-état », indiquent les chercheurs.

Les modèles espace-état aident en statistiques à modéliser des séries temporelles et des systèmes dynamiques. Ils décrivent un système par un ensemble d’équations, divisant le problème en un état latent non observable (des variables cachées désignant l’état interne) et des observations mesurables (des variables ou signaux). Ces modèles sont particulièrement utiles pour des applications comme le filtrage, le lissage et la prévision de données.

Les SSM héritent indirectement de ces capacités et ont fait leurs preuves pour traiter des signaux de données continues comme dans l’audio et la vidéo. C’était moins le cas pour le traitement de données textuelles.

Albert Gu et Tri Dao entendent boucher ce trou dans la raquette avec des modèles espace-état sélectifs (selective state space models).

Au lieu d’un mécanisme d’attention, Mamba est doté d’un mécanisme sélectif qui permet de se concentrer ou d’ignorer des éléments d’une séquence en entrée (à l’entraînement ou à l’inférence). Les chercheurs avaient alors remplacé les couches d’attention et de perceptron multicouche d’un Transformer par un bloc « simplifié ». En contrepartie, le SSM doit être entraîné comme un réseau de neurones récurrents, ce qui à la fois moins coûteux et moins adapté aux orientations des infrastructures actuelles (les GPU et les frameworks associés sont fortement optimisés pour les Transformers).

Mamba essaime. Outre le développement parallèle de SSM, il a inspiré certaines architectures hybrides SSM-Transformer, dont celles développées par A21Labs, Jamba. Il est un exemple dans la conception d’approches sensiblement différentes comme l’architecture Griffin de Google, qui a donné naissance à ReccurentGemma et mise plutôt sur une attention locale.

Avec le framework pensé pour créer Mamba 2, les chercheurs ont étendu les capacités de fusion entre SSM et les éléments cœurs des Transformer, tout en simplifiant les opérations d’entraînement. Ce framework, nommé dualité espace-état (space state duality ou SSD), « établit des liens entre les SSM, les matrices structurées et l’attention (linéaire) ».

Résultat, Mamba 2, toujours un réseau de neurones récurrent, est deux à huit fois plus véloce que Mamba premier du nom tout en étant compétitif face à l’architecture Transformer. Il est surtout plus efficient à l’entraînement.  

« Contrairement aux modèles Transformer, les modèles Mamba offrent l’avantage d’une inférence en temps linéaire et la possibilité théorique de modéliser des séquences de longueur infinie », résume Mistral AI, qui a été aidé par les deux chercheurs pour entraîner Codestral Mamba. « Ils permettent aux utilisateurs d’exploiter le modèle de manière intensive et d’obtenir des réponses rapides, quelle que soit la longueur de l’entrée ».

Cette architecture qui a fait ses preuves pour des données audiovisuelles et génomiques serait également intéressante pour la programmation, ajoute Mistral AI.

Résultat, Codestral Mamba Instruct, avec sa fenêtre de contexte de 256 000 tokens et ses 7,3 milliards de paramètres, recueille des scores proches de son grand frère Codestral 22B (22 milliards de paramètres) sur les benchmarks HumanEval (75 % vs 81,1 %), HumanEval C++ (59,8 % vs 65,2 %), HumanEval Java (57 % vs 63,3 %), Spider (58,8 % vs 63,5 %). Il parvient même à le battre contre le benchmark CruxE (57,8 % vs 51,3 %). Codestral 22B conserve une avance notable de plus de dix points face à MBPP et HumanEval Bash. 

Pour autant, et c’est sans doute le plus important pour Mistral AI, Codestral Mamba dépasse CodeGemma 1.1 7B, CodeLlama 7B, DeepSeek v 1.5 7B et, dans la plupart des CodeLlama 34B.

« C’est le modèle de code le plus performant pour sa taille, [il est] parfait pour les applications de type copilote », vante Théophile Gervet, chercheur chez Mistral AI, sur X (ex-Twitter).

Mathstral et Codestral Mamba sont disponibles sur HuggingFace, Ollama et sur La Plateforme, l’endroit où sont réunies les API payantes de la startup. Pour inférer ces deux LLM, Mistral AI recommande une machine équipée d’un GPU doté d’au moins 16 Go de VRAM.