Measuring Task-Agnostic Training Data Influence Across Language Model Pretraining

AI summary

Overview

Research area: Natural Language Processing — training data attribution and influence analysis for language model pretraining.

Technical level: Intermediate. The paper assumes familiarity with gradient-based training, influence functions, and checkpoint-based analysis, but the core idea (distance to final parameters) is explained geometrically.

Scope: The paper proposes a task-agnostic measure of training data influence defined as an example's reduction of squared L2 distance to the final pretrained parameters, estimates it from existing checkpoints without retraining, and applies it across 18 Pythia and PolyPythia configurations.

What This Paper Is About

Most methods that measure how much a training example matters define "mattering" relative to a specific downstream task, validation set, or loss. That is awkward for language model pretraining, which aims at broad general capabilities rather than one task, and it makes comparisons across intermediate checkpoints unreliable. The authors instead ask a task-free question: how much does a given example's gradient update move the model toward the final parameters of its own pretraining run? The goal is a trajectory-level view of which kinds of data are influential at which stages of pretraining.

Key Contributions

1. A task-agnostic reformulation of training data influence. An example's contribution is defined as the reduction in squared L2 distance to the final pretrained parameters, using the final model as a common reference point instead of a downstream metric.
2. A checkpoint-based estimator. The authors derive an approximation that lets example-level contribution be estimated post hoc from periodically saved checkpoints, without retraining or storing every training step.
3. Empirical characterization of stage-dependent influence. Applying the method across 18 Pythia and PolyPythia configurations, they report systematic temporal shifts in which data types are most aligned with the trajectory toward the final parameters.
4. Reliability and scope analysis. They validate the checkpoint approximation, test sensitivity to the reference endpoint, check consistency across model sizes, weight initializations and data orderings, and compare against a task-specific TracIn-style baseline.

Main Findings

• Contribution peaks in the middle of training. On Pythia-1.4B-Deduped, the mean contribution increases substantially after the early stage, becomes largest around roughly 40k steps, then gradually decreases toward the end of training. The standard deviation follows a similar pattern, peaking in the middle stage.
• "Opponents" grow late in training. Examples with negative contribution are called opponents because their updates move the model away from its final parameters. Their share stays around 0–1% for most of the trajectory before about 70k steps, then rises to roughly 6% around 90k steps and around 8–10% near the end of training.
• Learning rate alone does not explain contribution. The rise in opponent share occurs while the learning rate is already decaying smoothly, consistent with the definition in which contribution depends on alignment with the direction toward the final parameters as well as on update norm.
• Higher-perplexity text contributes more during the middle stage. Using perplexity computed with the final checkpoint of Pythia-12B-Deduped and five equal-sized PPL bins per interval, the highest-PPL bin's share of normalized contribution rises from approximately 20% to 25% in the middle stage, while the lowest-PPL bin falls from approximately 20% to 15% — a five-percentage-point redistribution.
• A literature-to-STEM crossover. Among the bottom 5% of examples by contribution, STEM-related domains such as Computers and Electronics and Science dominate in the first half of training, while Books and Literature becomes more common later. Among the top 5%, STEM domains are not necessarily dominant early but their share tends to increase in later stages.
• Checkpoint approximation works better later in training. On Pythia-70M-Deduped, exact versus approximate example-level contribution gives Pearson correlations of 0.622 (1k→2k), 0.808 (10k→11k), and 0.950 (100k→101k), with Spearman correlations of 0.592, 0.811, and 0.947; all are significantly different from zero under a two-sided test (p<0.05).
• The score is endpoint-dependent but locally stable. On Pythia-1.4B-Deduped, using the 140k checkpoint as reference instead of the final 143k checkpoint yields a mean Spearman correlation of 0.997 and mean top-5% overlap of 0.972 averaged over five intervals, whereas earlier references (30k, 70k, 120k) produce markedly different rankings.
• Cross-configuration consistency with scale effects. Across the six Pythia-Deduped sizes the broad temporal patterns are qualitatively consistent, but transitions occurring in the middle stage for larger models appear later for smaller models, and the magnitude of the dynamics tends to be smaller at larger scales.
• Data ordering matters more than initialization. PolyPythia runs show consistent domain-level dynamics across seeds, except the anomalous 410M seed-3 run that van der Wal et al. (2025) identify as an outlier with loss spikes and degraded performance. The decoupled PolyPythia-160M variants are highly consistent when only weight initialization is varied, while changing data ordering introduces somewhat greater variation.
• Task-specific attribution does not reproduce the crossover. A TracIn-style score computed against domain-wise language-modeling loss on held-out Pile validation sets shows attribution patterns that depend substantially on the chosen validation domain — for example, overrepresentation of the selected domain's examples among top scorers early in training, especially with Science as the target — and does not consistently recover the Literature-to-STEM crossover.

Methodology in Plain English

