Schichtweiser rekurrenter Router für Mischung von Experten

Layerwise Recurrent Router for Mixture-of-Experts

August 13, 2024

Autoren: Zihan Qiu, Zeyu Huang, Shuang Cheng, Yizhi Zhou, Zili Wang, Ivan Titov, Jie Fu

cs.AI

Zusammenfassung

Die Skalierung großer Sprachmodelle (LLMs) hat ihre Fähigkeiten in verschiedenen Aufgaben revolutioniert, doch dieses Wachstum muss mit effizienten Rechenstrategien einhergehen. Die Mixture-of-Experts (MoE) Architektur zeichnet sich durch ihre Fähigkeit aus, die Modellgröße zu skalieren, ohne die Trainingskosten signifikant zu erhöhen. Trotz ihrer Vorteile zeigen aktuelle MoE-Modelle oft eine Parameterineffizienz. Beispielsweise könnte ein vorab trainiertes MoE-basiertes LLM mit 52 Milliarden Parametern vergleichbare Leistungen erbringen wie ein Standardmodell mit 6,7 Milliarden Parametern. Als entscheidender Bestandteil von MoE weisen aktuelle Router in verschiedenen Schichten eigenständig Tokens zu, ohne historische Routing-Informationen zu nutzen, was potenziell zu suboptimalen Token-Expert-Kombinationen und dem Parameterineffizienzproblem führt. Um dieses Problem zu mildern, führen wir den Schichtweisen Rekurrenten Router für Mixture-of-Experts (RMoE) ein. RMoE nutzt eine Gated Rekurrente Einheit (GRU), um Abhängigkeiten zwischen Routing-Entscheidungen über aufeinanderfolgende Schichten herzustellen. Eine solche schichtweise Rekurrenz kann effizient parallel für Eingabetokens berechnet werden und führt zu verhandelbaren Kosten. Unsere umfangreichen empirischen Bewertungen zeigen, dass auf RMoE basierende Sprachmodelle konsistent eine Vielzahl von Basislinienmodellen übertreffen. Darüber hinaus integriert RMoE eine neue Rechenstufe, die orthogonal zu bestehenden Methoden ist und eine nahtlose Kompatibilität mit anderen MoE-Architekturen ermöglicht. Unsere Analysen führen die Gewinne von RMoE auf den effektiven Informationsaustausch zwischen Schichten zurück, der auch die Expertenauswahl und -vielfalt verbessert. Unser Code ist unter https://github.com/qiuzh20/RMoE verfügbar.

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