💻 computer science

Measuring Task-Agnostic Training Data Influence Across Language Model Pretraining

This paper introduces a task-agnostic method for measuring training data influence by quantifying how much individual examples reduce the distance to final model parameters, revealing that literature-related data drive early pretraining while STEM data become more influential in later stages across various model configurations.

Original authors: Yuto Nishida, Hirokazu Kiyomaru, Yusuke Oda, Takashi Kodama, Chaoran Liu, Daisuke Kawahara, Yusuke Miyao, Max Müller-Eberstein, Masaru Isonuma

Published 2026-08-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yuto Nishida, Hirokazu Kiyomaru, Yusuke Oda, Takashi Kodama, Chaoran Liu, Daisuke Kawahara, Yusuke Miyao, Max Müller-Eberstein, Masaru Isonuma

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a super-smart robot how to speak like a human. You don't just give it one book; you dump the entire internet into its brain. This process is called "pretraining," and it's the foundation of the AI models we use today. But here's the tricky part: once the robot is finished learning, nobody really knows which specific sentences or stories were the most important. Did it learn to be good at math because of a specific Wikipedia article? Did it learn to write poetry because of a specific novel?

Scientists have been trying to answer this by asking, "If I remove this one sentence, does the robot get worse at a specific test, like solving math problems or writing a poem?" This is like trying to figure out which ingredient made a cake taste good by tasting the cake after you've already eaten it. The problem is, AI models are supposed to be good at everything, not just one test. If you only check how well they do on math, you might miss the fact that they learned how to be funny from a different part of the data. So, the big question is: Can we see how the robot's brain changes over time without forcing it to take a specific test?

This paper introduces a clever new way to watch the robot learn, without needing a specific test. Instead of asking, "Did this sentence help the robot pass a math quiz?", the researchers asked, "Did this sentence help the robot get closer to its final, finished brain?" They treated the robot's final state as a destination on a map. Every time the robot read a new sentence, it took a tiny step. The researchers measured whether that step moved the robot closer to the finish line or if it accidentally took a step backward.

