BlackMamba: Miscela di Esperti per Modelli a Spazio degli Stati

BlackMamba: Mixture of Experts for State-Space Models

February 1, 2024

Autori: Quentin Anthony, Yury Tokpanov, Paolo Glorioso, Beren Millidge

cs.AI

Abstract

I modelli a spazio di stati (SSM) hanno recentemente dimostrato prestazioni competitive rispetto ai transformer nei benchmark di modellazione del linguaggio su larga scala, raggiungendo al contempo una complessità temporale e di memoria lineare in funzione della lunghezza della sequenza. Mamba, un modello SSM rilasciato di recente, mostra prestazioni impressionanti sia nella modellazione del linguaggio che nelle attività di elaborazione di sequenze lunghe. Parallelamente, i modelli a miscela di esperti (MoE) hanno mostrato prestazioni notevoli riducendo significativamente i costi computazionali e di latenza dell'inferenza, a scapito di un maggiore utilizzo di memoria. In questo articolo, presentiamo BlackMamba, una nuova architettura che combina il modello SSM Mamba con MoE per ottenere i vantaggi di entrambi. Dimostriamo che BlackMamba compete efficacemente sia con Mamba che con i baseline basati su transformer, superandoli in termini di FLOPs di inferenza e addestramento. Abbiamo addestrato completamente e reso open-source modelli BlackMamba da 340M/1.5B e 630M/2.8B su 300B token di un dataset personalizzato. Mostriamo che BlackMamba eredita e combina i vantaggi delle architetture SSM e MoE, unendo la generazione a complessità lineare degli SSM con l'inferenza economica e veloce dei MoE. Rilasciamo open-source tutti i pesi, i checkpoint e il codice di inferenza. Codice di inferenza disponibile su: 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