Capacity-Aware Mixture Law Enables Efficient LLM Data Optimization

TL;DR

This paper introduces a compute-efficient pipeline for optimizing data mixtures in large language models, leveraging a capacity-aware law to improve performance prediction and reduce costs.

Contribution

It proposes CAMEL, a capacity-aware mixture law, and a loss-to-benchmark prediction law, enabling efficient data mixture scaling and performance extrapolation for large models.

Findings

Reduced mixture optimization costs by 50%.

Improved downstream benchmark performance by up to 3%.

Validated the method on models up to 55B parameters.

Abstract

A data mixture refers to how different data sources are combined to train large language models, and selecting an effective mixture is crucial for optimal downstream performance. Existing methods either conduct costly searches directly on the target model or rely on mixture scaling laws that fail to extrapolate well to large model sizes. We address these limitations by introducing a compute-efficient pipeline for data mixture scaling. First, we propose CAMEL, a capacity-aware mixture law that models validation loss with the nonlinear interplay between model size and mixture. We also introduce a loss-to-benchmark prediction law that estimates benchmark accuracy from validation loss, enabling end-to-end performance prediction for the target model. Next, we study how to allocate a fixed compute budget across model scales to fit the law and reduce prediction error. Finally, we apply our method to Mixture-of-Experts models with up to 7B-A150M parameters to fit the law, and verify the optimal mixture derived from the law by extrapolating to a 55B-A1.2B target model. Compared to prior methods, we reduce mixture optimization costs by 50\% and improves downstream benchmark performance by up to 3\%.