TLDR: This paper introduces Data Mixing Optimization, a novel method for supervised fine-tuning (SFT) of Large Language Models (LLMs). It frames data mixing as an optimization problem to minimize validation loss by modeling effective data transfer and leveraging scaling laws. The method determines optimal, scale-dependent data weights, ensuring balanced performance across domains. Empirical results show it achieves near-optimal performance compared to computationally expensive grid search and improves validation loss and downstream task performance for popular SFT datasets, offering practical guidance for training general-purpose and domain-specific LLMs.
Large Language Models, or LLMs, are powerful AI systems capable of handling a wide array of tasks, from following instructions to solving complex math problems. A crucial step in making these models perform well is called supervised fine-tuning (SFT), where they are trained on specific datasets to refine their abilities. A key challenge in this process is deciding how to mix different types of data – like math, code, or general instructions – a practice known as data mixing.
Traditionally, finding the best way to mix data has been a trial-and-error process, which is highly inefficient given the vast number of possible combinations. While some research has focused on data mixing during the initial pre-training phase of LLMs, the fine-tuning stage has remained largely unexplored. This is a significant oversight, as fine-tuning not only helps models adopt a certain style but also enhances their knowledge and overall capabilities.
A new research paper, titled “Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models,” introduces a novel approach to tackle this problem. Authored by Yuan Li, Zhengzhong Liu, and Eric Xing, this paper frames data mixing as an optimization challenge. Their method aims to minimize the model’s validation loss – essentially, how well the model performs on unseen data – by intelligently determining the optimal proportions for different data domains.
The core of their approach involves modeling how effectively data from one domain can transfer knowledge to another, and leveraging what are known as “scaling laws” for fine-tuning. Scaling laws describe how a model’s performance changes with the amount of data it’s trained on. By conducting experiments with small-scale data mixtures, the researchers can fit parameters to these models and then mathematically derive the ideal weights for each data type.
The paper provides both mathematical proofs and practical results demonstrating that their algorithm achieves excellent performance across various domains. In controlled experiments, models trained with their optimized data weights performed almost as well as those trained with optimal weights found through a much more computationally expensive method called grid search. On average, the per-domain loss was only slightly higher (0.66%) than the best results from grid search.
One of the interesting findings is that the optimal data mixture is “scale-dependent,” meaning the best proportions change depending on the total amount of data available for fine-tuning. This challenges previous assumptions that optimal weights might be universal. Another important insight is that their algorithm ensures “no domains are left behind.” Even if a particular data type, like math data, is less effective at reducing overall loss, it still maintains a decent proportion in the mix. This prevents any single domain from dominating the training process and helps create a general-purpose model that performs robustly across all areas.
The researchers conducted extensive experiments, including mixing instruction following, math, and code data, and re-weighting popular fine-tuning datasets like Tulu3 and Orca. For Tulu3, their method significantly improved both validation loss and performance on various downstream tasks, such as general knowledge, instruction following, and coding. However, for the Orca collection, the improvements in validation loss didn’t always translate directly to better downstream task performance. This highlights the importance of ensuring that the validation datasets used during training closely match the real-world tasks the model is expected to perform.
The method also shows promise for training specialized LLMs. For instance, in a case study for a medical chatbot, the algorithm surprisingly recommended a mix with more general instruction following data than purely medical data. This aligns with the “cocktail effect,” where including diverse, related tasks can enhance performance on a specific target task, as general instruction data helps the LLM improve its fundamental language understanding and response generation skills.
Also Read:
- Optimizing Data Mixtures for Language Models with Bayesian Approaches
- A New Approach to Harmonize Language Model Training: Dynamic Weighting for Supervised Fine-Tuning and Reinforcement Learning
This research offers valuable insights and practical guidance for the supervised fine-tuning of large language models. By providing a systematic way to optimize data mixtures, it paves the way for developing more efficient and effective general-purpose and domain-specific LLMs. You can read the full paper here.
