MoE-Mamba: Эффективные селективные модели пространства состояний с использованием смеси экспертов

MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts

January 8, 2024

Авторы: Maciej Pióro, Kamil Ciebiera, Krystian Król, Jan Ludziejewski, Sebastian Jaszczur

cs.AI

Аннотация

Модели пространства состояний (SSM) стали серьезными конкурентами в области последовательного моделирования, бросая вызов доминированию трансформеров. В то же время смесь экспертов (MoE) значительно улучшила трансформерные языковые модели (LLM), включая недавние передовые модели с открытым исходным кодом. Мы предлагаем, что для раскрытия потенциала SSM в масштабировании их следует комбинировать с MoE. Мы демонстрируем это на примере Mamba, недавней модели на основе SSM, которая достигает впечатляющей производительности, сравнимой с трансформерами. Наша модель, MoE-Mamba, превосходит как Mamba, так и Transformer-MoE. В частности, MoE-Mamba достигает той же производительности, что и Mamba, за 2.2 раза меньше шагов обучения, сохраняя при этом преимущества Mamba в производительности на этапе вывода по сравнению с трансформером.

English

State Space Models (SSMs) have become serious contenders in the field of sequential modeling, challenging the dominance of Transformers. At the same time, Mixture of Experts (MoE) has significantly improved Transformer-based LLMs, including recent state-of-the-art open-source models. We propose that to unlock the potential of SSMs for scaling, they should be combined with MoE. We showcase this on Mamba, a recent SSM-based model that achieves remarkable, Transformer-like performance. Our model, MoE-Mamba, outperforms both Mamba and Transformer-MoE. In particular, MoE-Mamba reaches the same performance as Mamba in 2.2x less training steps while preserving the inference performance gains of Mamba against the Transformer.