Source Weighting Strategies

Training large language models often involves assembling a diverse collection of text data from sources such as web crawls (e.g., Common Crawl), books, code repositories, and specialized corpora. A critical aspect of this process is determining the proportional contribution of each data source to the training. Simply concatenating all available data and shuffling it uniformly might appear to be a straightforward approach, but it seldom produces optimal results. This is because different data sources inherently vary in terms of quality, relevance, and linguistic style. Source weighting provides precise control over the overall composition of the training data mixture, allowing the model to prioritize learning from specific text types over others.

The core idea is to assign a specific weight or probability to each data source. During training, when assembling a batch of data, examples are sampled from the different sources according to these predefined weights. Sources assigned higher weights will contribute more examples to the training process over time compared to sources with lower weights.

Defining Sampling Probabilities

Let's assume you have $k$ distinct data sources, $S_1, S_2, \dots, S_k$. Each source $S_i$ contains $N_i$ documents or examples. We want to define a probability $p_i$ of sampling an example from source $S_i$ at any given step. A common approach is to make this probability proportional to both the size of the source and an assigned weight $w_i$:

$$p_i \propto w_i \cdot N_i$$

Alternatively, and perhaps more directly controllable, you can define the desired proportion (or probability) $p_i$ for each source directly, such that $\sum_{i=1}^{k} p_i = 1$. The weights $w_i$ then simply become these target proportions $p_i$.

For instance, you might decide on the following mixture:

• Web Text (Source 1): $p_1 = 0.60$ (60%)
• Books (Source 2): $p_2 = 0.15$ (15%)
• Code (Source 3): $p_3 = 0.15$ (15%)
• Wikipedia (Source 4): $p_4 = 0.10$ (10%)

This means that, on average, 60% of the examples seen by the model in each training epoch (or over a large number of steps) will come from the Web Text dataset, 15% from Books, and so on.

Example proportions for sampling from different data sources during LLM pre-training.

Heuristics for Setting Weights

Choosing the right weights is often guided by a combination of factors and requires careful consideration and experimentation:

1. Data Quality: High-quality, well-curated sources like Wikipedia or peer-reviewed articles might receive higher weights relative to their size compared to noisy web crawl data. This steers the model towards learning more factual and well-structured language.
2. Domain Specialization: If you are building a model intended for a specific domain (e.g., a medical chatbot or a code generation assistant), you might increase the weight of relevant domain-specific datasets (e.g., PubMed abstracts or code repositories).
3. Data Volume vs. Influence: Very large datasets (like Common Crawl) can easily dominate the training mixture if sampled proportionally to their raw size. Assigning explicit weights prevents this, ensuring that smaller, high-quality, or domain-specific datasets still have a significant influence on the model's learning.
4. Preventing Contamination: If evaluating on specific benchmarks, you might down-weight or exclude data sources known to contain excerpts from those benchmarks to ensure fair evaluation.
5. Desired Model Capabilities: If fluency in conversational text is important, you might include and assign a reasonable weight to datasets containing dialogues or forum discussions. If factual grounding is crucial, structured text and encyclopedic sources might be weighted higher.

Implementation Approaches

Implementing source weighting typically involves modifying the data loading process. Instead of sampling uniformly from a single large dataset, the data loader needs to be aware of the different sources and their associated weights.

In PyTorch, you can achieve this by creating a custom IterableDataset or using samplers that accommodate weighted sampling across multiple underlying datasets. A simplified example might look like this:

import torch
import numpy as np
from torch.utils.data import IterableDataset, DataLoader

# Assume these are placeholder datasets for different sources
# In reality, these would load actual tokenized data
class DummyDataset(IterableDataset):
    def __init__(self, source_name, size):
        self.source_name = source_name
        self.size = size

    def __iter__(self):
        for i in range(self.size):
            # Yield dummy data: (example_data, source_identifier)
            yield torch.randn(512), self.source_name
            if i % 1000 == 0 and i > 0 : # Simulate potentially large dataset
                 print(f"Yielded {i} from {self.source_name}")

