BlackMamba: Смесь экспертов для моделей пространства состояний
BlackMamba: Mixture of Experts for State-Space Models
February 1, 2024
Авторы: Quentin Anthony, Yury Tokpanov, Paolo Glorioso, Beren Millidge
cs.AI
Аннотация
Модели пространства состояний (SSM) недавно продемонстрировали конкурентоспособные результаты по сравнению с трансформерами на крупномасштабных бенчмарках языкового моделирования, достигая при этом линейной временной и пространственной сложности в зависимости от длины последовательности. Mamba, недавно выпущенная SSM-модель, показывает впечатляющие результаты как в языковом моделировании, так и в задачах обработки длинных последовательностей. Одновременно модели смеси экспертов (MoE) продемонстрировали выдающуюся производительность, значительно снижая вычислительные затраты и задержку при выводе, хотя и за счет увеличения объема используемой памяти. В данной статье мы представляем BlackMamba — новую архитектуру, которая объединяет SSM Mamba с MoE, чтобы получить преимущества обеих подходов. Мы показываем, что BlackMamba конкурирует с базовыми моделями Mamba и трансформерами, а также превосходит их по количеству операций (FLOPs) при выводе и обучении. Мы полностью обучаем и открываем исходный код для моделей BlackMamba с 340M/1.5B и 630M/2.8B параметров на 300B токенов из собственного набора данных. Мы демонстрируем, что BlackMamba наследует и объединяет преимущества как SSM, так и MoE архитектур, сочетая генерацию с линейной сложностью от SSM с быстрым и экономичным выводом от MoE. Мы публикуем все веса, контрольные точки и код для вывода в открытом доступе. Код для вывода доступен по адресу: 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