Harnessing the Potential of Optimizing Data Mixtures via Bayesian Domain Reweighting

2026-07-30

AI Summary

This work addresses the challenge of optimizing domain mixture ratios in large language model pretraining, where existing methods often rely on strong structural assumptions—such as rank preservation or scaling laws—leading to unstable optimization or high computational costs. To overcome these limitations, the authors propose a Bayesian inference–based domain reweighting approach that learns a Gamma prior from observed data and infers optimal domain weights under a Dirichlet distribution, without requiring such restrictive assumptions. This method substantially reduces both data and computational requirements, enabling efficient identification of high-performing mixture strategies with significantly less training data. Empirical results demonstrate its superiority over current search-based function-fitting approaches in both efficiency and performance.

Abstract

The performance of Large Language Models (LLMs) is fundamentally influenced by the distributional composition of multi-domain pre-training data. While manual heuristics were prevalent in early models, they increasingly fail to capture the intricate synergies between domains as data complexity grows. To overcome the issue, a dominant approach seeks to fit a proxy function mapping between domain weights and their corresponding validation losses, and then find the optimal domain weights to minimize validation losses. These methods rely on strong structural assumptions, such as rank invariance or scaling laws, which are often violated, resulting in non-negligible estimation bias. A promising approach is to directly optimize the weighting scheme from data. However, it suffers from unstable optimization trajectory and prohibitive computational overhead, limiting its potential to search better domain weights configurations. This paper presents a Bayesian domain weighting method to infer the weights from a Dirichlet distribution via introducing Gamma prior information learned from observations. Experimental results demonstrate that proposed method could achieve stable and efficient domain weights learning, and identifies optimal mixtures while consuming substantially less data than search-based function-fitting methods, revitalizing optimization-based domain weighting for large-scale applications.

Problem

• domain weighting
• data mixture optimization
• Large Language Models
• validation loss minimization
• Bayesian inference

Innovation

• Bayesian domain reweighting
• Dirichlet distribution
• Gamma prior
• data mixture optimization
• large language models