Le développement de grands modèles de langage (LLM) basés sur des modèles de transformateur décodage-seul a joué un rôle crucial dans la transformation du domaine du traitement automatique des langues (NLP), ainsi que dans l’avancement de diverses applications d’apprentissage automatique, notamment l’apprentissage par renforcement, l’analyse de séries chronologiques, le traitement d’images et bien plus encore. Cependant, malgré leur scalabilité et leur solide performance, les LLM basés sur des modèles de transformateur décodage-seul font face à des limites importantes. Bien que ces modèles soient expressifs, le mécanisme d’attention dans les LLM dérivés des transformateurs nécessite des ressources computationnelles importantes pendant l’inférence et l’entraînement, nécessitant une quantité considérable de mémoire pour la longueur de la séquence et des opérations de calcul (FLOPs) qui augmentent de manière quadratique. Cette exigence computationnelle élevée limite la longueur du contexte des modèles, rend les tâches de génération autoregressive coûteuses à mesure que le modèle se scale, et entrave la capacité des modèles à apprendre à partir de flux de données continus et à traiter des séquences de longueur illimitée de manière efficace.

Récemment, les modèles d’espace d’état (SSM) ont démontré des capacités et des performances remarquables, rivalisant avec les modèles d’architecture de transformateur dans les benchmarks de modélisation à grande échelle, tout en atteignant une complexité de mémoire en fonction de la longueur de la séquence et une complexité temporelle linéaire. De plus, Mamba, un modèle d’espace d’état récemment publié, a montré des performances exceptionnelles dans une gamme de tâches de modélisation de langage et de traitement de séquences longues. Simultanément, les modèles de mélange d’experts (MoE) ont également montré des performances impressionnantes, tout en réduisant considérablement la latence et les coûts de calcul de l’inférence, bien qu’au prix d’une empreinte mémoire plus importante. En s’appuyant sur les modèles Mamba et MoE, cet article discutera de BlackMamba, une nouvelle architecture qui combine le modèle d’espace d’état Mamba avec les modèles MoE pour exploiter les avantages offerts par ces deux cadres. Les expériences menées sur BlackMamba ont démontré sa capacité à surpasser le cadre Mamba existant et les modèles de transformateur de référence en termes de FLOPs d’entraînement et de coût d’inférence. La performance exceptionnelle du cadre BlackMamba montre qu’il peut combiner efficacement les capacités des cadres Mamba et MoE, offrant une inférence rapide et rentable grâce à MoE, ainsi qu’une génération à complexité linéaire grâce à Mamba.

Cet article vise à couvrir en profondeur le cadre BlackMamba. Nous explorons le mécanisme, la méthodologie et l’architecture du cadre, ainsi que sa comparaison avec les cadres de génération d’images et de vidéos de pointe. Commençons.

BlackMamba : Introduction aux MoE pour les Modèles d’Espace d’État

L’évolution des grands modèles de langage (LLM), en particulier ceux basés sur des architectures de transformateur décodage-seul, a eu un impact notable sur le domaine du traitement automatique des langues (NLP) et s’est étendu à diverses applications d’apprentissage automatique, notamment l’apprentissage par renforcement, l’analyse de séries chronologiques, le traitement d’images et au-delà. Néanmoins, malgré leur scalabilité et leur solide performance, ces LLM basés sur des architectures de transformateur décodage-seul rencontrent des défis importants. Le mécanisme d’attention, une caractéristique clé des LLM basés sur les transformateurs, nécessite des ressources computationnelles importantes pour l’inférence et l’entraînement. Cela implique un besoin de mémoire qui augmente avec la longueur de la séquence et des opérations de calcul (FLOPs) qui augmentent de manière quadratique. De telles exigences computationnelles élevées limitent la longueur du contexte des modèles, rendent les tâches de génération autoregressive coûteuses à mesure que le modèle se scale, et entravent la capacité des modèles à apprendre à partir de flux de données continus ou à traiter des séquences de longueur illimitée de manière efficace.

Des efforts importants ont été déployés ces dernières années pour surmonter ces limites, et l’attention s’est portée sur la conception d’alternatives architecturales aux modèles de transformateur canoniques à attention dense, les modèles d’espace d’état (SSM) et les modèles de mélange d’experts (MoE) étant les architectures les plus prometteuses. Le principal avantage de l’utilisation de modèles d’espace d’état par rapport aux modèles d’architecture de transformateur est la complexité computationnelle linéaire par rapport à la longueur de la séquence d’entrée offerte par les SSM, par opposition à la complexité quadratique offerte par les transformateurs. Théoriquement, la complexité linéaire par rapport à la longueur de la séquence d’entrée permet aux modèles d’espace d’état de traiter des séquences plus longues que les modèles d’architecture de transformateur pour un budget de FLOPS ou d’opérations par seconde donné, et de rendre la génération autoregressive constante en calcul sans cache KV. Les modèles d’espace d’état récemment développés, notamment RWKV et Mamba, ont utilisé des noyaux de balayage parallèle pour cartographier les opérations récurrentes de manière efficace sur les GPU, facilitant ainsi l’entraînement de nouvelles architectures avec une efficacité comparable à celle des modèles de transformateur.

