BlackMamba: Mixture of Experts voor State-Space Modellen
BlackMamba: Mixture of Experts for State-Space Models
February 1, 2024
Auteurs: Quentin Anthony, Yury Tokpanov, Paolo Glorioso, Beren Millidge
cs.AI
Samenvatting
State-space modellen (SSM's) hebben recentelijk concurrerende prestaties laten zien ten opzichte van transformers bij grootschalige taalmodelleringsbenchmarks, terwijl ze lineaire tijd- en geheugencomplexiteit bereiken als functie van sequentielengte. Mamba, een recent uitgebracht SSM-model, toont indrukwekkende prestaties in zowel taalmodellering als taken voor het verwerken van lange sequenties. Tegelijkertijd hebben mixture-of-expert (MoE) modellen opmerkelijke prestaties getoond terwijl ze de reken- en latentiekosten van inferentie aanzienlijk verlagen, ten koste van een groter geheugenverbruik. In dit artikel presenteren we BlackMamba, een nieuwe architectuur die de Mamba SSM combineert met MoE om de voordelen van beide te benutten. We demonstreren dat BlackMamba concurrerend presteert ten opzichte van zowel Mamba als transformer-baselines, en overtreft in inferentie- en trainings-FLOPs. We trainen volledig en openbaren 340M/1.5B en 630M/2.8B BlackMamba-modellen op 300B tokens van een aangepaste dataset. We tonen aan dat BlackMamba zowel de voordelen van SSM- als MoE-architecturen erft en combineert, waarbij lineaire-complexiteit generatie van SSM wordt gecombineerd met goedkope en snelle inferentie van MoE. We geven alle gewichten, checkpoints en inferentiecode openbaar. Inferentiecode is beschikbaar op: 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