Table of Contents
Fetching ...

Scalable Influence and Fact Tracing for Large Language Model Pretraining

Tyler A. Chang, Dheeraj Rajagopal, Tolga Bolukbasi, Lucas Dixon, Ian Tenney

TL;DR

This work presents TrackStar, a scalable gradient-based method for training data attribution and influence tracing during large-language-model pretraining. TrackStar combines optimizer state correction, a task-aware Gauss-Newton Hessian proxy, and normalized gradient encodings with random projections to retrieve influential pretraining examples from massive corpora without subsampling. Across T-REx closed-set and C4 open-set evaluations, TrackStar outperforms prior gradient methods in measuring influence and, to a lesser extent, attribution, while revealing a persistent gap between factual entailment and causal influence that narrows as model scale increases. The findings show that while classical retrieval like BM25 excels at attribution, influence-based proponents often drive model predictions more directly, suggesting distinct yet complementary roles for data curation and diagnostic tooling in large-scale pretraining.

Abstract

Training data attribution (TDA) methods aim to attribute model outputs back to specific training examples, and the application of these methods to large language model (LLM) outputs could significantly advance model transparency and data curation. However, it has been challenging to date to apply these methods to the full scale of LLM pretraining. In this paper, we refine existing gradient-based methods to work effectively at scale, allowing us to retrieve influential examples for an 8B-parameter language model from a pretraining corpus of over 160B tokens with no need for subsampling or pre-filtering. Our method combines several techniques, including optimizer state correction, a task-specific Hessian approximation, and normalized encodings, which we find to be critical for performance at scale. In quantitative evaluations on a fact tracing task, our method performs best at identifying examples that influence model predictions, but classical, model-agnostic retrieval methods such as BM25 still perform better at finding passages which explicitly contain relevant facts. These results demonstrate a misalignment between factual *attribution* and causal *influence*. With increasing model size and training tokens, we find that influence more closely aligns with factual attribution. Finally, we examine different types of examples identified as influential by our method, finding that while many directly entail a particular fact, others support the same output by reinforcing priors on relation types, common entities, and names. We release our prompt set and model outputs, along with a web-based visualization tool to explore influential examples for factual predictions, commonsense reasoning, arithmetic, and open-ended generation for an 8B-parameter LLM.

Scalable Influence and Fact Tracing for Large Language Model Pretraining

TL;DR

This work presents TrackStar, a scalable gradient-based method for training data attribution and influence tracing during large-language-model pretraining. TrackStar combines optimizer state correction, a task-aware Gauss-Newton Hessian proxy, and normalized gradient encodings with random projections to retrieve influential pretraining examples from massive corpora without subsampling. Across T-REx closed-set and C4 open-set evaluations, TrackStar outperforms prior gradient methods in measuring influence and, to a lesser extent, attribution, while revealing a persistent gap between factual entailment and causal influence that narrows as model scale increases. The findings show that while classical retrieval like BM25 excels at attribution, influence-based proponents often drive model predictions more directly, suggesting distinct yet complementary roles for data curation and diagnostic tooling in large-scale pretraining.

Abstract

Training data attribution (TDA) methods aim to attribute model outputs back to specific training examples, and the application of these methods to large language model (LLM) outputs could significantly advance model transparency and data curation. However, it has been challenging to date to apply these methods to the full scale of LLM pretraining. In this paper, we refine existing gradient-based methods to work effectively at scale, allowing us to retrieve influential examples for an 8B-parameter language model from a pretraining corpus of over 160B tokens with no need for subsampling or pre-filtering. Our method combines several techniques, including optimizer state correction, a task-specific Hessian approximation, and normalized encodings, which we find to be critical for performance at scale. In quantitative evaluations on a fact tracing task, our method performs best at identifying examples that influence model predictions, but classical, model-agnostic retrieval methods such as BM25 still perform better at finding passages which explicitly contain relevant facts. These results demonstrate a misalignment between factual *attribution* and causal *influence*. With increasing model size and training tokens, we find that influence more closely aligns with factual attribution. Finally, we examine different types of examples identified as influential by our method, finding that while many directly entail a particular fact, others support the same output by reinforcing priors on relation types, common entities, and names. We release our prompt set and model outputs, along with a web-based visualization tool to explore influential examples for factual predictions, commonsense reasoning, arithmetic, and open-ended generation for an 8B-parameter LLM.

Paper Structure

This paper contains 29 sections, 4 equations, 3 figures, 9 tables.

Figures (3)

  • Figure 1: Top proponents from C4 using TrackStar given a factual query and model prediction. TrackStar is a gradient-based method that approximates influence on the model, which we show may not always be optimal for attribution, which involves finding examples which directly entail the target factual prediction.
  • Figure 2: Left, center: attribution (MRR) and influence (tail-patch) scores as a function of gradient projection dimensionality $d$ for different model sizes (§\ref{['sec:ablations']}). Right: attribution (MRR) scores throughout pretraining for different model sizes. As models improve, TrackStar influence becomes more similar to attribution (higher MRR; §\ref{['sec:infapproachesatt']}).
  • Figure 3: Proportions of TrackStar proponents (top 10 per query) retrieved from C4 that entail a prediction, contain both entities, or contain only one entity for a model prediction.