Les modèles de mélange d’experts (MoE) gagnent en popularité comme alternative aux transformateurs denses, car ils permettent une réduction significative des FLOPs d’inférence et d’entraînement, essentielle pour atteindre une qualité comparable à celle d’un modèle dense. Les modèles MoE fonctionnent en activant uniquement une sélection parcimonieuse des paramètres totaux lors d’un seul passage avant. Ils utilisent une fonction de routage pour déterminer quels « experts » sont appelés à l’action en fonction du contexte donné. Cette approche crée une séparation entre le coût computationnel de l’inférence et le nombre total de paramètres, permettant une meilleure performance dans un budget d’inférence fixe, bien qu’avec une empreinte mémoire plus importante.

BlackMamba : Architecture et Méthodologie

Modèles d’Espace d’État

Les modèles d’espace d’état appartiennent au groupe des modèles de séquence à complexité linéaire par rapport à la longueur de la séquence d’entrée. L’architecture des modèles d’espace d’état s’aligne plus sur les réseaux de neurones récurrents et les réseaux de neurones convolutionnels que sur l’architecture basée sur l’attention, et est inspirée d’un système dynamique continu qui mappe une fonction unidimensionnelle à travers un espace latent implicite. Un système dynamique linéaire rend les calculs parallèles efficaces en utilisant soit un balayage associatif, soit un balayage de convolution. Dans les scénarios pratiques, la nature récurrente des modèles d’espace d’état a été la raison pour laquelle ils n’ont pas encore été adoptés sur le matériel d’IA hautement parallèle comme les GPU. Cependant, l’émergence de SSM comme RWKV et Mamba a utilisé des noyaux de balayage parallèle pour cartographier les opérations récurrentes de manière efficace sur les GPU, facilitant ainsi l’entraînement de nouvelles architectures avec une efficacité comparable à celle des modèles de transformateur.

La complexité quadratique inhérente par rapport à la longueur de la séquence dans les transformateurs est une limitation bien connue qui entrave la compréhension et la raisonnement sur des contextes très longs. Les innovations récentes ont introduit l’idée d’étendre la longueur du contexte, permettant aux transformateurs d’être entraînés sur une échelle réalisable avant d’être appliqués à des contextes beaucoup plus longs pendant l’inférence. Malgré ces progrès, le processus d’inférence nécessite toujours une quantité considérable de ressources computationnelles et de mémoire, en particulier pour maintenir le cache KV, ce qui en fait une entreprise coûteuse en ressources. Les efforts de recherche récents se sont concentrés sur l’amélioration des capacités expressives des modèles d’espace d’état en incorporant des mécanismes de contrôle d’entrée dépendants, similaires aux matrices QKV trouvées dans les mécanismes d’attention.

Ces efforts visent à préserver la progression linéaire inhérente de la récursion d’espace d’état, permettant une exécution efficace par soit convolution, soit un processus de balayage sélectif. Cette approche réduit considérablement l’écart de performance avec les transformateurs dans les applications pratiques. Parmi ces progrès, Mamba se distingue comme un modèle d’espace d’état qui reflète les objectifs de la recherche antérieure, montrant des niveaux de performance impressionnants comparables à ceux des transformateurs à des échelles allant jusqu’à 2,8 milliards de paramètres. Il atteint cela en appliquant un contrôle d’entrée dépendant des entrées de la récursion du modèle d’espace d’état (SSM), tout en assurant un calcul efficace grâce à l’utilisation de noyaux de balayage sélectif sur mesure.

Modèles de Mélange d’Experts

Les modèles de mélange d’experts (MoE) réalisent une séparation entre le coût d’inférence et le nombre total de paramètres en activant sélectivement les paramètres pendant le passage avant. Au lieu d’utiliser tous les paramètres, ces modèles dirigent les jetons vers des experts MLP (Multilayer Perceptron) spécifiques. Idéalement, chaque expert est conçu pour traiter un type d’entrée particulier, avec un mécanisme de routage, essentiellement un réseau neural compact, qui détermine l’expert le plus approprié pour chaque jeton. Cette approche vise à préserver la puissance expressive globale d’un modèle avec un nombre équivalent de paramètres dans une configuration plus dense, mais avec des exigences computationnelles considérablement réduites. Typiquement, le routeur est une mise en correspondance des couches linéaires à partir des jetons aux indices d’experts, avec chaque expert étant un MLP standard. Cependant, les développeurs n’ont pas encore déterminé la méthode de formation optimale pour le routeur, car le problème d’affectation d’experts est non différentiable, et les modèles de mélange d’experts ont souvent des difficultés avec l’équilibrage de charge et la stabilité d’entraînement entre les différents experts pour l’efficacité du matériel.

