Search papers, labs, and topics across Lattice.
This paper introduces a Bayesian domain weighting method that optimizes the distribution of multi-domain pre-training data for Large Language Models (LLMs) by inferring weights from a Dirichlet distribution with Gamma prior information. This approach addresses the limitations of traditional function-fitting methods, which often suffer from estimation bias due to structural assumptions that do not hold in complex data environments. Experimental results show that the proposed method achieves stable and efficient learning of domain weights, enabling optimal data mixtures with significantly reduced data consumption compared to existing techniques.
Bayesian domain weighting can dramatically enhance LLM performance by efficiently optimizing data mixtures with less computational overhead and data usage.
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.