# Define sources and their desired sampling probabilities (weights)
sources = {
    "web": {"dataset": DummyDataset("web", 1_000_000), "weight": 0.60},
    "books": {"dataset": DummyDataset("books", 200_000), "weight": 0.15},
    "code": {"dataset": DummyDataset("code", 300_000), "weight": 0.15},
    "wiki": {"dataset": DummyDataset("wiki", 100_000), "weight": 0.10},
}

source_names = list(sources.keys())
source_weights = np.array([sources[name]["weight"] for name in source_names])
# Normalize weights to ensure they sum to 1 (optional if already normalized)
# source_weights /= source_weights.sum()

source_iters = {name: iter(sources[name]["dataset"]) for name in source_names}

class WeightedSourceSampler(IterableDataset):
    def __init__(self, source_names, source_weights, source_iters):
        self.source_names = source_names
        self.source_weights = source_weights
        self.source_iters = source_iters

    def __iter__(self):
        while True:
            # Choose a source based on weights
            chosen_source_name = np.random.choice(
                self.source_names, p=self.source_weights
            )

            try:
                # Get the next item from the chosen source's iterator
                item = next(self.source_iters[chosen_source_name])
                yield item
            except StopIteration:
                # If a source iterator is exhausted, recreate it (or handle as needed)
                print(f"Restarting iterator for {chosen_source_name}")
                self.source_iters[chosen_source_name] = iter(
                    sources[chosen_source_name]["dataset"]
                )
                # Optionally, break or implement logic for finite datasets
                # For large-scale pre-training, iterators often cycle indefinitely
                # Re-attempt fetching after reset:
                try:
                   item = next(self.source_iters[chosen_source_name])
                   yield item
                except StopIteration:
                   print(
                       f"Warning: Iterator for {chosen_source_name} "
                       f"immediately exhausted after reset."
                   )
                   # Decide how to handle this case, e.g., re-sample source or raise error
                   # For this example, we might just skip this turn.
                   continue 

# Create the combined dataset sampler
weighted_sampler_dataset = WeightedSourceSampler(
    source_names, source_weights, source_iters
)

# Use it with a DataLoader
# Note: num_workers > 0 requires careful handling with 
# IterableDatasets and iterators
# For simplicity, using num_workers=0 here. 
# Real implementations need multi-worker support.
data_loader = DataLoader(weighted_sampler_dataset, batch_size=4, num_workers=0) 

# Example of fetching a batch
print("Fetching a batch...")
batch = next(iter(data_loader))
# batch[0] contains the data tensors, batch[1] contains source names
print(f"Batch data shape: {batch[0].shape}") 
print(f"Batch source identifiers: {batch[1]}")

# Simulate fetching a few more batches to see source distribution
print("\nFetching more batches...")
for i in range(5):
     batch = next(iter(data_loader))
     print(f"Batch {i+1} sources: {batch[1]}")

PyTorch implementation sketch for sampling from multiple data sources based on predefined weights.

This example demonstrates the basic principle. Production-grade data loaders for LLMs often involve more sophisticated mechanisms for efficiency, handling distributed training, and managing extremely large datasets that don't fit in memory, potentially using streaming techniques.

Trade-offs

While powerful, source weighting introduces complexity:

• Tuning Difficulty: Finding the optimal weights usually requires significant experimentation and evaluation on downstream tasks, which can be computationally expensive.
• Bias Amplification: Over-weighting certain sources can lead the model to adopt biases present in that data. Conversely, under-weighting important but smaller datasets might hinder performance in specific areas.
• Dynamic Weighting: As discussed later in the context of curriculum learning and annealing, weights might not be fixed throughout training. They could change dynamically based on the training stage or model performance.

Source weighting is a fundamental technique for managing the composition of the massive datasets used in LLM pre-training. By carefully considering the quality, relevance, and volume of different data sources and assigning appropriate weights, engineers can better guide the learning process and shape the capabilities of the final model.