Architecture

Au cœur du BlackMamba se trouve un modèle de transformateur standard composé de blocs MLP et d’attention intercalés le long d’un flux résiduel. Maintenant, la plupart des modèles de mélange d’experts remplacent simplement les blocs de perceptron multicouche par une couche d’expert routée. D’un autre côté, le cadre BlackMamba remplace non seulement le bloc de perceptron multicouche dans le transformateur par une couche d’expert routée, mais remplace également la couche d’attention par une couche de modèle d’espace d’état Mamba. L’architecture du cadre BlackMamba est démontrée dans la figure suivante.

Entraînement et Données

Le modèle BlackMamba est entraîné sur plus de 300 milliards de jetons sur un jeu de données personnalisé et utilise la fonction d’activation SwiGLU pour les perceptrons multicouches des experts. Le cadre s’entraîne avec 8 experts, un nombre que les développeurs ont trouvé être l’équilibre et le compromis appropriés entre l’empreinte mémoire et le coût d’inférence du modèle. Le jeu de données personnalisé utilisé pour entraîner le cadre BlackMamba se compose d’un mélange de jeux de données open source existants, notamment Starcoder, SlimPajama, Pile, et plus encore. La table suivante démontre les poids de chaque jeu de données utilisé pour l’entraînement du cadre BlackMamba. Au total, il y a 1,8 billion de jetons dans le jeu de données.

BlackMamba : Résultats

Pour garantir une comparaison équitable entre Mamba et BlackMamba, les développeurs ont entraîné les deux modèles avec les mêmes paramètres d’entraînement sur les mêmes données d’entraînement. Le cadre BlackMamba est capable de surpasser à la fois Mamba et les modèles de transformateur pour une taille de modèle identique lors du passage avant, ainsi qu’en termes de FLOPs d’entraînement et d’inférence. La figure suivante démontre le temps nécessaire pour générer une séquence d’une longueur donnée de manière autoregressive à partir d’une invite initiale d’un jeton en fonction de la longueur de la séquence.

De plus, les avantages en termes de latence des modèles de mélange d’experts et de Mamba sont combinés dans le cadre BlackMamba, ce qui entraîne des temps d’inférence considérablement plus rapides par rapport aux modèles de transformateur, aux modèles Mamba purs et aux modèles MoE. De plus, l’avantage en termes d’inférence du cadre BlackMamba est directement proportionnel à la longueur des séquences, ce qui rend BlackMamba extrêmement efficace pour la génération de séquences longues. Ensuite, la figure suivante illustre le nombre de jetons attribués aux modèles BlackMamba avec 340 millions et 640 millions de paramètres, respectivement. Comme on peut le voir, la majorité des couches présentent un niveau élevé d’équilibre d’expertise en raison de l’algorithme Sinkhorn amélioré mis en œuvre par les modèles BlackMamba.

La table suivante présente les scores d’évaluation du cadre BlackMamba par rapport à une gamme de modèles de langage pré-entraînés open source. Comme on peut l’observer, le cadre BlackMamba est capable de rivaliser et de surpasser la majorité des cadres sur toutes les références. De plus, il est important de noter que les modèles qui surpassent BlackMamba ont un nombre considérablement plus élevé de paramètres, et l’écart de performance est minime, indiquant la capacité du cadre BlackMamba avec moins de paramètres.

Pensées Finales

Dans cet article, nous avons discuté de BlackMamba, une nouvelle architecture qui combine le modèle d’espace d’état Mamba avec les modèles de mélange d’experts pour exploiter les avantages offerts par ces deux cadres. Les expériences menées sur BlackMamba ont démontré sa capacité à surpasser le cadre Mamba existant et les modèles de transformateur de référence en termes de FLOPs d’entraînement et d’inférence. La performance exceptionnelle du cadre BlackMamba montre qu’il peut combiner de manière exceptionnelle les capacités des cadres Mamba et MoE, offrant une inférence rapide et rentable grâce à MoE, ainsi qu’une génération à complexité linéaire grâce à Mamba. Nous avons discuté de la manière dont l’architecture du cadre BlackMamba est capable de surpasser les modèles de langage à grande échelle entraînés, le cadre Mamba existant et les modèles de mélange d’experts en termes de FLOPs d’entraînement et de coût d’inférence. De plus, le cadre BlackMamba hérite également des FLOPs de génération et de la formation réduite à la fois des modèles de mélange d’experts et du cadre Mamba.