Data Pacing and Annealing Schedules

While static data mixtures provide a fixed recipe for training, large language models might benefit from a more dynamic approach to data presentation, akin to how a curriculum guides learning. Instead of feeding the model the same proportion of data sources throughout the entire training run, we can dynamically adjust the mixture or the sampling process itself over time. This is the core idea behind data pacing and annealing schedules. These techniques aim to optimize the learning process by controlling when and how different parts of the dataset are emphasized.

Data Pacing: Controlling the Rate of Exposure

Data pacing refers to the strategy of controlling the rate at which the model is exposed to different subsets or types of data during training. Think of it as managing the flow of information based on the model's learning progress. For example, you might start training primarily on a high-quality, curated corpus to establish strong foundational language understanding. As training progresses and the model becomes more capable, you can gradually increase the proportion of data from noisier sources like web crawls or specialized domains like code repositories.

The pacing can be based on various criteria:

1. Source Quality: Start with clean data, gradually introduce noisier or more complex data.
2. Domain Specificity: Begin with general text, then phase in domain-specific text if desired.
3. Data Difficulty: If you can estimate data difficulty (e.g., based on text complexity or rarity of concepts), you could pace the introduction from simpler to harder examples, although defining difficulty at scale is challenging.

Implementing data pacing often involves creating a schedule function that modifies the sampling weights or probabilities associated with different data sources based on the current training step or epoch.

import numpy as np

# Example data sources and initial weights
data_sources = {
    "curated_corpus": {"weight": 0.7, "path": "/path/to/curated"},
    "web_crawl": {"weight": 0.3, "path": "/path/to/web"},
    # Add more sources as needed
}

# Total training steps
total_steps = 1_000_000

def get_pacing_weights(current_step, total_steps, sources):
    """Calculates dynamic weights based on training progress."""
    progress = current_step / total_steps
    new_weights = {}

    # Example: Linearly decrease curated data weight, increase web crawl weight
    # Ensure weights still sum to 1 (or normalize afterwards)
    curated_weight = max(0.1, 0.7 - 0.6 * progress) # Don't go below 10%
    web_weight = 1.0 - curated_weight # Assuming only two sources for simplicity

    new_weights["curated_corpus"] = curated_weight
    new_weights["web_crawl"] = web_weight
    # Add logic for other sources if present

    # Normalize weights to ensure they sum to 1 (important!)
    total_weight = sum(new_weights.values())
    normalized_weights = {
        k: v / total_weight
        for k, v in new_weights.items()
    }

    return normalized_weights

# --- Inside the training loop ---
# current_training_step = ... get current step ...
# current_weights = get_pacing_weights(
#     current_training_step, total_steps, data_sources)
# Reconfigure the data sampler/loader with these new weights
# sampler.set_weights(current_weights)
# batch = next(data_loader)
# ... rest of training step ...

This snippet illustrates a simple linear pacing schedule. More complex functions (e.g., exponential, step-based) can be designed depending on the desired learning trajectory. The important part is dynamically adjusting the likelihood of sampling from each source as training advances.

Annealing Schedules for Data Sampling

Annealing, in the context of data sampling, typically refers to gradually changing a parameter that controls the shape or randomness of the sampling distribution over the data sources. This is closely related to the temperature-based sampling discussed previously.

Recall that temperature $T$ modifies the probabilities $p_i$ derived from weights $w_i$:

$$p_i = \frac{\exp(w_i / T)}{\sum_j \exp(w_j / T)}$$

An annealing schedule might involve starting with a higher temperature $T > 1$ and gradually reducing it towards $T = 1$ (or even slightly lower) over the course of training.

• High Temperature (Early Training): When $T > 1$, the probabilities $p_i$ become more uniform. This encourages exploration, meaning the model samples more evenly across different data sources, even those with lower initial weights. This can be beneficial early on to prevent the model from specializing too quickly on the dominant data sources.
• Lower Temperature (Later Training): As $T$ decreases towards 1, the sampling probabilities become closer to the distribution defined by the raw weights $w_i$. If $T < 1$, the distribution sharpens, further emphasizing the higher-weighted sources. This allows the model to consolidate its learning on the data considered most important or representative towards the end of training.

Annealing schedules can be applied to the temperature parameter or directly to the weights themselves (e.g., gradually increasing the difference between high and low weights). Common schedules include linear decay, cosine decay, or exponential decay of the temperature parameter or a similar modulation factor applied to the weights.

Example temperature annealing schedules over 1 million training steps, starting from T=2.0 and decaying towards T=1.0 using linear and cosine functions.

Combining Pacing and Annealing

Data pacing and annealing are not mutually exclusive. You can combine them to create sophisticated data feeding strategies. For instance:

• Use a pacing schedule to gradually change the target weights of different data sources over time (e.g., increasing the weight of domain-specific data).
• Simultaneously, use an annealing schedule on the sampling temperature to control the randomness of sampling based on those target weights (e.g., start with high temperature for exploration, decrease it later for exploitation of high-weighted sources).

This combination allows fine-grained control over both what data the model sees and how strictly it adheres to the prescribed mixture at different training stages.

Implementation Approaches

Implementing dynamic sampling schedules adds complexity to the training pipeline.

• Tracking Progress: The training loop needs to accurately track the current step or epoch to calculate the schedule's current state.
• Sampler/Loader Integration: The data loading mechanism must be flexible enough to update its sampling probabilities or source weights dynamically, potentially at every step or every few thousand steps. This might involve custom sampler logic. Frameworks like torch.utils.data.WeightedRandomSampler can be adapted, or custom iterators managing multiple datasets might be needed.
• Monitoring: Closely monitor training stability (loss curves, gradient norms) and evaluation metrics when using dynamic schedules. Abrupt changes in the data distribution can sometimes lead to temporary instability.
• Tuning: Finding the optimal pacing or annealing schedule often requires experimentation. The ideal schedule shape (linear, cosine, step) and its duration depend heavily on the datasets, model size, and training objectives.

While more complex than static sampling, data pacing and annealing schedules offer powerful tools for guiding the learning process of large language models. By thoughtfully controlling the data diet over time, you can potentially accelerate convergence, improve robustness, and better shape the final capabilities of the model to meet specific requirements.