MoE-Mamba: Modelos Eficientes de Espacios de Estados Selectivos con Mezcla de Expertos
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
Resumen
Los Modelos de Espacio de Estados (SSMs, por sus siglas en inglés) se han convertido en serios competidores en el campo del modelado secuencial, desafiando el dominio de los Transformers. Al mismo tiempo, la Mezcla de Expertos (MoE, por sus siglas en inglés) ha mejorado significativamente los modelos de lenguaje grandes (LLMs) basados en Transformers, incluyendo los modelos de código abierto más recientes y de vanguardia. Proponemos que, para desbloquear el potencial de los SSMs en términos de escalabilidad, estos deberían combinarse con MoE. Demostramos esto en Mamba, un modelo reciente basado en SSM que logra un rendimiento notable, similar al de los Transformers. Nuestro modelo, MoE-Mamba, supera tanto a Mamba como a Transformer-MoE. En particular, MoE-Mamba alcanza el mismo rendimiento que Mamba en 2.2 veces menos pasos de entrenamiento, manteniendo las mejoras en el rendimiento de inferencia de Mamba frente al 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.