They applied this method to 18 different versions of a famous AI model family, watching them learn through 154 different checkpoints (like snapshots of the robot's brain at different times). What they found was a fascinating story about how the robot's "diet" changes as it grows up.

In the very beginning, when the robot was just starting to learn, the most helpful sentences came from literature and stories. These were the "top contributors" that pushed the robot toward its final form. It was as if the robot needed to understand human stories and emotions first. However, as training continued into the middle and later stages, the story changed. The sentences that became most important were no longer stories; they were from STEM fields (Science, Technology, Engineering, and Math).

The researchers saw a clear "crossover." Early on, books and literature were the stars of the show. Later on, technical data about science and computers took over the spotlight. This suggests that the robot doesn't just learn "hard" things at the end; it actually needs a specific sequence of learning. It seems to need the foundation of human language and stories first, and only then does it really benefit from the complex, technical data to reach its full potential.

The authors also checked if their method was reliable. They compared their "distance-to-finish" measurement with the old "test-score" method and found that the old method gave very different answers depending on which test you picked. If you tested for math, you thought math data was important the whole time. If you tested for stories, you thought stories were important the whole time. Their new method, however, showed a consistent pattern across all the different models they tested, suggesting that this shift from stories to science is a real, fundamental part of how these AI models learn, not just an accident of how they were tested.

So, the main takeaway is that training an AI isn't just about feeding it a random mix of data. It's more like a curriculum. The paper suggests that the "best" data changes depending on how far along the AI is in its training. Early on, it needs stories; later on, it needs science. This gives us a new way to think about how to build better AI: maybe we should feed them stories first, and save the complex math for when they are ready for it.

Technical Summary: Measuring Task-Agnostic Training Data Influence Across Language Model Pretraining

Problem Statement
Existing methods for measuring training data influence in language models (LMs) typically rely on task-specific targets, such as loss on a specific test point or performance on a downstream benchmark. While useful for specific goals, these approaches face significant limitations in the context of pretraining, which aims to develop broad, general capabilities rather than optimize for a single task. Selecting downstream tasks or validation sets that are representative of a model's general capabilities throughout the entire pretraining trajectory is difficult. Furthermore, influence estimates tied to intermediate checkpoint performance on specific tasks are often inconsistent and difficult to compare across different stages of training, as early checkpoints may possess relevant capabilities without yet fully solving the target task.

Methodology
The authors propose a task-agnostic measure of training data influence defined by the extent to which an individual training example moves the model toward the final parameters of a given pretraining run.

Definition of Influence: The influence of a mini-batch $B_t$ is quantified as the reduction in the squared $L_2$ distance between the current parameters $\theta_t$ and the final pretrained parameters $\theta^*$ induced by the update at step $t$ . Mathematically, $Cont(B_t) = \|\theta^* - \theta_t\|_2^2 - \|\theta^* - \theta_{t+1}\|_2^2$ . This expands to $2\Delta_t^\top(\theta^* - \theta_t) - \|\Delta_t\|_2^2$ , where the first term measures alignment with the direction toward the final parameters, and the second term penalizes the squared norm of the update.
Example-Level Decomposition: Assuming standard Stochastic Gradient Descent (SGD), the mini-batch influence is decomposed into individual example contributions. The influence of an example $x_k^t$ is defined as $Cont(x_k^t) = 2\Delta_{t,k}^\top(\theta^* - \theta_t) - \Delta_{t,k}^\top \Delta_t$ . This formulation accounts for the interaction between examples within a mini-batch.
Checkpoint-Based Approximation: To make the method tractable for large-scale pretraining without retraining, the authors derive an approximation using periodically saved checkpoints. Instead of computing updates at every step, they approximate the state $\theta_t$ and update $\Delta_t$ using the parameters at the start ( $\theta_c$ ) and end ( $\theta_{c'}$ ) of a checkpoint interval. This allows for post-hoc estimation of influence from publicly available checkpoints.

Key Contributions

Task-Agnostic Reformulation: The paper reformulates training data influence across pretraining in terms of progress toward final parameters, removing the dependency on specific downstream tasks or validation sets.
Checkpoint-Based Estimation: The authors derive a practical approximation method that enables the estimation of example-level influence from intermediate checkpoints without retraining the model from scratch.
Empirical Characterization: The study characterizes the evolution of data contribution throughout pretraining, revealing systematic, stage-dependent variations across different data types.
Reliability Analysis: The work establishes the reliability of the proposed analysis across checkpoint approximations, reference endpoints, and various model configurations (sizes, initializations, and data orderings).

Experimental Results
The method was applied to 18 configurations from the Pythia and PolyPythia suites, covering model sizes from 70M to 12B parameters and utilizing 154 checkpoints per run.

Contribution Dynamics: The mean contribution increases significantly after the early training stage, peaks around 40k steps, and then gradually decreases. The standard deviation follows a similar pattern. Notably, the share of "opponent" examples (those with negative contribution, moving the model away from final parameters) remains low (<1%) during early and middle stages but rises to 8–10% in the late stages.
Text Difficulty: Higher perplexity (PPL) texts (more difficult for the model) account for a larger share of normalized contribution during the middle stage of pretraining, whereas contribution is more evenly distributed across difficulty levels in early and late stages.
Domain Shift (The Crossover): A systematic temporal shift in influential data domains was observed:
- Early Training: Literature-related data (e.g., Books, Hobbies) are more strongly aligned with the trajectory toward final parameters.
- Late Training: STEM-related data (e.g., Science, Computers and Electronics) become increasingly aligned with the trajectory.
- This qualitative crossover is broadly consistent across model configurations, though the timing of transitions may vary with model scale (larger models show transitions earlier).
Validation: The checkpoint-based approximation shows strong correlation (Pearson $r > 0.94$ in late stages) with exact example-level contributions computed via retraining on smaller models. The results are robust to variations in weight initialization and data ordering, though data ordering introduces slightly more variation than initialization.
Comparison with Task-Specific Methods: Unlike task-specific TracIn-style analyses, which yield attribution patterns heavily dependent on the chosen validation domain, the proposed task-agnostic measure consistently reveals the Literature-to-STEM crossover regardless of the specific downstream target.

Significance and Claims
The paper claims to provide a tractable, trajectory-level view of how influential data change throughout pretraining. By decoupling influence analysis from specific downstream tasks, the authors offer a perspective that complements existing task-specific analyses.

The results suggest that the training examples most strongly moving the model toward its final parameters are not static across pretraining. This finding motivates stage-aware approaches to pretraining data composition and curriculum design. Specifically, the observed shift toward STEM data in later stages provides empirical support for recent LLM pretraining recipes that increase the allocation of math and code data in later training stages. The authors position their work as a tool to guide more nuanced pretraining strategies that adapt data composition across training stages rather than treating a single data mixture as equally suitable throughout the entire process.

The authors remain modest regarding causal claims, noting that their analysis characterizes how contribution varies across training rather than directly establishing the causal effect of changing data mixtures on downstream performance. They also acknowledge limitations, such as the reliance on squared parameter-space distance (which is sensitive to parameterization) and the assumption of SGD for decomposition despite the use of adaptive optimizers in practice.

More like this