The authors treat each 2048-token training sequence as one example. For a pretraining run, they take the final model parameters as a fixed target. At any point in training, the model sits at some squared distance from that target. A training update is considered influential to the extent that it shrinks that distance.

Concretely, mini-batch contribution is written as the distance before an update minus the distance after it, which algebraically splits into two parts: how well the update points toward the final parameters, minus a penalty for the update's squared size. The example-level version distributes a mini-batch update across its examples, assuming standard SGD for this decomposition; the authors note that real pretraining uses adaptive optimizers such as Adam, so the per-example view is an approximation. Contributions can be negative, which is what defines "opponents."

Because storing every training step is impractical, they replace step-level quantities with checkpoint-level ones: all examples in an interval between two saved checkpoints are treated as if processed at the earlier checkpoint state. This makes the measure applicable post hoc to publicly released checkpoints.

In experiments, they sample examples directly from the actual training data stream associated with each checkpoint interval rather than from held-out or proxy corpora. They analyze the contribution distribution (1,000 sampled examples per interval), its relationship to perplexity (binning into five equal PPL groups and min–max normalizing contributions within each interval), and its relationship to textual domain (using the NeMo Curator Domain Classifier, which assigns one of 26 domains, and sampling 100 examples per domain per interval, i.e., 2,600 examples per interval, focusing on the top and bottom 5% by contribution).

Why This Matters

Impact on research. The paper offers a way to study pretraining data influence without committing to a downstream task or validation set, which the authors argue is a partial view at best. It complements existing attribution work — influence functions, Data Shapley, TracIn, TRAK, LoGra — and connects to observations that influence patterns evolve during training (Wang et al., 2025), that causal influence and factual attribution diverge (Chang et al., 2025), and that attribution evaluation is sensitive to task-specific design (Jiao et al., 2025). It also frames the observed dynamics in relation to reported Pythia training phases.

Real-world applications:

• Stage-aware data mixture design. The finding that STEM-related data become increasingly prominent among high contributors later in training lends trajectory-level support to recipes that increase math and code allocation in later pretraining stages.
• Data curation and valuation. A task-free contribution score could help rank or filter pretraining corpora without repeatedly running downstream benchmarks.
• Diagnosing training dynamics. The opponent share offers a signal for detecting late-training examples whose updates oppose progress toward the final model.
• Provenance and dataset auditing. The measure can be applied post hoc to released checkpoints, which lowers the barrier to auditing large public pretraining runs.

Industry relevance. Pretraining runs are expensive and data mixtures are often fixed in advance. A measurement that reuses already-saved checkpoints rather than requiring retraining is directly usable by teams that release or operate large models, and the crossover result speaks to a heuristic — shifting toward STEM data in late training — that current large-model recipes already adopt.

Future Directions

1. Function-aware contribution measures. The authors note that squared parameter-space distance is sensitive to parameterization and does not directly measure behavioral change; predictive divergence or information-geometry alternatives are suggested but described as substantially more expensive at pretraining scale.
2. Extending exact validation. The checkpoint approximation was validated only on three intervals of Pythia-70M-Deduped; exact validation at larger model scales remains computationally expensive and is left open.
3. Beyond SGD assumptions. Accounting exactly for adaptive optimizer states and their history would require a substantially more involved decomposition than the per-example SGD-based one used here.
4. Causal intervention and broader generality. The analyses characterize how contribution varies across training rather than establishing the causal effect of changing the data mixture, and all experiments stay within closely related model families and pretraining data. The authors call for direct intervention-based evaluation of stage-aware data-composition strategies, and note that disentangling domain effects from correlated properties such as text difficulty remains an important direction.

Target Audience

Researchers and engineers working on data attribution, data valuation, and pretraining data curation; practitioners designing stage-aware data mixtures or curricula for large language models; and anyone studying training dynamics of public model suites such as Pythia and PolyPythia. Readers should be comfortable with gradient-based training and influence-function concepts, though the core geometric intuition is accessible without a deep mathematical background.

Authors’ abstract

Measuring training data influence consistently across language model pretraining is challenging. It is difficult to select downstream tasks or validation sets representative of a model's general capabilities, and reliance on task performance at intermediate checkpoints complicates comparisons across training. We propose a measure of training data influence that does not require selecting a downstream task or validation set as the attribution target. Specifically, we define an example's influence by how much its gradient update reduces the squared distance to the final parameters of a given pretraining run, and estimate this quantity from intermediate checkpoints without retraining. Applying the method to 18 configurations from the Pythia and PolyPythia suites, we find systematic temporal changes in influential data. Early in training, literature-related data are more strongly aligned with the trajectory toward the final parameters, whereas STEM data become more strongly aligned in later stages. This qualitative crossover is broadly consistent across model configurations. Our results provide a tractable trajectory-level view of how influential data change throughout pretraining, complementing influence analyses defined with respect to specific downstream tasks or validation sets.