Laaggewijze Recurrent Router voor Mixture-of-Experts

Samenvatting

De schaalvergroting van grote taalmodellen (LLMs) heeft hun mogelijkheden in diverse taken revolutionair verbeterd, maar deze groei moet gepaard gaan met efficiënte computationele strategieën. De Mixture-of-Experts (MoE)-architectuur springt eruit vanwege het vermogen om de modelgrootte te schalen zonder de trainingskosten significant te verhogen. Ondanks hun voordelen vertonen huidige MoE-modellen vaak parameterinefficiëntie. Zo kan een vooraf getraind MoE-gebaseerd LLM met 52 miljard parameters vergelijkbaar presteren als een standaardmodel met 6,7 miljard parameters. Als een cruciaal onderdeel van MoE, wijzen huidige routers in verschillende lagen tokens onafhankelijk toe zonder gebruik te maken van historische routeringsinformatie, wat mogelijk leidt tot suboptimale token-expertcombinaties en het parameterinefficiëntieprobleem. Om dit probleem te verlichten, introduceren we de Layerwise Recurrent Router for Mixture-of-Experts (RMoE). RMoE maakt gebruik van een Gated Recurrent Unit (GRU) om afhankelijkheden tussen routeringsbeslissingen over opeenvolgende lagen te creëren. Een dergelijke laaggewijze recurrentie kan efficiënt parallel worden berekend voor invoertokens en introduceert verhandelbare kosten. Onze uitgebreide empirische evaluaties tonen aan dat RMoE-gebaseerde taalmodellen consistent beter presteren dan een reeks basismodellen. Bovendien integreert RMoE een nieuwe rekentrap die orthogonaal is aan bestaande methoden, waardoor naadloze compatibiliteit met andere MoE-architecturen mogelijk is. Onze analyses schrijven de winsten van RMoE toe aan de effectieve gelaagde informatie-uitwisseling, die ook de expertselectie en diversiteit verbetert. Onze code is te vinden op 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