ToReMi: Topic-Aware Data Reweighting for Dynamic Pre-Training Data Selection

TL;DR

ToReMi tackles the challenge of efficient pre-training data selection for large language models by introducing a two-stage, topic-aware reweighting framework. It dynamically adjusts topic-level weights during training to emphasize high-loss topics initially and downweight noisy topics later, thereby balancing learning progression with noise reduction. Empirical results using GPT-2 124M on a 30B-token corpus show accelerated perplexity reduction across 12 Paloma domains and improves several downstream tasks, with performance that is robust across a range of hyperparameters. The work highlights topic-aware data weighting as a practical lever to improve both the efficiency and generalization of large-scale pre-training.

Abstract

Pre-training large language models (LLMs) necessitates enormous diverse textual corpora, making effective data selection a key challenge for balancing computational resources and model performance. Current methodologies primarily emphasize data quality metrics and mixing proportions, yet they fail to adequately capture the underlying semantic connections between training samples and quality disparities within individual domains. We introduce ToReMi (Topic-based Reweighting for Model improvement), a novel two-stage framework that dynamically adjusts training sample weights according to their topical associations and observed learning patterns. Our comprehensive experiments reveal that ToReMi variants consistently achieve superior performance over conventional pre-training approaches, demonstrating accelerated perplexity reduction across multiple domains and enhanced capabilities on downstream evaluation tasks. Code is available at https://github.com/zxx000728/ToReMi.

Figures (7)

• Figure 1: The framework of ToReMi, a two-stage, topic-based reweighting method for dynamic pre-training data selection and model improvement. During each training interval, training samples are reweighted based on their topic labels and previous training dynamics.
• Figure 2: Topic distribution of the 30B-token dataset organized by Wikipedia taxonomy.
• Figure 3: The log perplexity for different methods on the Paloma test dataset across 12 domains. ToReMig refers to ToReMi with directly generated topic labels, and ToReMis refers to ToReMi with topic labels selected from Wikipedia taxonomy.
• Figure 4: Performance difference between the standard method and ToReMi with various stage transition points. Red indicates performance improvement over the standard model, while blue indicates degradation.
• Figure 5: Performance difference between the standard method and ToReMi variants with various weight upper bounds. Red indicates performance improvement over the standard model, while blue indicates degradation.