Capacity-Aware Mixture Law Enables Efficient LLM Data Optimization

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\%.

Large language models learn from many kinds of text, such as web pages, code, math problems, and educational material. Choosing how much of each kind to use is important, because the wrong data mix can waste computing power and lead to weaker models. The hard part is that a mix that works well for a small model may not work well for a much larger one, and testing many mixes on the large model is very expensive. We introduce CAMEL, a method that predicts how a data mix will affect both validation loss and benchmark performance as model size changes. CAMEL views training as a capacity allocation problem: as models grow, they do not improve equally on every type of data, so the best data mix should also change. We also study how to spend a small experiment budget, finding that testing more mixes at the smallest and largest model sizes gives better predictions than spreading tests evenly. On mixture-of-experts language models, CAMEL finds strong data mixes for a 55B-parameter target model while using less than half of the compute required by prior methods. This can help researchers train better language models with fewer trial runs and lower computational cost.