Dynamic Data Mixing Maximizes Instruction Tuning for Mixture-of-Experts (2406.11256v1)

Abstract: Mixture-of-Experts (MoE) models have shown remarkable capability in instruction tuning, especially when the number of tasks scales. However, previous methods simply merge all training tasks (e.g. creative writing, coding, and mathematics) and apply fixed sampling weights, without considering the importance of different tasks as the model training state changes. In this way, the most helpful data cannot be effectively distinguished, leading to suboptimal model performance. To reduce the potential redundancies of datasets, we make the first attempt and propose a novel dynamic data mixture for MoE instruction tuning. Specifically, inspired by MoE's token routing preference, we build dataset-level representations and then capture the subtle differences among datasets. Finally, we propose to dynamically adjust the sampling weight of datasets by their inter-redundancies, thus maximizing global performance under a limited training budget. The experimental results on two MoE models demonstrate the effectiveness of our approach on both downstream knowledge & reasoning tasks and open-ended queries. Code and models are available at https://github.com/Spico197/MoE-SFT .

An Analysis of Dynamic Data Mixing in Mixture-of-Experts Models

The paper addressed the challenges associated with instruction tuning in Mixture-of-Experts (MoE) models, particularly as the number of tasks increases. The primary contribution of this research is the introduction of a novel dynamic data mixing strategy that optimizes the process of instruction tuning by dynamically adjusting the sampling weights of datasets based on inter-redundancies.

Overview

MoE models, by design, incorporate multiple experts, enabling them to scale effectively with diverse tasks. Traditionally, these models apply a static data mixing strategy that overlooks the variance in the contribution of different datasets as the model evolves during training. This can lead to inefficiencies and fails to capitalize on the full potential of MoE architectures. The authors propose a dynamic data mixture strategy to address these shortcomings, one that is responsive to the task-specific needs as training progresses.

Methodology

The proposed method integrates the concept of token routing preference inherent in MoE architectures. Specifically, the authors leverage gate loads—metrics indicating the routing of tokens among experts—to derive dataset-level representations. These representations facilitate the calculation of L2 distances between datasets, effectively quantifying their inter-redundancies relative to the model’s current state. By incorporating these redundancies, the sampling weights are adjusted dynamically to enhance the alignment abilities of the MoE models.

Experimental Evaluation

The paper reports experiments conducted using two MoE models: MoLM 700M-4E and LLaMA-MoE 3.5B-2E. The models were fine-tuned using a combination of four diverse instruction datasets, each aiming to represent different domain challenges such as open-ended dialogue, task-oriented instructions, mathematical problem solving, and code generation tasks. The dynamic sampling strategy demonstrated superior performance across a variety of downstream knowledge and reasoning tasks as well as open-ended question answering scenarios. In particular, the evaluation on benchmarks including MMLU, BBH, GSM8K, MBPP, and MT-Bench confirmed the efficacy of the dynamic method.

Numerical Outcomes

One of the significant findings is the improvement in performance metrics across tasks, particularly evident on datasets like GSM8K and MBPP, which are heavily task-aligned (math and code-related, specifically). For instance, the dynamic sampling weights led to GSM8K performance of 11.90% on LLaMA-MoE, surpassing other methods significantly.

Implications and Future Directions

The implications of this research are noteworthy for the development of adaptive instruction tuning strategies in MoE architectures. This dynamic approach not only optimizes resource allocation during training but also potentially reduces the computational redundancy. The method’s reliance on internal state characteristics of the model, such as gate loads, signals a shift towards more nuanced, self-regulating training algorithms that could vastly improve efficiency and scalability.

Looking ahead, this approach could inspire further exploration into adaptive mechanisms that transcend traditional instruction tuning boundaries. Integrating additional model signals beyond gate loads and applying this dynamic strategy to larger, more complex models could yield even more robust performance gains. Moreover, expanding upon this work could involve exploring proxy models or alternative structures to estimate redundancies, thus refining the dynamic data mixing protocol.

In conclusion, the dynamic data mixture strategy presented in this paper provides a compelling pathway for advancing instruction tuning in MoE models, aligning with the broader trajectory of enhancing model efficiency and performance through adaptive learning mechanisms.