Suivez nous sur

Intelligence artificielle

BlackMamba : mélange d'experts pour les modèles d'état et d'espace

mm

Publié le

 on

BlackMamba : mélange d'experts pour les modèles d'état et d'espace

Le développement de grands modèles de langage (LLM) construits à partir de modèles de transformateur uniquement décodeurs a joué un rôle crucial dans la transformation du domaine du traitement du langage naturel (NLP), ainsi que dans l'avancement de diverses applications d'apprentissage en profondeur, notamment apprentissage par renforcement, analyse de séries chronologiques, traitement d’images et bien plus encore. Cependant, malgré leur évolutivité et leurs performances élevées, les LLM construits à partir de modèles de transformateurs à décodeur uniquement sont toujours confrontés à des lacunes importantes. Bien qu'expressif, le mécanisme d'attention dans les LLM dérivés de transformateurs nécessite des ressources de calcul élevées pendant l'inférence et la formation, ce qui nécessite une mémoire importante pour la longueur de la séquence et les FLOP quadratiques. Cette exigence de calcul élevée limite la longueur du contexte des modèles de transformateur, rendant les tâches de génération autorégressive proportionnellement coûteuses avec l'échelle, et entrave l'apprentissage à partir de flux de données continus et la capacité de traitement de séquence véritablement illimité.

