BlackMamba: Mixture of Experts für State-Space-Modelle
BlackMamba: Mixture of Experts for State-Space Models
February 1, 2024
Autoren: Quentin Anthony, Yury Tokpanov, Paolo Glorioso, Beren Millidge
cs.AI
Zusammenfassung
State-Space-Modelle (SSMs) haben kürzlich bei groß angelegten Sprachmodellierungs-Benchmarks eine wettbewerbsfähige Leistung im Vergleich zu Transformern gezeigt, während sie eine lineare Zeit- und Speicherkomplexität in Abhängigkeit von der Sequenzlänge erreichen. Mamba, ein kürzlich veröffentlichter SSM-Modell, zeigt beeindruckende Leistungen sowohl in der Sprachmodellierung als auch bei Aufgaben zur Verarbeitung langer Sequenzen. Gleichzeitig haben Mixture-of-Expert-Modelle (MoE) bemerkenswerte Leistungen gezeigt, während sie die Rechen- und Latenzkosten der Inferenz erheblich reduzieren, allerdings auf Kosten eines größeren Speicherbedarfs. In diesem Artikel präsentieren wir BlackMamba, eine neuartige Architektur, die das Mamba-SSM mit MoE kombiniert, um die Vorteile beider Ansätze zu nutzen. Wir zeigen, dass BlackMamba sowohl gegen Mamba als auch gegen Transformer-Baselines wettbewerbsfähig abschneidet und in Bezug auf Inferenz- und Trainings-FLOPs übertrifft. Wir trainieren vollständig und veröffentlichen 340M/1,5B und 630M/2,8B BlackMamba-Modelle auf 300B Token eines benutzerdefinierten Datensatzes. Wir zeigen, dass BlackMamba die Vorteile sowohl der SSM- als auch der MoE-Architekturen erbt und kombiniert, indem es die lineare Komplexität der Generierung von SSM mit der kostengünstigen und schnellen Inferenz von MoE verbindet. Wir veröffentlichen alle Gewichte, Checkpoints und Inferenz-Code Open-Source. Inferenz-Code unter: 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