BlackMamba: Mezcla de Expertos para Modelos de Espacio de Estados
BlackMamba: Mixture of Experts for State-Space Models
February 1, 2024
Autores: Quentin Anthony, Yury Tokpanov, Paolo Glorioso, Beren Millidge
cs.AI
Resumen
Los modelos de espacio de estados (SSM, por sus siglas en inglés) han demostrado recientemente un rendimiento competitivo frente a los transformadores en benchmarks de modelado de lenguaje a gran escala, logrando una complejidad lineal en tiempo y memoria en función de la longitud de la secuencia. Mamba, un modelo SSM lanzado recientemente, muestra un rendimiento impresionante tanto en tareas de modelado de lenguaje como en el procesamiento de secuencias largas. Simultáneamente, los modelos de mezcla de expertos (MoE, por sus siglas en inglés) han mostrado un rendimiento notable mientras reducen significativamente los costos computacionales y de latencia en la inferencia, a expensas de una mayor huella de memoria. En este artículo, presentamos BlackMamba, una arquitectura novedosa que combina el SSM Mamba con MoE para obtener los beneficios de ambos. Demostramos que BlackMamba compite favorablemente tanto con Mamba como con los transformadores de referencia, y supera en FLOPs de inferencia y entrenamiento. Entrenamos completamente y liberamos modelos BlackMamba de 340M/1.5B y 630M/2.8B en 300B tokens de un conjunto de datos personalizado. Mostramos que BlackMamba hereda y combina los beneficios de las arquitecturas SSM y MoE, integrando la generación de complejidad lineal de SSM con una inferencia rápida y económica de MoE. Liberamos todos los pesos, puntos de control y código de inferencia como código abierto. Código de inferencia en: 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