Ces derniers temps, 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 des tests de modélisation à grande échelle tout en atteignant une complexité de mémoire en fonction de la longueur de séquence et du temps linéaire. De plus, Mamba, un modèle spatial 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 Mixture of Expert (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 détriment d'une empreinte mémoire plus importante. S'appuyant sur les modèles Mamba et MoE, cet article abordera BlackMamba, une nouvelle architecture qui combine le modèle spatial d'état Mamba avec les modèles MoE pour tirer parti des avantages offerts par les deux frameworks. Les expériences sur BlackMamba ont démontré sa capacité à surpasser le framework Mamba existant et les lignes de base du transformateur en termes de formation des FLOP et d'inférence. Les performances exceptionnelles du framework BlackMamba montrent qu'il peut combiner efficacement les capacités des frameworks Mamba et MoE, offrant une inférence rapide et rentable à partir de MoE avec une génération de complexité linéaire à partir de Mamba.

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

BlackMamba : une introduction au MoE pour les modèles spatiaux d'états

La progression des grands modèles de langage (LLM), en particulier ceux basés sur des architectures de transformateurs uniquement décodeurs, a notamment influencé le Traitement du langage naturel (PNL) domaine et étendu à diverses applications d’apprentissage profond, notamment l’apprentissage par renforcement, l’analyse de séries chronologiques, le traitement d’images et au-delà. Néanmoins, malgré leur évolutivité et leurs performances robustes, ces LLM basés sur un transformateur et un décodeur uniquement rencontrent des défis notables. Le mécanisme d’attention, une caractéristique clé des systèmes basés sur les transformateurs LLMs, nécessite des ressources informatiques étendues à la fois pour l’inférence et la formation. Cela implique un besoin de mémoire qui augmente avec la longueur de la séquence et les opérations de calcul (FLOP) qui augmentent quadratiquement. De tels besoins de calcul intensifs limitent la longueur du contexte des modèles, augmentent les coûts des tâches de génération autorégressive à mesure que le modèle évolue et entravent la capacité des modèles à apprendre à partir de flux de données continus ou à traiter efficacement des séquences de longueur illimitée. 

Des efforts importants ont été déployés ces dernières années pour tenter de surmonter ces limitations, et l'attention s'est portée sur la conception d'alternatives architecturales aux modèles canoniques de transformateur d'attention dense, les modèles SSM et MoE étant les architectures candidates les plus prometteuses. Le principal avantage obtenu en favorisant les modèles spatiaux d'états par rapport aux modèles d'architecture de transformateur est la complexité de calcul linéaire en ce qui concerne la longueur de séquence d'entrée offerte par les SSM, par opposition à la complexité quadratique offerte par les transformateurs. Théoriquement, la complexité de calcul linéaire par rapport à la longueur de la séquence d'entrée permet aux modèles spatiaux d'état de traiter des séquences plus grandes que les modèles d'architecture de transformateur pour un budget FLOPS ou d'opérations à virgule flottante par seconde donné, et de rendre la génération autorégressive constante dans le calcul sans cache KV. Les modèles spatiaux d'état récemment développés, notamment Mamba, RetNet et quelques autres, ont démontré une inférence et une formation efficaces sur de longues séquences, ainsi que des performances de tâches de modélisation de langage compétitives pour des transformateurs ayant des propriétés de mise à l'échelle similaires. D'autre part, le mélange d'architectures de modèles experts gagne en popularité comme alternative aux transformateurs denses car il facilite une réduction significative des FLOP d'inférence et de formation essentiels pour obtenir une qualité comparable à celle d'un modèle dense. Les modèles MoE (Mixture of Experts) fonctionnent en activant uniquement une sélection clairsemée de l’ensemble des paramètres au cours d’un seul passage vers l’avant. Ils utilisent une fonction de routage pour déterminer quels « experts » sont appelés à agir en fonction du contexte donné. Cette approche crée une séparation entre le coût de calcul de l'inférence et le nombre total de paramètres, permettant des performances améliorées dans un budget d'inférence fixe, bien qu'avec un nombre accru de paramètres et un besoin en mémoire plus important.

Cette avancée architecturale offre des avantages notables par rapport aux transformateurs traditionnels et représente une direction passionnante pour un développement ultérieur. Nous postulons que l'intégration de ces améliorations dans un modèle combiné Mamba-MoE pourrait accélérer considérablement les capacités et l'efficacité de la modélisation du langage au-delà de celles des modèles de transformateur standard. Les avantages attendus d'une architecture Mamba-MoE par rapport à un modèle de transformateur dense traditionnel incluent :

Mamba: Atteint une complexité de calcul linéaire par rapport à la longueur de la séquence d’entrée pour les phases de formation et d’inférence. Il permet à la génération autorégressive de se produire dans un laps de temps constant et avec une utilisation constante de la mémoire.

Ministère de l'Environnement : Offre une vitesse d'inférence et une efficacité de calcul de formation comparables à celles d'un modèle de base plus petit et dense, tout en conservant un niveau de qualité de modèle qui rivalise avec celui d'un modèle avec un nombre de paramètres équivalent à celui de la version plus dense.

Cela étant dit, il est essentiel d'affirmer que les modèles d'architecture de transformateur sont toujours à la pointe de la technologie et ont démontré des performances constantes et remarquables sur les tâches de modélisation de langage et les tâches de traitement de séquence. À la base, l'architecture du transformateur utilise une auto-attention qui effectue une comparaison quadratique tout-à-tout des similitudes du produit scalaire entre les intégrations de différents jetons dans une séquence et effectue une carte linéaire avec un vecteur de sortie. Le modèle de transformateur se compose de blocs d'auto-attention empilés entre des blocs MLP ou Multi-Layer Perceptron qui consistent en outre en un MLP à deux couches avec une fonction d'activation donnée. 

BlackMamba : Architecture et Méthodologie

Modèles d'espace d'état

Les modèles spatiaux d'états appartiennent au groupe des modèles de séquence avec une complexité linéaire par rapport à la longueur de la séquence d'entrée. L'architecture des modèles spatiaux d'état s'aligne davantage sur les réseaux de neurones récurrents et les réseaux de neurones convolutifs plutôt que sur l'architecture basée sur l'attention, et s'inspire d'un système dynamique continu qui cartographie 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 des scénarios pratiques, la nature récurrente des modèles spatiaux d’état est la raison pour laquelle ils doivent encore être adoptés sur du matériel d’IA hautement parallèle comme les GPU. Cependant, l'émergence de MSS comme RWKV et Mamba ont utilisé des noyaux d'analyse parallèle pour mapper efficacement les opérations récurrentes sur les GPU, facilitant ainsi la formation de nouvelles architectures avec une efficacité comparable à celles obtenues par les modèles de transformateur. 

La complexité quadratique inhérente à la longueur des séquences au sein des transformateurs est une limitation bien connue qui entrave le raisonnement et la compréhension sur des contextes très longs. Des innovations récentes ont introduit l'idée d'étendre la longueur du contexte, permettant ainsi aux transformateurs d'être formés à une échelle réalisable avant d'être appliqués à des contextes beaucoup plus longs lors de l'inférence. Malgré ces progrès, le processus d'inférence nécessite toujours une quantité considérable de ressources de calcul et de mémoire, en particulier pour maintenir le cache Key-Value (KV), ce qui en fait une entreprise gourmande en ressources. Les efforts de recherche récents se sont concentrés sur l'amélioration des capacités d'expression des modèles d'espace d'état en incorporant des mécanismes de déclenchement dépendants des entrées, semblables aux matrices Requête, Clé, Valeur (QKV) trouvées dans les mécanismes d'attention. 

Ces efforts visent à préserver la progression intrinsèquement linéaire de la récursion dans l'espace d'état, permettant une exécution efficace par convolution ou par un processus d'analyse sélective. Cette approche réduit considérablement la disparité de performances avec les transformateurs dans les applications pratiques. Parmi ces avancées, Mamba se distingue comme un modèle d'espace d'état qui reflète les objectifs de recherches antérieures, montrant des niveaux de performances impressionnants comparables à ceux des transformateurs à des échelles allant jusqu'à 2.8 milliards de paramètres. Il y parvient en appliquant un contrôle dépendant des entrées aux entrées de la récursion du modèle d'espace d'état (SSM), tout en garantissant un calcul efficace grâce à l'utilisation de noyaux d'analyse sélective sur mesure.

Mélange de modèles 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 direct. Au lieu d'utiliser tous les paramètres, ces modèles dirigent les jetons vers des experts spécifiques en Perceptron multicouche (MLP). Idéalement, chaque expert est conçu pour traiter un type particulier d'entrée, avec un mécanisme de routage, essentiellement un réseau neuronal compact, déterminant 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 de calcul considérablement réduites. En règle générale, le routeur est un mappage des couches linéaires depuis les jetons vers les indices experts, chaque expert étant simplement un transformateur standard Perceptron multicouche. Cependant, les développeurs n'ont pas encore trouvé la méthode de formation optimale pour le routeur, car le problème d'affectation des experts est indifférenciable et les modèles mixtes d'experts ont souvent du mal à équilibrer la charge et à stabilité de formation entre différents experts pour l'efficacité matérielle. 

Architecture

À la base, BlackMamba utilise un modèle de transformateur standard composé de blocs MLP entrelacés et de blocs d'attention ajoutés en séquence le long d'un flux résiduel. Désormais, la majorité des modèles Mixture of Expert remplacent simplement les blocs de perceptrons multicouches par une couche experte routée. D'autre part, le framework BlackMamba remplace non seulement le bloc perceptron multicouche dans le transformateur par une couche experte routée, mais remplace également la couche d'attention par une couche Mamba State Space Model. L'architecture du framework BlackMamba est illustrée dans la figure suivante. 

Formation et ensemble de données

Le modèle BlackMamba est formé sur plus de 300 milliards de jetons sur un ensemble de données personnalisé et utilise la fonction d'activation SwiGLU pour les perceptrons multicouches experts. Le framework s'entraîne avec 8 experts, un nombre que les développeurs ont trouvé comme étant le bon équilibre entre l'empreinte mémoire et le coût d'inférence du modèle. L'ensemble de données personnalisé utilisé pour entraîner le framework BlackMamba se compose d'un mélange d'ensembles de données open source déjà existants, notamment Starcoder, SlimPajama, Pile, etc. Le tableau suivant montre les poids de chacun des ensembles de données utilisés pour la formation du framework BlackMamba. Au total, l’ensemble de données contient 1.8 XNUMX milliards de jetons. 

BlackMamba : Résultats

Pour garantir une comparaison équitable entre Mamba et BlackMamba, les développeurs ont formé les deux modèles avec les mêmes paramètres de formation sur les mêmes données de formation. Le framework BlackMamba est capable de surpasser les modèles Mamba et Transformer pour une taille de modèle de passe avant identique au moment de l'inférence ainsi que pour la formation d'opérations à virgule flottante par seconde. La figure suivante montre le temps nécessaire pour générer une séquence d'une longueur donnée de manière autorégressive à partir d'une invite initiale à un jeton en fonction de la longueur de la séquence. 

De plus, les avantages en matière de latence des modèles Mixture of Expert et Mamba sont combinés dans le framework BlackMamba, ce qui entraîne des temps d'inférence nettement plus rapides par rapport aux modèles de transformateur, aux modèles Mamba purs et aux modèles MoE. De plus, l'avantage d'inférence du framework 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. En continuant, la figure suivante illustre le nombre de jetons attribués aux modèles BlackMamba avec respectivement 340 millions et 640 millions de paramètres. Comme on peut le constater, la majorité des couches démontrent un niveau élevé d'équilibre expert grâce à l'algorithme Sinkhorn amélioré mis en œuvre par les modèles BlackMamba. 

Le tableau suivant couvre les scores d'évaluation du framework BlackMamba par rapport à une gamme de modèles de langage open source pré-entraînés. Comme on peut l'observer, le framework BlackMamba est capable de rivaliser et de surpasser la majorité des frameworks dans toutes les lignes de base. En outre, il convient de noter que les modèles qui surpassent BlackMamba ont un nombre de paramètres considérablement plus élevé et que l'écart de performances est minime, ce qui indique la capacité du framework BlackMamba avec moins de paramètres. 

Réflexions finales

Dans cet article, nous avons parlé de BlackMamba, une nouvelle architecture qui combine le modèle spatial d'état Mamba avec un mélange de modèles experts pour profiter des avantages offerts par ces deux frameworks. Des expériences sur BlackMamba ont démontré qu'il surpasse le framework Mamba existant et les lignes de base du transformateur en termes de formation des FLOP et d'inférence. Les performances exceptionnelles du framework BlackMamba démontrent qu'il est capable d'hériter et de combiner exceptionnellement bien les capacités des frameworks Mamba et MoE puisqu'il combine l'inférence rapide et bon marché de MoE avec la génération de complexité linéaire de Mamba. Nous avons expliqué comment l'architecture du framework BlackMamba est capable de surpasser les grands modèles de langage fortement formés, le framework Mamba existant et le mélange de modèles experts en termes de FLOP de formation et de coût d'inférence. De plus, le framework BlackMamba hérite également de la génération FLOP et de la formation réduite du mélange de modèles experts et du framework Mamba simultanément. 

 

« Ingénieur de profession, écrivain de cœur ». Kunal est un rédacteur technique avec un amour et une compréhension profonds de l'IA et du ML, dédié à la simplification de concepts complexes dans ces domaines grâce à sa documentation engageante et informative.