Mixture-of-Skills: Learning to Optimize Data Usage for Fine-Tuning Large Language Models (2024)

Abstract: Large language models (LLMs) are typically fine-tuned on diverse and extensive datasets sourced from various origins to develop a comprehensive range of skills, such as writing, reasoning, chatting, coding, and more. Each skill has unique characteristics, and these datasets are often heterogeneous and imbalanced, making the fine-tuning process highly challenging. Balancing the development of each skill while ensuring the model maintains its overall performance requires sophisticated techniques and careful dataset curation. In this work, we propose a general, model-agnostic, reinforcement learning framework, Mixture-of-Skills (MoS), that learns to optimize data usage automatically during the fine-tuning process. This framework ensures the optimal comprehensive skill development of LLMs by dynamically adjusting the focus on different datasets based on their current learning state. To validate the effectiveness of MoS, we conduct extensive experiments using three diverse LLM backbones on two widely used benchmarks and demonstrate that MoS substantially enhances model performance. Building on the success of MoS, we propose MoSpec, an adaptation for task-specific fine-tuning, which harnesses the utilities of various datasets for a specific purpose. Our work underlines the significance of dataset rebalancing and present MoS as a powerful, general solution for optimizing data usage in the fine-tuning of LLMs for various purposes.

## introduction large language models (llms) have demonstrated their extraordinary capabilities and are expected to proficiently master a diverse range of skills @xcite @xcite @xcite @xcite , such as writing, reasoning, chatting, coding, and more, through supervised fine-tuning (sft) and reinforcement learning from human feedback (rlhf) on an extensive collection of datasets from various sources @xcite @xcite . each dataset contributes unique elements to the model's skill set, but this diversity also brings challenges. one common challenge in fine-tuning models across multiple datasets is dealing with their heterogeneity (where different datasets exhibit distinct characteristics) and imbalance (where datasets vary in size and accessibility), making the fine-tuning process highly challenging. to address this challenge, recent approaches often cap dataset usage to prevent models from being overwhelmed by large datasets, but this limits the utilization of all available data @xcite @xcite . previous multilingual research adjusts dataset distribution heuristically with a temperature term τ , which requires extensive hyperparameter tuning and overlooks dataset interactions and model learning dynamics @xcite . this leads us to a critical research question: is there a better way to optimize the data usage? building on the research question posed, we first confirm that adjusting the dataset usage properly can significantly enhance model performance (see to validate the effectiveness of mos, we conduct extensive experiments using three diverse model backbones: qwen1.5-0.5b @xcite , gemma-2b @xcite , and llama-3-8b (ai@meta, 2024), on two widelyused benchmarks: mmlu @xcite and mt-bench @xcite . our empirical results demonstrate that mos substantially improves the models' overall performance. our analysis indicates that our model not only effectively learns optimal data utilization but also accelerates training convergence by 2.2×. additionally, it demonstrates robustness against variations in sampling priors and integrates seamlessly with advanced instance selection methods. furthermore, we explore the application of mos in task-specific fine-tuning. we show that mos, with minor reward modifications known as mospec, can be effectively used for task-specific fine-tuning. our contributions are summarized as follows: heuristic balancing by temperature instead of merging all datasets into a single training mixture, a common approach is to adjust the sampling probability of texts in different languages using a temperature term τ @xcite . specifically, the sampling probability of the i-th dataset is q(i) consequently, τ = 1 corresponds to proportional sampling, equivalent to equation 2. conversely, τ = ∞ corresponds to uniform sampling, where smaller datasets are up-sampled to match the largest dataset. the loss function becomes: differentiable data selection (dds) @xcite propose a general framework that automatically re-weighs training instances to enhance model performance, utilizing a validation set d dev . this framework consists of two components: the model θ θ θ and the scorer network ψ ψ ψ. the scorer network ψ ψ ψ is designed to calculate a sampling probability for each training instance, which reflects its impact on validation performance. training instances that have a greater similarity with d dev are allocated a higher probability, thus increasing their likelihood of being selected for model θ θ θ updates. mos is inspired by dds but has key differences. firstly, mos focuses on rebalancing datasets, unlike dds, which reweighs training instances. secondly, mos does not require prior knowledge of downstream applications, whereas dds relies on validation set feedback, risking overfitting to that specific validation set. thirdly, dds uses the same architecture for the scorer network ψ ψ ψ and model θ θ θ, limiting its scalability, while mos uses a simple mlp model as its scorer network (see section 3.1). ## mixture-of-skills in this section, we provide a detailed overview of mixture-of-skills (mos). we begin by outlining the reinforcement learning framework employed in mos (section 3.1). following this, we discuss the reward functions (section 3.2). ## experiments we present our experimental setup (section 4.1) and main results (section 4.2) in this section. ## analysis in this section, we conduct an in-depth analysis of mos using llama-3-8b. our analysis encompasses the learned dataset distribution of mos, the training convergence speed of mos, the impact of sampling priors on mos, and its compatibility with the instance selection methods. ## conclusion in this work, we address the critical challenge of optimizing data usage during the fine-tuning process of llms. we propose a general, model-agnostic reinforcement learning framework, mixture-of-skills (mos), that dynamically adjusts dataset usage to enhance model performance with three novel rewards. through extensive experiments on three diverse model backbones and two widelyused benchmarks, we demonstrate that mos significantly improves overall model performance. additionally, we explore the application of mos in taskspecific fine-tuning, leading to the development of mospec. our experiments show that models finetuned with mospec on various datasets outperform those trained solely on task-specific datasets. in summary, mos provides a powerful and flexible solution to the challenges of dataset heterogeneity and imbalance in the fine-tuning of llms.