Nexus: Specialization meets Adaptability for Efficiently Training Mixture of Experts (2408.15901v1)

Abstract: Efficiency, specialization, and adaptability to new data distributions are qualities that are hard to combine in current LLMs. 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.

• The paper demonstrates that Nexus adapts dense, domain-specific models into a unified Mixture of Experts framework using an adaptive routing mechanism.
• It introduces an efficient upcycling process that leverages pre-trained dense experts, significantly reducing the computational overhead of full model retraining.
• Experimental results reveal up to an 18.8% performance improvement when incorporating new domain experts, underscoring the framework’s robust adaptability.

An Overview of "Nexus: Specialization meets Adaptability for Efficiently Training Mixture of Experts"

The paper "Nexus: Specialization meets Adaptability for Efficiently Training Mixture of Experts" addresses a significant yet intricate problem in the training of LLMs: combining efficiency, specialization, and adaptability. The authors propose "Nexus," a novel architecture that extends the Mixture of Experts (MoE) framework, embedding it with capabilities that streamline its adaptation to new data distributions and improve the training process's overall efficiency.

Core Contributions and Methodology

Nexus is an innovative MoE architecture focused on adaptively integrating specialized dense expert models into a flexible MoE framework. The central innovations within Nexus are encapsulated in its adaptive routing mechanism and its ability to seamlessly extend with new experts. Here are the main contributions delineated in the paper:

1. Adaptive Router Based on Domain-Specific Data: Nexus introduces a router that learns to project the embedding of each data domain into an expert embedding. This helps maintain the specialization derived from the initial domain-specific training of dense experts.
2. Flexible Upcycling of Dense Experts: By leveraging the learned domain-to-expert embedding projections, Nexus can flexibly add new experts post the initial upcycling phase. This design bypasses the need for extensive re-training of the entire MoE, significantly reducing computational overhead.
3. Enhanced Performance Metrics: The proposed architecture shows a relative gain of up to 2.1% over the baseline in the initial upcycling phase and an 18.8% relative gain when incorporating new experts. This demonstrates Nexus's superior adaptability and specialization retention.

Technical Specifics

Initial Upcycling

The paper's methodology starts with training dense models on specific domains from the SlimPajama dataset. These experts are then 'upcycled' into an MoE framework. During this phase, the original feed-forward network (FFN) layers are copied to form the MoE's experts. Additionally, averaging is used for merging non-FFN parameters across experts.

Adaptive Router Mechanism

The Nexus router, crucial to its adaptability, uses a two-layer MLP to project domain embeddings to expert embeddings. This projection mechanism allows for high specialization accuracy, as the experts are aligned closely with their initial training domains.

Efficient Extendability

Extending the MoE with new experts under Nexus involves a calculated merging of parameters and minimal finetuning tokens ensuring that new domains can be integrated swiftly without compromising the model's initial performance.

Experimental Setup and Results

The paper evaluates Nexus using two primary seed models: one with 470M parameters and another with 2.8B parameters. The experiments span three phases: training specialized expert LMs, MoE training, and the extension of the MoE model with new experts.

Key Results

• For the MoE created using the 470M parameter model, Nexus achieved a 5.8% relative gain in performance over the seed model, and 3.2% over the upcycled MoE with a linear router.
• For the 2.8B parameter model, Nexus outperformed the baseline approaches in most evaluation categories, showing substantial improvement in knowledge retrieval tasks.
• When extending the MoE with a new code expert, Nexus demonstrated an 18.8% relative gain in the new domain tasks over the baseline, with minimal performance loss in general tasks, underlying its adaptive capabilities.

Implications

The implications of Nexus are broad, particularly for practitioners aiming to develop and deploy adaptable LLMs efficiently:

• Practical Deployment: For real-world applications where model maintenance and adaptability to new data are critical, Nexus provides an efficient framework to incorporate updated domain knowledge without extensive retraining.
• Scalable Specialized Models: This methodology facilitates the creation of scalable models that maintain high specialization across diverse domains, thus offering fine-tuned performance for specific tasks.

Future Developments

Given Nexus's promising results, future work could explore:

• Broader Domain Exploration: Automatically clustering and identifying new domains for expert training could further enhance Nexus’s adaptability and performance.
• Scalability Enhancements: Investigating the performance and efficiency of Nexus in even larger model scales and more extensive data sets to validate its robustness.
• Optimization of Embedding Projections: Fine-tuning the domain-to-expert embedding projections could further improve the router's efficiency and accuracy.

Conclusion

In summary, the paper "Nexus: Specialization meets Adaptability for Efficiently Training Mixture of Experts" introduces a pivotal architecture that addresses pressing issues in MoE training. By merging domain-specific dense experts into a unified and adaptable MoE framework, Nexus presents a significant step towards creating more efficient and flexible LLMs.