Training Data Distribution Estimation for Optimized Pre-training Data Management

Large language models (LLMs) have demonstrated exceptional performance across a wide range of tasks and domains, with data preparation playing a critical role in achieving these results. Pretraining data typically combines information from multiple domains. To maximize performance when integrating data from various domains, determining the optimal data distribution is essential. However, state-of-the-art (SOTA) LLMs rarely disclose details about their pretraining data, making it difficult for researchers to identify ideal data distributions. In this paper, we introduce a new approach, data distribution estimation, which enables the automatic estimation of pretraining data distributions by analyzing the generated outputs of LLMs. We provide rigorous theoretical proofs, practical algorithms, and preliminary experimental results for data distribution estimation. Based on these findings, we offer valuable insights into the challenges and future directions for effective data distribution estimation and data management. The source code, data, and other artifacts are available at https://github.com/yangyajie0625/data_detection