Nexus: Specialization meets Adaptability for Efficiently Training Mixture of Experts

TL;DR

Nexus is an enhanced Mixture of Experts architecture that enables efficient upcycling of dense models into MoEs with adaptive routing, allowing easy addition of new experts and improved specialization for diverse tasks.

Contribution

The paper introduces Nexus, a novel MoE architecture with adaptive routing that facilitates flexible expert addition and better specialization without extensive retraining.

Findings

Nexus achieves up to 2.1% improvement in initial upcycling performance.

Nexus attains an 18.8% relative gain when extending with new experts.

Flexible expert addition enhances adaptability and specialization in MoE models.

Abstract

Efficiency, specialization, and adaptability to new data distributions are qualities that are hard to combine in current Large Language Models. The Mixture of Experts (MoE) architecture has been the focus of significant research because its inherent conditional computation enables such desirable properties. In this work, we focus on "upcycling" dense expert models into an MoE, aiming to improve specialization while also adding the ability to adapt to new tasks easily. We introduce Nexus, an enhanced MoE architecture with adaptive routing where the model learns to project expert embeddings from domain representations. This approach allows Nexus to flexibly add new experts after the initial upcycling through separately trained dense models, without requiring large-scale MoE training for unseen data domains. Our experiments show that Nexus achieves a relative gain of up to 2.1% over the baseline for initial upcycling, and a 18.8% relative gain for extending the MoE with a new expert by using limited finetuning data. This flexibility of Nexus is crucial to enable an open-source ecosystem where every user continuously assembles their own MoE-mix according to their needs.