MoE-Mamba: Modelos de Espaço de Estados Seletivos Eficientes com Mistura de Especialistas
MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts
January 8, 2024
Autores: Maciej Pióro, Kamil Ciebiera, Krystian Król, Jan Ludziejewski, Sebastian Jaszczur
cs.AI
Resumo
Os Modelos de Espaço de Estados (SSMs) tornaram-se concorrentes sérios no campo de modelagem sequencial, desafiando a dominância dos Transformers. Ao mesmo tempo, a Mistura de Especialistas (MoE) melhorou significativamente os LLMs baseados em Transformers, incluindo modelos open-source recentes de última geração. Propomos que, para desbloquear o potencial dos SSMs em termos de escalabilidade, eles devem ser combinados com MoE. Demonstramos isso no Mamba, um modelo recente baseado em SSM que alcança um desempenho notável, semelhante ao dos Transformers. Nosso modelo, MoE-Mamba, supera tanto o Mamba quanto o Transformer-MoE. Em particular, o MoE-Mamba atinge o mesmo desempenho do Mamba em 2,2 vezes menos etapas de treinamento, mantendo os ganhos de desempenho de inferência do Mamba em relação ao Transformer.
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.