BiMix: Bivariate Data Mixing Law for Language Model Pretraining (2405.14908v3)

Abstract: LLMs have demonstrated remarkable capabilities across various tasks, primarily attributed to the utilization of diversely sourced data. However, the impact of pretraining data composition on model performance remains poorly understood. This paper introduces $\textbf{BiMix}$, a novel bivariate data mixing law that models the joint scaling behavior of domain proportions and data volume in LLM pretraining. $\textbf{BiMix}$ provides a systematic framework for understanding and optimizing data mixtures across diverse domains. Through extensive experiments on two large-scale datasets, we demonstrate $\textbf{BiMix}$'s high accuracy in loss extrapolation (mean relative error < 0.2%) and its generalization to unseen mixtures (R${}^{2}$ > 0.97). Optimization of domain proportions yields superior model performance compared to existing methods. Furthermore, we establish entropy-based measures as efficient proxies for data mixing, offering a computationally lightweight strategy. Our work contributes both theoretical insights into data mixing dynamics and practical tools for enhancing LLM training efficiency, paving the way for more effective scaling strategies in LLM development.

• The paper introduces a bivariate scaling law that models how data quantity and mixing proportions affect LLM performance.
• It employs entropy-driven proxies, including conditional entropy, to efficiently predict validation loss and reduce computational costs.
• Empirical results on diverse datasets like The Pile and SlimPajama validate the method with high R² and PCC, enabling optimized data curation.

A Unified Scaling Law for Efficient Data Mixing in LLMs

Introduction

The paper under review introduces a novel approach to the integration of diverse datasets in the training of LLMs. The motivation stems from the need to overcome the inefficiencies associated with traditional heuristic-based methods for data mixing, which tend to be sub-optimal and resource-intensive. Instead, the paper proposes a principled strategy utilizing a unified scaling law to model the effects of data quantity and mixing proportions on model performance. This framework is designed to optimize data curation, thereby enhancing training efficiency.

Unified Scaling Law

The core contribution of the paper is a new bivariate scaling law aimed at accurately quantifying the impact of two critical factors: data quantity and domain mixing proportions. The scaling law, denoted as $\mathcal{L}(s, r)$, is mathematically formulated to predict validation loss based on the number of training steps $s$ and the domain proportion $r$. The law is empirically validated through extensive experiments, demonstrating its predictive power and practical applicability.

Entropy-Driven Proxies

To efficiently determine the coefficients of the scaling law, the paper leverages entropy-driven proxies. Specifically, the conditional entropy of the dataset is utilized to estimate the importance of various data domains. This approach allows for the creation of optimized data mixtures without prior training, significantly reducing computational costs.

Experimental Setup and Results

The experiments are conducted using two prominent domain-diverse datasets: The Pile and SlimPajama. Three types of data mixtures are evaluated: Baseline, DoReMi, and entropy-driven proxies (including Conditional Entropy, Shannon Entropy, Joint Entropy, and Von Neumann Entropy). The models are assessed on multiple one-shot generative tasks and perplexity metrics.

The results indicate that the conditional entropy-driven data mixtures not only achieve comparable or superior performance to the DoReMi method but do so with greater efficiency. The proposed scaling law is shown to fit the empirical data well, with high R\textsuperscript{2} and Pearson Correlation Coefficient (PCC) values across most domains.

Practical Implications

The scaling law provides several practical benefits:

1. Mixture Selection: The law can predict validation losses for new data mixtures, enabling informed decisions about data selection without expensive preliminary experimentation.
2. Proportion Optimization: The law can be extended to directly optimize domain proportions, tailoring data mixtures to specific training scenarios and requirements.

Implications for Future Research

The introduction of a predictive scaling law for data mixing has significant implications for the development of LLMs. The ability to optimize data use efficiently is crucial as models grow in size and complexity. Future research could focus on refining the entropy proxies and further validating the scaling law across different model architectures and datasets. Additional exploration into the applicability of the scaling law in various AI domains beyond natural language processing is also a promising direction.

Conclusion

The paper makes a substantial contribution by proposing a unified scaling law that effectively models the relationship between data quantity, mixing proportions, and model performance. The use of entropy-driven proxies for efficient data mixing further enhances the practical utility of this approach. The findings offer a balanced trade-off between computational effort and model quality, setting a foundation for more cost-effective LLM development strategies.

This research highlights the potential of principled approaches to data integration in advancing the capabilities of LLMs, paving the way for future innovations in artificial intelligence.