BlackMamba : Mélange d'Experts pour les Modèles à Espace d'États
BlackMamba: Mixture of Experts for State-Space Models
February 1, 2024
Auteurs: Quentin Anthony, Yury Tokpanov, Paolo Glorioso, Beren Millidge
cs.AI
Résumé
Les modèles à espace d'états (SSM) ont récemment démontré des performances compétitives par rapport aux transformateurs sur des benchmarks de modélisation du langage à grande échelle, tout en atteignant une complexité temporelle et mémoire linéaire en fonction de la longueur de la séquence. Mamba, un modèle SSM récemment publié, montre des performances impressionnantes à la fois dans la modélisation du langage et dans le traitement de longues séquences. Parallèlement, les modèles à mélange d'experts (MoE) ont montré des performances remarquables tout en réduisant significativement les coûts de calcul et de latence lors de l'inférence, au prix d'une empreinte mémoire plus importante. Dans cet article, nous présentons BlackMamba, une nouvelle architecture qui combine le SSM Mamba avec MoE pour bénéficier des avantages des deux. Nous démontrons que BlackMamba performe de manière compétitive par rapport aux modèles de référence Mamba et transformateurs, et surpasse en termes de FLOPs d'inférence et d'entraînement. Nous avons entièrement entraîné et rendu open-source des modèles BlackMamba de 340M/1.5B et 630M/2.8B sur 300 milliards de tokens d'un jeu de données personnalisé. Nous montrons que BlackMamba hérite et combine les avantages des architectures SSM et MoE, alliant une génération à complexité linéaire du SSM à une inférence rapide et peu coûteuse du MoE. Nous rendons publics tous les poids, points de contrôle et le code d'inférence. Code d'inférence disponible à l'adresse : https://github.com/Zyphra/BlackMamba
English
State-space models (SSMs) have recently demonstrated competitive performance
to transformers at large-scale language modeling benchmarks while achieving
linear time and memory complexity as a function of sequence length. Mamba, a
recently released SSM model, shows impressive performance in both language
modeling and long sequence processing tasks. Simultaneously, mixture-of-expert
(MoE) models have shown remarkable performance while significantly reducing the
compute and latency costs of inference at the expense of a larger memory
footprint. In this paper, we present BlackMamba, a novel architecture that
combines the Mamba SSM with MoE to obtain the benefits of both. We demonstrate
that BlackMamba performs competitively against both Mamba and transformer
baselines, and outperforms in inference and training FLOPs. We fully train and
open-source 340M/1.5B and 630M/2.8B BlackMamba models on 300B tokens of a
custom dataset. We show that BlackMamba inherits and combines both of the
benefits of SSM and MoE architectures, combining linear-complexity generation
from SSM with cheap and fast inference from MoE. We release all weights,
checkpoints, and inference code open-source. Inference code at:
https://github.com/Zyphra/BlackMamba