BlackMamba: Mistura de Especialistas para Modelos de Espaço de Estados
BlackMamba: Mixture of Experts for State-Space Models
February 1, 2024
Autores: Quentin Anthony, Yury Tokpanov, Paolo Glorioso, Beren Millidge
cs.AI
Resumo
Modelos de espaço de estados (SSMs, na sigla em inglês) recentemente demonstraram desempenho competitivo em relação aos transformers em benchmarks de modelagem de linguagem em grande escala, ao mesmo tempo em que alcançam complexidade linear de tempo e memória em função do comprimento da sequência. Mamba, um modelo SSM lançado recentemente, mostra desempenho impressionante tanto em modelagem de linguagem quanto em tarefas de processamento de sequências longas. Simultaneamente, modelos de mistura de especialistas (MoE, na sigla em inglês) têm demonstrado desempenho notável enquanto reduzem significativamente os custos computacionais e de latência da inferência, às custas de uma maior pegada de memória. Neste artigo, apresentamos o BlackMamba, uma arquitetura inovadora que combina o SSM Mamba com o MoE para obter os benefícios de ambos. Demonstramos que o BlackMamba tem desempenho competitivo em relação tanto ao Mamba quanto aos baselines de transformers, e supera em FLOPs de inferência e treinamento. Treinamos totalmente e disponibilizamos publicamente modelos BlackMamba de 340M/1.5B e 630M/2.8B em 300B de tokens de um conjunto de dados personalizado. Mostramos que o BlackMamba herda e combina os benefícios das arquiteturas SSM e MoE, unindo a geração de complexidade linear do SSM com a inferência rápida e econômica do MoE. Disponibilizamos todos os pesos, checkpoints e código de inferência como código aberto. Código de inferência em: 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