Roteador Recorrente Camada a Camada para Mistura de Especialistas

Resumo

A escalabilidade de grandes modelos de linguagem (LLMs) revolucionou suas capacidades em várias tarefas, no entanto, esse crescimento deve ser acompanhado por estratégias computacionais eficientes. A arquitetura Mixture-of-Experts (MoE) destaca-se por sua capacidade de aumentar o tamanho do modelo sem aumentar significativamente os custos de treinamento. Apesar de suas vantagens, os modelos MoE atuais frequentemente exibem ineficiência de parâmetros. Por exemplo, um LLM baseado em MoE pré-treinado com 52 bilhões de parâmetros pode ter desempenho comparável a um modelo padrão com 6,7 bilhões de parâmetros. Sendo uma parte crucial do MoE, os roteadores atuais em diferentes camadas atribuem tokens de forma independente sem aproveitar informações de roteamento histórico, potencialmente resultando em combinações subótimas de token-especialista e no problema de ineficiência de parâmetros. Para mitigar esse problema, introduzimos o Roteador Recorrente Camada a Camada para Mixture-of-Experts (RMoE). O RMoE utiliza uma Unidade Recorrente com Portas (GRU) para estabelecer dependências entre decisões de roteamento em camadas consecutivas. Essa recorrência camada a camada pode ser calculada de forma eficiente e paralela para tokens de entrada e introduz custos negociáveis. Nossas extensas avaliações empíricas demonstram que os modelos de linguagem baseados em RMoE superam consistentemente uma variedade de modelos de referência. Além disso, o RMoE integra uma etapa de computação inovadora, ortogonal aos métodos existentes, permitindo compatibilidade perfeita com outras arquiteturas MoE. Nossas análises atribuem os ganhos do RMoE ao seu eficaz compartilhamento de informações entre camadas, o que também melhora a seleção e diversidade de especialistas. Nosso código está em 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