エキスパートの混合（Mixture-of-Experts）のための階層型リカレントルータ

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モデルはパラメータ効率が低いことが多い。例えば、520億パラメータを持つ事前学習済みのMoEベースのLLMは、67億パラメータの標準モデルと同等の性能を示すことがある。MoEの重要な要素である現在のルーターは、異なる層で独立してトークンを割り当て、過去のルーティング情報を活用しないため、最適でないトークンとエキスパートの組み合わせが生じ、パラメータ効率の問題を引き起こす可能性がある。この問題を緩和するため、我々はLayerwise Recurrent Router for Mixture-of-Experts（RMoE）を提案する。RMoEは、Gated Recurrent Unit（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