Source-linked AI summary
Pushing Mixture of Experts to the Limit: Extremely Parameter Efficient MoE for Instruction Tuning
Ted Zadouri, Ahmet Üstün, Arash Ahmadian, Beyza Ermiş, Acyr Locatelli, Sara Hooker
TL;DR
Conventional MoEs are difficult to fine-tune at scale because their many experts impose substantial parameter and computational demands. The paper combines MoE routing with lightweight PEFT experts, and reports stronger performance than standard PEFT methods while approaching or matching full fine-tuning with less than 1% of model parameters updated. The evaluation supports this behavior on unseen tasks, while remaining limited to text-to-text models and fine-tuning.
Problem
Applying MoEs to instruction fine-tuning is challenging because fully fine-tuning their many parameters is computationally demanding, motivating a PEFT-based alternative.
Method
The paper introduces Mixture of Vectors and Mixture of LoRA, which use lightweight PEFT adapters as experts in an MoE framework.
Results
The proposed methods outperform standard PEFT techniques on unseen tasks and are on par with full fine-tuning while updating less than 1% of model parameters.
Takeaways & Limitations
Parameter-efficient MoEs can retain strong instruction-tuning performance across model scales while updating only a tiny fraction of the underlying model.
Takeaways & Limitations
The evaluation focuses on text-to-text T5-style models and fine-tuning, leaving decoder-only models and pre-training unexplored.
Abstract
from arXiv · showhide
The Mixture of Experts (MoE) is a widely known neural architecture where an ensemble of specialized sub-models optimizes overall performance with a constant computational cost. However, conventional MoEs pose challenges at scale due to the need to store all experts in memory. In this paper, we push MoE to the limit. We propose extremely parameter-efficient MoE by uniquely combining MoE architecture with lightweight experts.Our MoE architecture outperforms standard parameter-efficient fine-tuning (PEFT) methods and is on par with full fine-tuning by only updating the lightweight experts -- less than 1% of an 11B parameters model. Furthermore, our method generalizes to unseen tasks as it does not depend on any prior task knowledge. Our research underscores the versatility of the mixture of experts architecture, showcasing its ability to deliver robust performance even when subjected to rigorous parameter constraints. Our code used in all the experiments is publicly available here: https://github.com/for-ai/parameter-efficient-moe.
1 Introduction
MoE offers conditional computation but conventional architectures are costly to fine-tune at modern language-model scales. This work adapts MoE with lightweight PEFT experts, achieving strong unseen-task performance while updating very few parameters.
- Motivation: Conventional MoEs require updating and storing many expert parameters, making full fine-tuning computationally infeasible for most practitioners at modern LLM scales.The paper frames this as the central challenge for applying MoEs to instruction fine-tuning.
- Motivation: The paper asks whether MoEs can be combined with PEFT methods such as (IA)3 and LORA in a more realistic, parameter-limited setting.This setting also inherits optimization challenges from MoE training while restricting the number of updated parameters.
- Approach: The proposed framework introduces Mixture of Vectors and Mixture of LoRA, using modular lightweight experts to extend MoE to severely constrained environments.These methods adapt the MoE architecture rather than fully fine-tuning all model parameters.
- Evaluation: Across T5 models from 770M to 11B parameters and 12 tasks from 55 datasets, the experiments report consistent results for the proposed methods.The evaluation spans multiple model scales and instruction-tuning tasks.
- Results: The proposed instruction-tuning methods consistently outperform traditional parameter-efficient methods on unseen tasks across model sizes and expert configurations.MoV improves over standard (IA)3 by up to 14.57% at 3B and 8.39% at 11B.
- Results: MoV achieves highly competitive performance with full fine-tuning while updating only 0.32% of a 3B model and 0.86% of an 11B model.The reported comparison covers 8 unseen tasks and uses substantially fewer updated parameters.
2 Methodology
The framework combines MoE routing with lightweight PEFT adapters, training only experts and routers while keeping the dense backbone frozen. MoV uses (IA)3 vectors, whereas MoLORA uses LORA adapters, with soft merging enabling additional memory benefits.
- Framework: The instruction-tuning setup separates training tasks from held-out evaluation tasks and replaces full-model updates with parameter-efficient mixture-of-experts adaptation.The pretrained model is fine-tuned on T_train and evaluated zero-shot on unseen T_eval tasks.
- Framework: PEFT methods restrict updates to a limited number of parameters to reduce the computational burden of full LLM fine-tuning.The paper evaluates the approach with both (IA)3 vectors and LORA adapters.
- Mixture-of-Experts design: MoE layers use a router and multiple experts, with softmax gating scores combining expert outputs for each input representation.The router is a trainable dense layer followed by softmax, and the experts are parameterized functions.
- Mixture-of-Experts design: The proposed architecture replaces full dense experts with lightweight (IA)3 or LORA adapters while freezing pretrained dense weights and training experts and routers.The resulting variants are called Mixture of Vectors (MoV) and Mixture of LORA (MoLORA).
- Routing and efficiency: Soft merging averages linear PEFT experts before applying the resulting combined transformation, and Figure 2 illustrates this operation for MoV.Only a small fraction of parameters is updated through the vectors and router in each multi-head attention block.
- Routing and efficiency: Freezing most parameters reduces gradient-computation and optimizer-state memory during training, while storing only one backbone copy improves inference memory efficiency.Traditional MoE models require many copies of full feed-forward blocks or model replicas, whereas these methods retain a single backbone copy.
3 Experiments
The experiments evaluate parameter-efficient MoE variants across T5 model sizes, instruction-tuning tasks, baselines, and routing ablations. They measure zero-shot performance on held-out datasets and compare soft merging with discrete top-k routing.
- Experimental setup: The experiments use P3, a collection of 62 datasets covering a wide variety of instruction-tuning tasks.Tasks are converted into prompt templates following the T5 procedure.
- Experimental setup: The study evaluates T5 models ranging from 770M to 11B parameters using Adafactor, sequence lengths of 1024 inputs and 256 targets, and 500K training steps.All parameter-efficient MoE variants use batch size 32.
- Experimental setup: MoV and MoLORA are compared with fully fine-tuned T0 and standard (IA)3 and LORA baselines.The comparison includes both the full-fine-tuning baseline and standard parameter-efficient methods.
- Evaluation: Zero-shot evaluation uses eight held-out datasets spanning natural-language inference, multiple-choice question answering, story completion, coreference resolution, and word-sense disambiguation.The paper reports median accuracy per dataset and an average across datasets.
- Ablations: The study runs ablations on routing inputs, comparing token embeddings with externally computed sentence embeddings for their effects on routing and downstream generalization.Sentence embeddings are produced with Sentence-T5, while the main methods use intermediate token embeddings.
- Ablations: Soft merging averages all experts within a routing block, whereas discrete top-k routing selects only the k experts chosen by the router.The experiments evaluate k = {1, 2} and also test load balancing with an auxiliary loss.
4 Results and Discussion
Across model sizes and routing choices, parameter-efficient MoE variants improve on standard PEFT methods while approaching or matching full fine-tuning with very small update budgets. Soft merging and larger expert mixtures generally provide the strongest unseen-task performance, while routing specialization appears on both seen and unseen tasks.
- MoV and MoLORA significantly outperform standard (IA)3 and LORA on zero-shot unseen-task performance at the T5-3B scale.
- 14.57% improvement over dense (IA)3 is achieved by MoV with 30 experts, adding only 0.018% parameters per expert.
- MoV-10 and MoLORA-10 are on par with full fine-tuning, with MoV-10 updating only 0.32% of model parameters.
- At 3B and 11B scales, MoV-60 is approximately on par with full fine-tuning while updating less than 1.3% of parameters.
- Increasing experts generally improves unseen-task performance: MoV-60 reaches 53.63 at 770M and 64.08 at 11B, versus 52.47 and 62.3 with 10 experts.
- Soft merging all 10 experts achieves 59.93 average median accuracy, exceeding top-1 and top-2 routing at 54.92 and 57.45.
- Expert routing differs across tasks: seen and unseen tasks activate experts at different magnitudes and orderings, indicating task-dependent specialization.
- A learning rate of 3e−4 yields higher performance than larger tested rates by stabilizing parameter-efficient expert training and preserving expert diversity.
5 Related Work
Related work frames this paper at the intersection of MoE routing, instruction tuning, and parameter-efficient adaptation. The paper distinguishes its approach by combining lightweight PEFT experts with multi-task instruction tuning and scaling across larger task and data collections.
- Mixture-of-Experts: MoE increases model capacity by selectively activating components while keeping computation near that of a dense model.
- Mixture-of-Experts: Prior MoE research studies routing methods ranging from random and weighted averaging to sparse selection of one or k experts.
- Instruction Tuning: Instruction tuning fine-tunes language models on paired prompts and responses across tasks to improve zero-shot performance on unseen tasks.
- Instruction Tuning: The paper applies instruction fine-tuning to diverse inputs while targeting computational efficiency without compromising zero-shot performance.
- Parameter-Efficient Fine-tuning: PEFT methods include adapters, soft prompts, bias-only updates, LORA low-rank matrices, and (IA)3 task-specific vectors.
- Parameter-Efficient Fine-tuning: The authors choose (IA)3 and LORA as PEFT components because they balance performance and parameter efficiency.
- Parameter-Efficient Fine-tuning: Earlier MoE-PEFT studies include task-specific adapter mixtures and deterministic routing, whereas this work scales data and task volume to address instability.
6 Conclusion
The framework introduces MoV and MoLORA, combining MoE architectures with lightweight parameter-efficient experts for computationally limited instruction tuning. Across unseen tasks, it matches full fine-tuning and outperforms parameter-efficient baselines while updating less than 1% of model parameters, but evaluation is limited to text-to-text models and fine-tuning.
- MoV and MoLORA introduce lightweight mixture-of-experts methods for instruction tuning in extremely computationally limited settings.The methods combine MoE architectures with parameter-efficient experts.
- The framework outperforms parameter-efficient techniques and remains compatible with quantization and other efficiency-enhancing strategies.The paper reports rigorous ablations across model sizes, token versus embedding representations, and routing mechanisms.
- The evaluation focuses on text-to-text T5 models and does not extend to decoder-only GPT-style models.The authors leave decoder-only evaluation for future work.
A.1 Zero-Shot Evaluation for P3 dataset
The appendix evaluates proposed methods across T5 model sizes and unseen P3 tasks, including comparisons of expert counts at 770M, 3B, and 11B scales. Results are reported using mean and median scores for each evaluation set.
- The evaluation compares MoV and MoLoRA across 770M, 3B, and 11B T5 models on unseen P3 tasks.The 770M and 3B evaluations vary the number of experts, while the 3B setup also experiments with top-k routing.
- Mean and median scores are reported for every evaluation set derived from the P3 dataset.P3 covers a range of tasks used for zero-shot evaluation.
- The 11B evaluation demonstrates the efficacy of mixture-of-PEFT experts at the largest available T5 checkpoint.The largest evaluated checkpoint is T5-XXL with 11B parameters.
A.2 Token vs. Sentence Embeddings for Routing
The routing comparison evaluates whether routers receive tokens directly or sentence embeddings. For T5-XL, token routing consistently performs better on both mean and median measures, and the table frames token routing as preferable to sentence-embedding routing across model parameters.
- Token routing consistently outperforms sentence-embedding routing for T5-XL on both mean and median performance.The ANLI dataset is excluded from the embedding dataset.
- The comparison examines token routing against the stronger inductive bias imposed by sentence embeddings across model parameters.