Послойный рекуррентный маршрутизатор для смеси экспертов.

Layerwise Recurrent Router for Mixture-of-Experts

August 13, 2024

Авторы: Zihan Qiu, Zeyu Huang, Shuang Cheng, Yizhi Zhou, Zili Wang, Ivan Titov, Jie Fu

cs.AI

Аннотация

Масштабирование крупных языковых моделей (LLM) революционизировало их возможности в различных задачах, однако этот рост должен быть сопровожден эффективными вычислительными стратегиями. Архитектура Mixture-of-Experts (MoE) выделяется своей способностью масштабировать размер модели без значительного увеличения затрат на обучение. Несмотря на их преимущества, текущие модели MoE часто демонстрируют неэффективность параметров. Например, предварительно обученная LLM на основе MoE с 52 миллиардами параметров может выполнять сравнимо с моделью стандартного типа с 6,7 миллиарда параметров. Будучи важной частью MoE, текущие маршрутизаторы в различных слоях независимо назначают токены, не используя историческую информацию о маршрутизации, что потенциально приводит к субоптимальным комбинациям токенов и экспертов, а также проблеме неэффективности параметров. Для устранения этой проблемы мы представляем Layerwise Recurrent Router for Mixture-of-Experts (RMoE). RMoE использует воротной рекуррентный блок (GRU) для установления зависимостей между решениями по маршрутизации на последовательных слоях. Такая рекуррентность по слоям может быть эффективно параллельно вычислена для входных токенов и вносит незначительные затраты. Наши обширные эмпирические оценки показывают, что языковые модели на основе RMoE последовательно превосходят спектр базовых моделей. Более того, RMoE интегрирует новый этап вычислений, ортогональный существующим методам, обеспечивая безупречную совместимость с другими архитектурами MoE. Наши анализы приписывают успехи RMoE его эффективному обмену информацией между слоями, что также улучшает выбор и разнообразие экспертов. Наш код доступен по ссылке 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