Enrutador Recurrente por Capas para Mezcla de Expertos

Resumen

La escalabilidad de los modelos de lenguaje grandes (LLMs) ha revolucionado sus capacidades en diversas tareas, sin embargo, este crecimiento debe ir de la mano de estrategias computacionales eficientes. La arquitectura de Mezcla de Expertos (MoE) destaca por su capacidad de escalar el tamaño del modelo sin aumentar significativamente los costos de entrenamiento. A pesar de sus ventajas, los modelos MoE actuales a menudo muestran ineficiencia en los parámetros. Por ejemplo, un LLM basado en MoE pre-entrenado con 52 mil millones de parámetros podría tener un rendimiento comparable a un modelo estándar con 6.7 mil millones de parámetros. Siendo una parte crucial de MoE, los enrutadores actuales en diferentes capas asignan tokens de forma independiente sin aprovechar la información de enrutamiento histórica, lo que podría llevar a combinaciones subóptimas de tokens-experto y al problema de ineficiencia de parámetros. Para aliviar este problema, presentamos el Enrutador Recurrente por Capas para Mezcla de Expertos (RMoE). RMoE aprovecha una Unidad Recurrente con Compuertas (GRU) para establecer dependencias entre decisiones de enrutamiento a lo largo de capas consecutivas. Esta recurrencia por capas puede calcularse eficientemente en paralelo para tokens de entrada e introduce costos negociables. Nuestras extensas evaluaciones empíricas demuestran que los modelos de lenguaje basados en RMoE superan consistentemente a una variedad de modelos de referencia. Además, RMoE integra una etapa de cálculo novedosa, ortogonal a los métodos existentes, lo que permite una compatibilidad perfecta con otras arquitecturas de MoE. Nuestros análisis atribuyen las mejoras de RMoE a su efectivo intercambio de información entre capas, lo que también mejora la selección y diversidad de expertos. Nuestro código está disponible en https://github.com/qiuzh20/RMoE

English

The scaling of large language models (LLMs) has revolutionized their capabilities in various tasks, yet this growth must be matched with efficient computational strategies. The Mixture-of-Experts (MoE) architecture stands out for its ability to scale model size without significantly increasing training costs. Despite their advantages, current MoE models often display parameter inefficiency. For instance, a pre-trained MoE-based LLM with 52 billion parameters might perform comparably to a standard model with 6.7 billion parameters. Being a crucial part of MoE, current routers in different layers independently assign tokens without leveraging historical routing information, potentially leading to suboptimal token-expert combinations and the parameter inefficiency problem. To alleviate this issue, we introduce the Layerwise Recurrent Router for Mixture-of-Experts (RMoE). RMoE leverages a Gated Recurrent Unit (GRU) to establish dependencies between routing decisions across consecutive layers. Such layerwise recurrence can be efficiently parallelly computed for input tokens and introduces negotiable costs. Our extensive empirical evaluations demonstrate that RMoE-based language models consistently outperform a spectrum of baseline models. Furthermore, RMoE integrates a novel computation stage orthogonal to existing methods, allowing seamless compatibility with other MoE architectures. Our analyses attribute RMoE's gains to its effective cross-layer information sharing, which also improves expert selection and diversity. Our code is at https://github.com/qiuzh20/RMoE