Scalable Influence and Fact Tracing for Large Language Model Pretraining
Authors: Tyler A. Chang
Dheeraj Rajagopal
Tolga Bolukbasi
Lucas Dixon
Ian Tenney
Venue: ICLR 2025 Poster (2025)
Keywords: training data attribution
LLM pretraining
influence functions
factual knowledge
Primary area: interpretability and explainable AI
TL;DR: We present a gradient-based influence method that scales to LLM pretraining, and we demonstrate its ability to retrieve examples that influence LLM factual predictions.

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.

Published as a conference paper at ICLR 2025
SCALABLE INFLUENCE AND FACT TRACING FOR LARGE LANGUAGE MODEL PRETRAINING
Tyler A. Chang, 1,2- Dheeraj Rajagopal,' Tolga Bolukbasi,' Lucas Dixon,' Ian Tenney'
tachang@ucsd.eduu, (zajagopald, tolgab, Idixon, ftteeneelll @google. 'Google DeepMind
UC San Diego
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 cu- ration. 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- basend methods to work effectively at scale, allowing us to retrieve influential examples for an 8B- parameter language model from pretraining corpus of over 160B tokens with no need for subsampling or pre-filtering Our method combines several techniques, including optimizer state correction, task-specific Hessian approximation, and normalized encodings, which we find to be critical for performance at scale. In quantitative evaluations on fact tracing task, our method performs best at iden- tifying examples that influence model predictions, but classical, model-agnostic retrieval methods such as BM25 still perform better at finding passages which ex- plicitly contain relevant facts. These results demonstrate misalignment between factual attribution and causal influecce 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 and model outputs, along with web-based visualization tool to explore influential examples for factual predictions, commonsense reason- ing, arithmetic, and open-ended generation for an 8B SB-parameter LLM
INTRODUCTION
Modern large language models (LLMs) perform extraordinarily well at wide variety of natural language tasks, but exactly how they leverage training data achieve such capabilities is not well understood. One promising avenue of research to study this is training data aftribution (TDA), which aims to identify influential training examples for given model predictions. When successful, TDA can serve method to both inspect and intervene on the training process. For example, it can enable training data curation targeted at specific tasks (Engstrom et al. 2024), reduction of data contamination (Mozes et al. 2023), and better transparency into model predictions (Grosse et al. 2023
Choe et al. 2024).
However, the steep computational cost of applying TDA approaches to LLMs has limited work in this area. TDA approaches have achieved promising results identifying examples that influence LLM behavior during fine-tuning (Akyurek al. 2022 Park et al., 2023 Xia al. 2024). but there is
significant evidence that much of an LLMM's knowledge and capabilities originates from pretraining
(Hoffmaan et al., 2022 Chang & Bergen 2024). Previous work applying TDA to pretraining has
either focused on small models (Engstrom et al. (2024) or extremely few target queries, e.g. retriev- ing influential examples from a significantly subsampled corpus for less than 100 model predictions Grrsse et al. 2023 Choe et al. 2024] Ruis et al. 2024). In our work, we scale TDA experiments LLMS up to 8B parameters, thousands of model predictions, and corpora up 160B tokens.
Work done as a student rescarcher at Google Research and Google DeepMind.
https: / thub. www. / p-ooirreeetettaeannninn http://b...oo//aaaooooorerrrraaa

---

Published as a conference paper at ICLR 2025
Prompt: 21st Century Breakdown is of the following genre:
Output: Punk Rock
#1: 21st Century Breakdown proves that Green Day still has the power rally troops,..
Pretraining Dataset
TrackStar
year ago today, 21st Century Breakdown released, the first
#9: 21st Century Breakdown is the eighth studio album American punk rock band Green Day.
Figure 1: Top proponents from C4 using TrackStar given a factual query and model prediction. TrackStar is gradient-basee based method that approximates inffuence the model, which show may not always be optimal for attribution, which involves finding examples which directly entail the target factual prediction.
Specifically we propose TrackStar, gradient-based influence method that combines innovations from previous work and scales to large setups ($3), while still supporting efficient retrieval of influ- ential examples and quantitative evaluation. At scale, our method significantly outperforms previous influence methods at retrieving pretraining examples that entail a fact (attribution) and examples that influence specific factual predictions (influence) ($5) Ş6). Importantly, ment between attribution and influence; classical methods such as BM25 are better at retrieving examples that entail factual predictions, but those examples are not necessarily those that most af- fect model predictions ($5.T). We show that influence grows closer to attribution as models scale, both in parameters and training tokens ($5.3).
demonstrate misalign-
To provide insights into where attribution and influence misalign, we include analyses of examples that have high influence for factual predictions in LLMs ($7). For example, rather than containing a full statement of a fact of interest, many examples appear to support priors on relation types, common entities, or names.
2 BACKGROUND: TRAINING DATA ATTRIBUTION
Researchers have proposed variety of methods measure the influence of individual training examples on output model metrics (e.g. loss on target datasets). Some of the best results come from simulation-basse methods such as Datamodels (Ilyas et al.) (2022) and Simfluence (Guu al., 2023). which estimate contributions based on multiple training runs with different data subsets. However, this not tractable for LLM pretraining, which is expensive to perform even once. For computational tractability, we focus gradient-based methods that use parametric approximations, based on a single model, to predict how a model's behavior would change under the removal (or addition) of specific training examples.
To quantify the influence from training example 2m an evaluation (query) example Zg- most gradient-based influence methods compute some version of Hessian-correcter dot product of model gradients for ². and 2g loss (Koh & Liang, 2017 Schioppa al, Park al. 2023 Grosse et al. 2023). Here, H- is the inverse Hessian matrix the loss with respect to model parameters. The resulting dot product approximates changes in Z loss from adding example im during training. practice, existing methods use variety of methods to approximate
LV -
2022
the Hessian H1, notably the autocorrelation matrix ""' E [[(0x|0] of per-example gradients in TRAK (Park et al. 2023; Sagun et al. 2018) or the closely-related approximate curvature matrix G of EK-FAC (Grosse et al. 2023). These methods also tie closely back to TracIn (Pruthi et al. 2020),
which computes gradient dot products aggregated over model checkpoints; recent work using TracIn has included gradient second moment correction (Akyurek et al.) (2022) Xia et al. 2024). equivalent to diagonal Hessian approximation ($3)
2

---

Published as a conference paper at ICLR 2025
However, when applying influence methods to LLM pretraining, additional bottleneck is the scale of model parameters and pretraining data. Given [0] model parameters (on the order of billions), we have loss gradients [2)2) E |R and inverse Hessian H-1 E [| To rank pretraining examples, the gradient dot product must be computed between every pretraining example 2m (on the order of millions to billions) and the target query 2q. For this reason, previous work has focused on identifying influential examples during fine-tuning (Akyurek al., 2022 Park al. 2023 Xia itt al., 2024) or during pretraining for smaller models (Engstrom et al. 2024). Closest to our work has been that of Grosse et al. (2023), Choe et al. (2024), Ruis et al. (2024), who look at LLMs
with billions of parameters; however, they perform retrieval on a small subset of the pretraining corpus and report evaluation on only a small number of queries.
METHOD: TRACKSTAR
Here, introduce TrackStar, gradient-based influence method that combines innovations from previous work that scales effectively to large settings. Following the gradient-based methods in 2 we compute the influence between two examples as the dot product between the projected and corrected model gradients for those examples. Because correction terms vary substantially in previ- ous work, here we describe the motivation behind our specific approach, with ablation experiments in $.2 Specifically, given a query bq (an input prompt with corresponding model completion or desired completion) and a training example Zm, we compute the influence from 2m to Zg as:
Io (zmi 2q) (Zm (zg(z))
where Gz(z) is the projected, Hessian-corrected, and unit normalized gradient for example given model parameters 0:
Go(z) Go(z) ||Go())]12 Go(z) R- Pa
VoLoss(2,0) 0)
VV
In turn, we define:
Loss gradient: V VoLoss(z,0)) 0) E RR As in previous work, we compute the loss gradient for each example with respect model parameters 0 (Pruthi et al. 2020 Akyurek et al.) 2022 Han & Tsvetkov] [022 Grosse et al. 2023) Xia et al. 2024). The original TRAK method uses the multi-class margin function gradient; we evaluate different output functions in $A.2.1 but we find that loss gradients perform best empirically, For each example z, we sum gradients over target tokens; for training examples, this comprises all tokens in the example. If an input prompt is given, we only include tokens in the model completion. In contrast to previous work that selects only a subset of layers (Akyurek et al. 2022, Ych et al. 2022; Grosse et al. (023). we compute gradients with respect to all model parameters 0 except the token embedding layer, pooled into layer blocks ($A.2.2) before dimensionality reduction with random projection.
Second moment estimate V E RRO To account for high-magnitude gradient components that might dominate gradient dot products (e.g. for outlier model components; Timkey & van Schijndel 202T: Puccetti et al., 2022). we correct by an estimate of the expected magnitude of the loss gradient with respect to each model parameter. Formally, for each parameter z 8, 0))²) of the gradient with respect to x. These estimates are used by common optimizers such as Adafactor (Shazeer & Stern (2018) and Adam (Kingma& & Ba 2015), and as such the gradients corrected by V can be seen also as more faithful the model's training process (Pruthi et al. 2020). We use the estimates of V computed by Adafactor, which can be efficiently applied by element-wise multiplication.
this an estimate of the second moment IE_ (VVV.(VVLLs
Notably, dividing by the square root second moment vV is equivalent to using only the di- Newton Hessian approximation (gradient autocorrelation matrix R around zero; Sagun et al. (0188) from previous work ($2). Unlike TRAK (Park et al., 2023 Engstrom et al. 2024), which corrects gradients by the autocorrelation after random projection, the opti- mizer second moment estimates allow us to apply the correction per parameter before random projection, enabling more granular correction of individual outlier components.
agonal of the
Random projection: : Pa E dd : As in TRAK (Park et al. 2023 Engstrom et al. 2024) and as described in the original TracIn paper (Pruthi et al. 2000), we use Gaussian random projection to reduce the dimensionality of the full model gradients. To improve projection efficiency, we use
3

---

Published as a conference paper at ICLR 2025
the two-sided projection from Pruthi al. 2220)) which is equivalent to the recently proposed
LoGra method of dimensionality reduction (Choe et al. 2024; $A.2.2). This approach contrasts with GrthSer et al. (202)), who use either un un-projected model gradients or query-specifie low-rank reductions. Un-projectedd model gradients are too large to store for all pretraining examples, and query-specific approximations require that all pretraining example gradients he re-computed if new query is considered. Random projections allow the projected gradient to be computed and saved exactly once per pretraining example, usable for all future retrievals. We use projection dimensionality d = 216, but we experiment with lower dimensionalities in $5.2
Hessian approximation R E Rdxd; We follow Park et al. (202)) and Engstrom et al. (2024)
in using the autocorrelation matrix R T of per- per-example gradients as a Gauss- approximation to the loss Hessian. We compute and apply this after optimizer state correction and random projection. For efficiency, we enforce a block-diagonal structure ($/..2.3) which allows us to efficiently compute R- (Eq. 2). However, using this method, we still find that retrievals suffer from common proponents that largely mimic the task template. Thus, departing from previous work, consider mixing approach where the matrix Raain estimated from pretraining example gradients is mixed with Roval derived from evaluation example gradients:
R Xeeal + (1 Ruain
This mixture allows R 2 downweight maghittae components that are common for eval- uation examples in a task, such as components corresponding to the task template. We select the mixing parameter such that the top ~ 1000 task-specific gradient components (out of 65K) are downweighted (details in A.233). For T-REx closed set experiments ($5), we use À = 0.90; for
C4 open set experiments C we use À 0.99 H
Unit normalization: To compute cosine similarity, we unit normalize both input vectors in Equation I. as in Akyurek et al. (2022), Han & Tsvetkov (2022). Choe et al. (2024). and Xia al. (2024). This reduces the effect of outlier training examples that have high overall gradient magnitudes (Barshan et al.- 2020 Han & Tsvetkov (022) and thus appear as common pro-
ponents before unit normalization. Unit norming is equivalent to l-relative influence (Barshan)
et al. 2020), which identifies training examples that maximize the query example loss change while constraining the overall loss change.
Using the dot product of unit normed vectors G=(=) (Eq. 2 to quantify the influence (Zm, =g) of each training example 2m on query 2g. are able to retrieve the top k training examples for 2g. We refer to these top-ranking examples as proponents of 2g (Pruthi et al. 2020 In the spirit of TracIn and TRAK, we refer to our proponent ranking method Track or "TrackStar".
4 MODELS, DATASETS, AND METRICS
We apply TrackStar to identify proponent examples for factual predictions in LLMs up to 8B param- eters. Unless otherwise specified, we first build an index of the projected gradients for all candidate examples of interest (e-g. the pretraining set), then compute the Hessian approximation R using a single pass over the data. For gradient- based methods, we perform exact scoring between each query and all candidate examples, with no need for lexical pre-filtering (c.f. [Akyurek et al., 2022] Park et al. 2023: Grosse et al. 2023). All elements of this approach have compute cost linear in model and dataset size ($2).
4.1 MODELS
We pretrain decoder-only language model on two epochs of English C4 (Raffel et al. 2020) for three model sizes: 154M, 1B, and 8B parameters, using the architecture described in Chowdhery al. 2023). Using our tokenizer, English C4 consists of 160B tokens across 365M examples (short passages and documents). Model details are in $A.3) We focus primarily on the 8B model, but we study the smaller models for dimensionality ablations (55. and scaling ($5.).
Descrihed in A.2.3 C4 requires larger À hecause the pretraining example gradients tend to he larger than those for T-REx sentences, due longer sequence lengths.
4

---

Published as a conference paper at ICLR 2025
4.2 FACT TRACING DATASET
We focus on identifying proponent examples for factual predictions in LLMs. For factual recall prompts, we use the filtered T-REx dataset from KILT (Petroni et al. (02T). which consists of entity- relation-entity triples such as (Carleton College, country, USA). We merge the KILT dataset back into the original T-REx dataset (Elsahar et al., 2018), obtaining a dataset of facts with corresponding entity IDs, surface form aliases, and Wikipedia abstract sentences containing each fact For each fact. there are an average of 2.5 "ground gground-truth" entailing sentences out of 19.6M sentences total. After filtering out ambiguous prompts (e.g. multiple correct target entity IDs), we have total of 1.2M fact triples covering 97 relation types. Following Kandpal et al. (023), we also annotate facts with the count (fact frequency) of pretraining examples from C4 that mention both entities.
We manually write natural language templates for all 97 relation types, so that a left-to-right LLM can predict the target fact completion, e.g. "Carleton College located in the following coun- try: USA". We mark predictions as correct using string matching after lowercasing and stopword removal, considering multiple possible correct surface form aliases as provided by KILT. Random chance accuracy on this task is close to 0% due to the large number of possible completions, our 154M, 1B, and 8B models achieve 16.3%, 28.0%, and 32.4%, respectively. For factual queries in all following evaluations, we use a fixed sample of 5415 facts drawn randomly from this set but bal- anced for fact frequency and down-sampling common target entities (such as "USA" or "English "), and limiting facts that are incorrectly predicted by all models. Each query contains an input prompt and ground truth target text. Dataset details and pre-processing are described in $A.4
4.3 EVALUATION METRICS
After retrieving proponents for a given query, we consider traditional fact tracing metrics (MRR and recall, altribution; Akkurek et al.) (022) and a tail-patch metric that quantifies the effect of proponents on the model itself (influence).
Attribution metries: MRR and Recall@1N 10. These measure whether we retrieve examples that logically support (entail) a fact. For T-REx evaluations (55). entailing Wikipedia abstract sentences are annotated in the dataset ($ (4.2). For C4 evaluations ($6). because passages are not annotated, we use an entailment model to score whether a candidate passage contains factual information support- ing the query. Specifically,
a fine-tuned T5 1 1B (Raffel et al. 2020) model trained on ANLI (Nie et al. 2020) and synthetic data as described in Gekhhman et al. (2023). Because C4 passages can be up 2048 tokens, we split the input C4 passages by sentences using regex matching. and take the maximum entailment score using a sliding window of three sentences as the input premise and the facttull query as the hypothesis. We mark a proponent as entailing when this entailment score 2 0 9 For MRR, we compute the mean reciprocal rank of the top-ranked "entailing" proponent for each fact. For recall@@ 10, we compute the proportion of facts for which an entailing proponent appears in the top 10 proponent retrievals.
Influence metric: incremental training probability increase (tail-patch). Both of the previous metrics assume that the top-ranked proponents should be those that entail the target fact. However, these proponents are not necessarily the examples that would most influence the model to make that prediction. Following Koh & Liang (017)) most work influence methods derives metrics experiments such as linear datamodeling score ([yas)
from leave-one-out or other ablate-and-retran et al. 2022). However, this is not tractable for full-scale LLM pretraining,
we instead estimate the additive contribution from incremental training. Starting from the final model checkpoint and maintaining all pretraining hyperparameters, patch step-on a single retrieved proponent and measure the change in probability of the target sequence. We then average this change across the top k 10 proponents for each query, and across all queries in the evaluation set. Results for different l are reported in $A.5.2
take single training step-which we call tail-
Because T-REx is automatically scraped, there is some inherent noise in the "ground truth" labels.
'aased on blinded manual annotation of 100 C4 passages marked as entailing and non- entailing respectively. we find false positive rate of 13% and false negative rate of 11 11%. Of the incorrect annotations, most (54%) are fuzzy entailments where passage implies but technically may not entail fact.
5

---

Published as a conference paper at ICLR 2025
4.4 BASELINE METHODS
We compare four ranking methods to retrieve proponent examples: BM25 Kamphuis et al. 2020). Gecko embeddings (Lee et al., 2024), TRAK (Park al. 2023 Engstrom al.) 2024), and our method, TrackStar. As described in 5.2, we also implicitly compare to methods from several pre- vious studies by ablating different correction terms from TrackStar. As a classical retrieval baseline, BM25 is pag-of-words method that ranks candidates according to query terms appearing in each document, with TF-IDF weighting ("Lucene accurate" version; Kamphuis et al. 2020). Gecko is a text embedding model trained with a two-step distillation procedure and contrastive loss (Lee al.) 2024 TRAK a gradientt -ased influence method that has been shown to perform well in previous work; the recommended version of TRAK in Park al. (2023) and Engsttrm et al. (2024) is similar to TrackStar but uses the multi-class margin gradient and a -pssssis--peccffc Hessian approximation, with optimizer second moment correction unit norm
T-REX CLOSED SET EVALUATION
We first evaluate TtackStar at fact tracing and influence on the T-REx dataset ($5.1) and we find that it outperforms prior gradient-based TDA methods ($5.2). For each query we retrieve the top 10 highest scoring candidate sentences from the set of 19.6M Wikipedia abstract sentences in T-REx.
5.1 T-REX CLOSED SET RESULTS
In T our results show that TrackStar outperforms other variants at identifying proponent exam- ples for both attribution (identifying fact-entailing proponents) and influence (tail-patching scores) in the 8B model. Notably, previous work (Akyurek et al.) 2022: Park et al. 2023) on this task con- siders only small candidate sets of 300 "distractors" per query. but we retrieve from the entire set of 19. 19.6M sentences. We observe that this makes a significant difference in performance: while our TRAK replication achieves an MRR of 0.401 when replicating the smaller setting (similar to as re- ported by Park et al. 2023), it only achieves an MRR of 0.001 (and negative tail patch scores) the full setting evaluated here. We find that this poor performance is largely due to the use of multi-class margin function gradients with correction applied per example rather than per token ($A.2.J) along with the lack of unit normalization (Table Experiment vs. Experiment 2). This result highlights that methods that perform well in small settings and classification tasks may be susceptible to noise in larger settings with LLMs.
We also find that classical, modell-agosste retrieval approaches (BM25 and Gecko) are significantly better than gradient-based influence methods at attribution (i.e. retrieving proponents that entail a given fact; MRR and recall; Table D]. Despite this, proponents from these classical methods still have much less influence on model predictions than proponents from influence methods. Tail- patching on proponents from BM25 (tail-patch score +0.41%) increase target fact probabilities by 2.2x less than proponents from TtackStar (tail-patch score +0.90%). Even tail-patching on "ground truth" entailing sentences from T-REx (tail- (tail-patch score +0.52%) results in much smaller probability changes than proponents from TrackStar This result suggests a distinction between attribution and influence; while classical, model-agnostic retrieval methods perform well at attribution (which we note is highly lexical and well-suited for string matching methods such as BM25; Wang et al., (202)) influence methods better predict how a proponent might affect the model and may better reflect the model's ' reasoning ($7)
5.2 ABLATION EXPERIMENTS
Correction terms: In Table I we conduct ablations (labeled by experiment numbers) to determine the source of improvements for TrackStar over other gradient-based influence methods. Many of these ablations are close or equivalent to methods from previous work. For example, Experiment
'ee use the d = 768 embedding model textembeddi -eg.teee@ddigggkekkk000000000000 available on Google Cloud. 'TRAK has one additional difference, that multiplies scores by Q 1- p. where p is the probability of the candidate sequence, averaged over tokens. We find Q ranges from 0.6 to 1.0 across candidate sequences, and it has little effect on MRR, recall, and tail-patch scores. Theoretical motivations are described in 1 This echoes results from Park et al. (023) for fact fine-tuning.

---

Published as a conference paper at ICLR 2025
Exp. Optim.
R
Unit norm
MRR
Recall 10
Taic-patch
T-REx gold
Gold
Gold
+0.52%
BM25
0.592
0.773
+0.41%
Gecko
0.620
0.794
+0.31%
TRAK
0.001
0,001
-0022%
0.064
0.114
+0.35%
0.266
0.358
+0.65%
0.290
0.399
+0.85%
0.300
0,413
+0.71%
0.295
0,406
+0.87%
TrackStar
Mixed
0.365
0.496
+0.90%
Table 1: Results on T-REx closed set evaluation for the 8B-model (5 Our setup is significantly
more difficult than prior work, retrieving from all 20M T-REx abstract sentences rather than 300
"distractors" per fact. Classical retrieval results are reported in the top section for comparison. We note TRAK uses multi-class margin function gradients rather than loss gradients; details in A..2.
0.35
Bb
Teippatch preb. mncerre Hat
0.35
0.30
1b
154m
0.25
1.0
0.30
0.8
0.6
154m
0.20
0.05
0.00
25
25
1408)
(605)
(1505]
(S208)
Projection dimarsions
Projection dimensions
C4 upochs (rakking takens))
Figure 2: Left, center: attribution (MRR) and influence (tail-patch) scores as a function gradi- projection dimensionality d for different model sizes ($5.2). Right: attribution (MRR) scores throughout pretraining for different model sizes. As models improve, TrackStar influence becomes more similar to attribution (higher MRR; 5.3)
corresponds closely Pruthi al. 2020). without the use of multiple checkpoints, while Experi-
ment corresponds to Han & Tsverkov 2022) and Han et al. (2023)
Experiment 3 is the method of
Choe et al. (2024). and Experiment 4 corresponds to Akyurek et al. 2022 and Xia et al. (024 8
First, we find that unit normalization is key to good MRR, recall, and tail-patch scores (Experiments 1-2; also verified by ablating unit normalization from other experiments). Optimizer second mo- ment correction provides additional improvement, similar to including the Hessian approximation R (Experiments 3-4); intuitively, both methods perform Hessian correction by approximating the Hes- sian with gradient variance terms ($3) Using optimizer correction and R together produces fairly similar results to either optimizer correction or alone (Experiments 3-5). any case, the mixed task-specific Hessian approximation R (downweighting common gradient components for the fact completion task; (3) provides substantial improvements, particularly for MRR and recall.
In fact, theoretical basis, it may be somewhat surprising that unit normalization and task specific Hessian approximation improve tail-patch scores at all. These two normalizations encourage the scoring method to focus proponents specific to the target fact by downweighting proponents that affect loss overall (or loss for the overall task). These downweighted proponents might still be expected have high tail-patch scores (large effects target fact probability), despite their significant effect on loss overall, because tail-patch scores do not constrain the overall loss change induced by proponent. The fact that these downweighted proponents actually have lower tail- patch scores (i.c. lower tail-patch scores for proponents before unit normalization and task-specific Hessian correction) indicates that these corrections actually have an overall denoising effect. Despite their motivation based on maximizing target influence under overall loss change constraints, these corrections actually find proponents that maximize target influence even without such constraints.
Projection dimensionality: Higher gradient projection dimensionality d results in higher fidelity representations, but both the memory required to store projected gradients and the retrieval cost to
Xia al. 2024 also normalize Adam first moment estimates, but these are available in Adafactor.

---

Published as a conference paper at ICLR 2025
Method
MRR
Recall@ 10
Tail-patch
BM25
0.687
0.845
+0.83%
Gecko
0.636
0.826
+0.54%
GRADIENT DOT PRODUCT
0.003
0.015
+0.04%
GRADIENT COSINE
0.252
0.393
+1.95%
TrackStar
0.338
0,515
+2.11%
Table 2: Results retrieving proponents from all of C4 (.6) GRADIENT COSINE ablates the Hessian approximation from TrackStar, and GRADIENT DOT PRODUCT further ablates unit normalization.
compute dot products increase (linearly) with d. To balance these considerations, in Figure 2 (left, center), we plot how attribution (MRR) and influence (tail-patch) scores improve with projection dimensionality for TrackStar ? For smaller models (154M and 1B parameters), scores have begun to plateau around
216 Although the 8B-parameter model has not yet plateaued, we restrict our experiments to d 216 due to memory limitations, the index of 365M C4 examples at d = 2 is already 87TB size ($6). Furthermore, we note that our factual knowledge task likely requires large number of dimensions due the inherent difficulty of compressing individual facts about large number of entities; fewer dimensions may be sufficient for other tasks.
5.3 INFLUENCE APPROACHES ATTRIBUTION
In $5.1 we demonstrated distinction between attribution (identifying proponents that entail fact, quantified by MRR and recall) and influence (identifying proponents that influence the model pre- diction, quantified by tail- tail-patch scores). We find that TrackStar attribution scores improve as models increase in both parameters and training tokens (Figure 2 right). This indicates that models im- prove, the influential examples that TrackStar identifies align more with attribution, suggesting that more capable models rely more examples that actually entail facts for factual predictions. Or course, it does not appear that these measures converge entirely; even for our largest model, tail- patching ground truth entailing proponents results in much smaller target probability changes than TrackStar proponents ($5.1). Thus it appears that as models improve, they are more likely to entailing examples to learn facts, but there are still many proponents that have large effects on the model despite non-entailment. We present a deeper analysis of these types of proponents in $7
6 C4 OPEN SET EVALUATION
While the T-REx setting is useful for controlled experiments, in practical use LLMs are trained from large, unsupervised corpora containing wide variety of passages that the model may learn from. To extend TDA to this setting, we apply TrackStar to identify influential examples for our 8B-parameter model from all 365M passages in the C4 corpus. In this scenario, our candidate all training examples that the LLM has ever seen. These candidate sequences are often much longer than T-REx sentences (up 2048 tokens), and they cover many domains rather than just Wikipedia.
(C4) consists of
For the same set of 5.4K factual queries with ground truth targets as $ we retrieve proponents using TrackStar, BM25, and Gecko. As C4 much larger corpus and is proportionally more expensive to compute gradients over, we perform a more limited set of ablations as shown in Table 2 We focus on the two corrections that had the largest effects in $5.2 mixed task-specific Hessian approximation R (ablated in GRADIENT COSINE in Table and unit normalization (further ablated in GRADIENT DOT PRODUCT in Table 2 10 Both these ablated methods still significantly outperform original TRAK in
5
As before, we quantify performance using MRR, recall, and tail-patch scores.
6.1
C4 OPEN SET RESULTS
In line with $5.1 TrackStar has better tail-patch (influence) scores than all other methods, and bet- ter MRR and recall (attribution) than other gradient methods (Table 2). Again, find that unit
Park et 2023) find that performance may decrease if projection dimensionality is large, but experiments do not appear to he close to that limit, likely due to the high dimensionality [8] of L.M gradients. !oGRADEENT COSINE equivalent to Experiment in Table and GRADIENT DOT PRODUCT is equivalent Experiment with optimizer second moment correction.
8

---

Published as a conference paper at ICLR 2025
For correct targets
For incorrect predictions
10*
10"
Entaling (MS)
Contains both entitie
Rict frequency
10
freqeen
10
Contains target entiyy only
entities
Partial match target only
10a
10a
Contains Input entity only
Partial match input only
101
10¹
of above
10
10"
0.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Proportion propanents
Propartion proponents
Figure 3: Proportions of TrackStar proponents (top 10 per query) retrieved from C4 that entail prediction, contain both entities, or contain only one entity for a model prediction.
normalization (cosine) is critical to performance, and the addition sask-specific Hessian approxi- mation R in TrackStar further improves results. In particular, lack of unit normalization often leads to retrieval of long irrelevant examples, as these tend to have large gradient magnitudes.
Also in line with $5 BM25 and Gecko perform much better than TrackStar for attribution (MRR 0.687 and 0.636 vs. 0.338). Still, we note overall high MRR and recall scores given the difficulty of the task: for 51.5% of facts, TrackStar retrieves an entailing example in the top 10 out of 365M C4 examples. Additionally, TrackStar performs over 2.5x better than BM25 and Gecko at retrieving examples that influence model predictions (tail-patch score +2.11% % vs. +0.83% and +0.54%). This reinforces the distinction between attribution and influence: examples that entail facts-sure such as proponents retrieved by BM25 are often not the examples that most affect a model's predictions.
7 HEADROOM ANALYSIS
In $5 and we showed that TrackStar makes significant progress towards attribution for LLM pretraining, outperforming other methods at retrieving influential pretraining examples, but it still falls short of classical baselines on the fact tracing task (attribution). To better understand why, we look in more detail at the types of proponents retrieved by our method. We find that much of the headroom can be explained by proponents which reflect priors or partial matches, multi-hop reasoning, or idiosyncrasies of the fact tracing task such as ambiguous entities or alternative correct answers. Examples of these proponents are included in Table A.T and Table A.2 The full set of proponents can be viewed at fnttps:: //g //g/tth..b.. cwm/pa wwwe//peettrrree ining-tda including proponents for additional evaluation tasks as described in $A.6
Priors and partial matches: In an ideal world, model predictions would be informed by examples that fully describe (entail) that fact (Chang et al. 2024), and these would appear as the top propo- nents. However, the model's output also be informed by priors, such the probability that the answer to any question about language being " "English" or city being "New York". In Figure B (left), examine the distribution of proponents (top 10 per query) from TrackStar based on their relation to the query: (1) full entailment (as in (6). (2) containing the input and target entity but without entailing the fact, (3) containing only one of the entities, (4) containing neither entity. The latter three categories are estimated by matting and we also consider partial-matches where the proponent matches at least one non-stopword of an entity (e.g. a last name). Additionally, we stratify these categories by how frequently the fact appears in the pretraining corpus.
We find that majority of proponents fall into one of the full- or partial-mate match categories, with close to 40% of proponents entailing the query on facts and 40% more with at least a partial match to one of the entities. While is not clear why partial matches sometimes score more highly than entailing examples, it is plausible that especially when the target entity is mentioned, these examples would contribute to the model's prior P(y) of generating the target string; 62% of proponents contain at least partial match to the target entity. For less common entities (e.g. less frequent facts), partial matches appear more frequently, such matching the first name in "Alemayehu Shumye" in Table A.I These may represent fallback reasoning, where the model does not know about the specific entity but makes a guess based on statistical associations, such between names and nationalities.
(10² to 10xh occurrences in C4)

---

Published as a conference paper at ICLR 2025
Interestingly, observe drop in entailing and antity-matchinn proponents for the highest fre- quency facts (2 10a occurrences). In part, this appears due to data artifacts, as some examples in T-REx are string matches for common nouns, such as "City is located in the following country: United States", and is unclear how a model should interpret these queries. Additionally, there may be saturation effects (Pruthi al., 2020), where these common facts (such as that Munich is city in Germany in the "Munich Symphony Orchestra" example in Table A.1) are learned early in training, and gradients at the final checkpoint are diminished.
Multi-hop reasoning: We also observe some cases where there is multi-step reasoning path be- tween a proponent text and the query. For example the "Carel Steven Adama van Sheltema" example in Table A.I] prompt asks about a person's 's native language, and the first proponent pas- sage states that they were born Amsterdam. While not not strietly entailing, this does provide support for the model to plausibly guess that the language is "Dutch".
Noisy proponents: However, some retrieved proponents are simply noisy. In relatively small number of cases, examples are entirely irrelevant (e.g. "Cadillac" in Table(A.1]. In these cases, we often find that TrackStar and other gradient-basee methods retrieve long or repetitive passages which bear minimal relation to the query, or which repeat the target string many times (e.g. "Présent" example in Table ; this may also reflect priors as discussed above). We suspect that this may be because training examples have gradients aggregated across an entire passage. Repeating statements that are tangentially relevant can add up to a substantial similarity score, and if better examples receive lower scores due to saturation or other effects, these distractors may rank top proponents.
7.1 DEBUGGING INCORRECT PREDICTIONS
Above, examined proponent retrievals only for ground truth targets, even when the actual model prediction would have been incorrect. In model-debugging scenario, however, it may be useful to attribute the incorrect predictions of the model so as to understand where this behavior was learned Grosse et al.) 2023] Nguyen et al. (024) and uncover mislabeled examples or other issues with the training corpus.
Figure 3 (right) we consider the subset of 1592 queries from our evaluation set that the 8B model gets wrong, we retrieve proponents from C4 using TrackStar for the (incorrect)
model prediction, and we categorize the proponents using the same scheme as 7
We observe that substantial fraction (28.5%) of these answers actually do have entailing passages (7.1% of all proponents in Figure 3 right), and on closer inspection, we find that many of these correspond to alternative correct answers. For example, for "Ricky Gervais works as: comedian" the model predicts "actor" " (he is both), and for "Victoria Park is located in the following country: Australia" " the model predicts " "United Kingdom" (there is a Victoria Park in both). However, much larger fraction of proponents consist of partial matches, suggesting that when the model is wrong, it is often making guess based partial information, e.g. using of the reasoning strategies or priors described in $7
8 CONCLUSION
In this paper, we explore data attribution (influence) methods at the scale of C4 pretraining for an 8B-parameter LLM, pushing the frontier of TDA capabilities closer to full pretraining attribution for modern LLMss Our best method, TrackStar, outperforms previous gradient- methods both retrieving examples that entail a fact (attribution) as well as examples that influence model pre- dictions (influence). Despite this, we find that classical, model-agnostic retrieval methods such as BM25 still perform better on attribution metrics. While this may be partially due to the highly lexi- cal nature of the fact tracing task (Wang et al.] [021] Akkurek et al.] 2022). it also demonstrates that attribution and influence may not always be fully aligned, both due to headroom in the method and the fact that different types of non-entailing examples can be influential on a model's prediction. We do find, however, that influence and attribution are more closely aligned as models improve. This suggests that TDA results may even align further with human intuition for newer generations of LLMs, potentially enabling practical usage of this technique to debug model predictions and better understand the connection between training data and model behavior.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Ekin Akyurek, Tolga Bolukbasi, Frederick Liu, Binbin Xiong, Ian Tenney, Jacob Andreas, and Kelvin Guu. Towards tracing knowledge in language models back to the training data. In Findings of the Association for Computational Linguisties: EMNLP 2022, pp. Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL
Abu
https://acle anthhhlllantttcccoaa low..lo..............o.....o.... (tre:22finniiggse Findings-emt 1p. 180
Elnaz Barshan, Marc-Etienne Brunet, and Gintare Karolina Dziugaite. RelatlF: Identifying explana- tory training samples via relative influence. In Proceedings of the Twenty Third International Conference Artificial Intelligence and Statisties, volume 108, pp. 1899-1909, 2020. URL.
htttps://proceedi ngs ml wtt.g.ss...............p
press http://////aaaaaa//a0
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: Reasoning about physical commonsense in natural language. In AAAI 2020, pp. 7432-7425 2020. URL
https://arxiv.. wwtp://wwwwggaaa///...... httpss//araav..rrrra////111111.1
Hoyeon Chang, Jinho Park, Seonghyeon Ye, Sohee Yang. Youngkyung Seo, Du-Seong Chang, and Minjoon Seo. How do large language models acquire factual knowledge during pretraining? arXiv, 2024.
Tyler A. Chang and Benjamin K. Bergen. Language model behavior: A comprehensive survey. Computational Linguistics, 50(1):293-35 2024.
11504
Sang Keun Choe, Hwijeen Ahn, Juhan Bae, Kewen Zhao, Minsoo Kang. Youngseog Chung, Adithya Pratapa, Willie Neiswanger, Emma Strubell, Teruko Mitamura, Jeff Schneider, Eduard Hovy, Roger Grosse, and Eric Xing. What is your data worth to GPT? scale data valuation with influence functions. arXiv, 2024. URL httpt:tt/////. /tttpp:://axivv.. wwt.oo////aa.o.. 2405 13954
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tssyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito. David Luan, Hyeontaek Lim, Barret Zoph. Alexander Spiridonov. Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, An- drew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowyez. Erica Mor- eira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. PaLM: Scaling language modeling with pathways. Jour- nal of Machine Learning Research, 2023, URL httpt::t////ttttt /ttp://aavvv..//.. org/ abs 2204 02311
Hady Elsahar, Pavlos Vougiouklis, Arslen Remaci, Christophe Gravier, Jonathon Hare, Frederique Laforest, and Elena Simperl. T-REx: A large scale alignment of natural language with knowledge base triples. In Proceedings of the Eleventh International Conference on Language Resources and Evaluation (LREC 2018), Miyazaki, Japan, May 2018. European Language Resourees Associa- tion (ELRA). URL httpp: http:///laaaatthoooo/o/ org/ 1544
Logan Engstrom, Axel Feldmann, and Aleksander Madry. DsDm: Model-aware dataset selection with datamodels. In Proceedings of the 41st International Conference Machine Learning, pp. 12491-12526. PMLR, 21-27 Jul 2024. URL lhtttt:t////ttt htttps://proceeiingg .mlr .tt.eingsmmm.pp.sss/ press/ v235/engstrom2e4 html
Zorik Gekhman, Jonathan Herzig, Roee Aharoni, Chen Elkind, and Idan Szpektor. Teacheacher: Learning factual consistency evaluation with large language models. In Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pp. 2073-2070, Singapore, December 2023. Association for Computational Linguistics. URL https:// acl antho wwt.//22..0....... emn pp-ma n. 127

---

Published as a conference paper at ICLR 2025
Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamilé Lukošiūté, Karina Nguyen, Nicholas Joseph. Sam McCandlish, Jared Kaplan, and Samuel R. Bowman. Study- ing large language model generalization with influence functions. arXiv, 2023. URL https: llaaxiv.org/aaain 2308 03296
Kelvin Guu, Albert Webson, Ellie Pavlick, Lucas Dixon, Ian Tenney, and Tolga Bolukbasi. Simflu- ence: Modeling the influence of individual training examples by simulating training runs, arXiv, 2023. URLbttts. http///axvvvvv.. org. /a///ab//// 2tt///333..111
Xiaochuang Han and Yulia Tsvetkov. Influence tuning:
Demoting spurious correlations via in- stance attribution and instance-driven updates. In Findings of the Association for Computational Linguistics: EMNLP 2021, pp. 4398 4409, Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. URL httpss//aac.aatbbooooo.o////////
Findings
www.unnggsemm.....mmm emnlp....p......m 374
Xiaochuang Han and Yulia Tsvetkov. ORCA: Interpreting prompted language models via locating supporting data evidence in the ocean of pretraining data, arXiv, 2022. WRRLhttps:
org/abs/2205 2205. 12600
Xiaochuang Han, Daniel Simig, Todor Mihaylov, Yulia Tsvetkov, Asli Celikyilmaz, and Tianlu Wang. Understanding in-context learning via supportive pretraining data. In Proceedings of the 6l Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pp. 12660-12673, Toronto, Canada, July 2023, Association for Computational Linguisties. URL
https: attps::/acllantholooy.coo org
wwwpor/g22020. 202 acl long. 708
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack W. Rac, and Laurent Sifre. Training compute-optimal large language models. In Advances in Neural Information Processing Systems, 2022. URL https::////// wttp//twwiii......aaaaoooo abs http://bb////.... 220p.22233...22 15556
Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclere, and Aleksander Madry. Data- models: Predicting predictions from training data. arXiv, 2022. URL hwwphhtps///aee ptttpp://aaxiv.orgg abs/2203 00622
Chris Kamphuis, Arjen de Vries, Leonid Boytsov, and Jimmy Lin. Which BM25 do you mean?
A large-scale reproducibility study of scoring variants. In Advances in Information Retrieval: 42nd European Conference on Research, ECIR 2020, Lishon, Portugal, April 14-17, 2020, Proceedings, Part 11, pp. 28 34, Berlin, Heidelberg, 2020. Springer-Verlag. URL https:/7
doi. org.orggg/o.ooo. 10. 1007 978 3- 0200000--455 45442-5. 5_
Nikhil Kandpal, Haikang Deng, Adam Roberts, Eric Wallace, and Colin Raffel. Large language models struggle to learn long-tail knowledge. In International Conference on Machine Learning, 2023. URL htttttts////rtttaa www.axwww.....aaaaoooooa abs / 01221.00411
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. URL. http:tts///axxii.orr////////////
6980
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In Proceedings of the 34th International Conference on Machine Learning, pp. 1885-1894, 2017.
Utthhttps:/ /atp:i:://///vii////a//sasaa .rg/ abs/ 1703. 04730
Taku Kudo and John Richardson. SentencePiece: A simple and language independent sub- word tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Con- ference Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66-71, Brussels, Belgium, November 2018. Association for Computational Linguistics. URL https://aaes http://cc..aattllllooooooo//e/ee
12

---

Published as a conference paper at ICLR 2025
Jinhyuk Lee, Zhuyun Dai, Xiaoqi Ren, Blair Chen, Daniel Cer, Jeremy R. Cole, Kai Hui, Michael Boratko, Rajvi Kapadia, Wen Ding, Yi Luan, Sai Meher Karthik Duddu, Gustavo Hernnndez Abrego, Weiqiang Shi, Nithi Gupta, Aditya Kusupati, Prateek Jain, Siddhartha R. Jonnalagadda, Ming-Wei Wei Chang, and Iftekhar Naim. Gecko: Versatile text embeddings distilled from large language models. arXiv, 2024. httppsss/////xiiv.ooo abpprr/aaa//240..... 2203.22327
Maximilian Mozes, Tolga Bolukbasi, Ann Yuan, Frederick Liu, Nithum Thain, and Lucas Dixon. Gradient-based automated iterative recovery for parameter-efficient tuning. arXiv, 2023, URL
https:///rxxit org/ als
Elisa Nguyen, Johannes Bertram, Evgenii Kortukov, Jean Y Song, and Seong Joon Oh. Towards research researach in training data attribution for human-centered explainable AI. arXiv, 2024.
user-focused
URL wtttttttttaxiiiiis www.zii..rrgoo......o abs/ 16978
Yixin Nie, Adina Williams, Emily Dinan, Mohit Bansal, Jason Weston, and Douwe Kiela. Adversar- ial NLI: A new benchmark for natural language understanding. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 4885-4901, Online, July 2020. org.:020000g/00000
Association for Computational Linguistics. URL
acl 1l1-main.444 441
Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry. TRAK: Auributing model behavior at scale. Proceedings of the 40th International Con- ference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, m.r.press/
pp. 27074-27113, PMLR, 23-29 Jul 2023. URL
22G2/park23c.oo html
Fabio Petroni, Tim Rocktäschel, Sebastian Riedel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, and Alexander Miller. Language models as knowledge bases? In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International
Joint Conference on Natural Language Processing (EMNLP-IJCNLP,, pp. 2463-2473, Hong
Kong, China, November 2019. Association for Computational Linguistics. URL https:
//aclantholoon /tq//19-12555
Fabio Petroni, Aleksandra Piktus, Angela Fan, Patrick Lewis, Majid Yazdani, Nicola De Cao, James Thorne, Yacine Jernite, Vladimir Karpukhin, Jean Maillard, Vassilis Plachouras, Tim Rocktäschel, and Sebastian Riedel. KILT: A benchmark for knowledge intensive language tasks. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2523 2544, Online, June 2021. Association for Computational Linguistics. URL https://///.... //acl anthohooho hogyy. org/2021.
naacl ma mm..i...000 200
Garima Pruthi, Frederick Liu, Mukund Sundararajan, and Satyen Kale. Estimating training data influence by tracing gradient descent. In Advances in Neural Information Processing Systems, 2020. URL/attps://arxe wwtp:oo/gg/..../.... 08484
Giovanni Puccetti, Anna Rogers, Aleksandr Drozd, and Felice Oell'Orletta, Outlier dimensions that disrupt transformers are driven by frequency. In Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 1286-1388 1304, Abu Dhabi, United Arab Emirates, December 2022.
Association for Computational Linguistics. URL
findingg-eemnlp 99.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(1), 2020. URL https: /ttps://arxiv.
orq/abs/1910.100 10683
Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor, Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthorne, Aitor Lewkowycz, Alex Salcianu, Mare van Zee, Jacob Austin, Sebastian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jasmijn Bastings, Jannis Bu- lian, Xavier Garcia, Jianmo Ni, Andrew Chen, Kathleen Kenealy, Jonathan H. Clark, Stephan Lee, Dan Garrette, James Lee-Thorp, Colin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma,
13

---

Published as a conference paper at ICLR 2025
Alexandre Passos, Jeremy Maitin-Shepatte Noah Fiedel, Mark Omernick, Brennan Saeta, Ryan Sepassi, Alexander Spiridonov, Joshua Newlan, and Andrea Gesmundo. Scaling up models and data with t5x and seqio. arXiv, 2022. Utthttts:: http:p:::/axxivoooo
abs/ 2203.
Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S. Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series, 2011. URL 2ttp::/cc..aa.a..../////////////
Subhro Roy and Dan Roth. Mapping to declarative knowledge for word problem solving. Trans-
actions of the Association for Compulational Linguistics, 6:159-172, 2018. URL https:
//aclanthology. wwtg/ii-g///1//1
Laura Ruis, Maximilian Mozes, Juhan Bae, Siddhartha Rao Kamalakara, Dwarak Talupuru, Acyr Locatelli, Robert Kirk, Tim Rocktäschel, Edward Grefenstette, and Max Bartolo. Procedural knowledge in pretraining drives reasoning in large language models. arXiv, 2024. URL htLhttps:
//arxiv. abs/ /2111. 12580
Levent Sagun, Utku Evei, Ugur Guney, Yann Dauphin, and Leon Bottou, Empirical analysis of the hessian of over pee-parametrized neural networks. arXiv, 2018. URLLttpp://axxxvvorrg
abs/1706. 04454
Andrea Schioppa, Polina Zablotskaia, David Vilar Torres, and Artem Sokolov. Scaling up influence functions. In AAA/-22, 2022. https: //tppp::///aiv.../// org/ wwwporggg/bbbbaaeeeeeeee abs/ 212..11211 03052
Noam Shazeer and Mitchell Stem. Adafactor: Adaptive learning rates with sublinear memory cost. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 4596-4604. PMLR, 10-15 Jul 2018. URL httpp:://proceedingg: www.m.rrrreeee press wttp:///eaaaeee...aaaaaa. html
William Timkey and Marten van Schijndel. All bark and no bite: Rogue dimensions in transformer
language models obscure representational quality. In Proceedings of the 2021 Conference on
Empirical Methods in Natural Language Processing, pp. 4527-4546, Online and Punta Cana,
Dominican Republic, November 2021. Association for Computational Linguistics. URL https:
V/Aclantholoog.orrr//2/00 ott.lpyo.......222......mmmmmm.. www.:/min.mmmm.......mmmm.......
Shuai Wang, Shengyao Zhuang, and Guido Zuccon. BERT-based dense retrievers require interpola- tion with BM25 for effective passage retrieval. In Proceedings of the 2021 ACM SIGIR Interna- tional Conference on Theory of Information Retrieval, ICTIR 21, pp. 317-324, New York, NY, USA, 2021. Association for Computing Machinery. URL https //wp:///cc.o///ooo/ /tt.:://///....//100.1....
3471158. 3472233
Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. LESS:
Selecting influential data for targeted instruction tuning. In International Conference on Machine Learning (ICML), 2024. URL /ttppp://axxii.orrrt http:///wxxx...ott////////////// /ttp///b222222.433333
Chih- Kuan Yeh, Ankur Taly, Mukund Sundararajan, Frederick Liu, and Pradeep Ravikumar. First is better than last for language data influence. In Advances in Neural Information Processing Sys- tems, volume 35, pp. 32285-32298, 2022. URL littps:// //axxivivvvv/////////aa///aa//aa v. org/ hbtp://..//222.2//////// 11844
A APPENDIX
A.1 QUALITATIVE EXAMPLES
Examples of top passage retrievals from TrackStar are shown in Tables A.I and A.2 below.
14

---

Published as a conference paper at ICLR 2025
Example proponents from TrackStar:
Présent is in the following language: - English (incorrect, groundtruth: French)
Proponent retrieval #3 (non-entailing):
Sanry, estry only available Deutsch, Eiglish and All Languages Sony, entry only available Noderlunds, Slovenina, Franais, Polski. English and All Langunges, Sorry, this entry is only available in English and All Langugess Sorry, this entry is only available in English and All Languages. Sorry, this anty is only available in English and All Languges. Sorry, this enty is only weailsble in Polski, Enghish and
Victoria Park is located in the following country: - United Kingdom
(incorrect", groundtruth: Australia)
Proponent retrieval #1 (entailing);
Victoria Park the noothem part of Londons bast end hectares mendows, trees and formal gardens set around two lakes, The Regents Canal runs along the south and west sides of the purk and is pleasant to walk along especially on summer day. The park was donated to the people by Queen Victoria, was the fins public purk and opened the public 1845, There number of good bars restaurants the nunthrr edge of the park on Grove Road.
Ricky Gervais works as: - actor (incorrect*, groundtruth: comedian)
Proponent retrieval #1 (entailing):
Guing to soe Ricky in Turento fomerrow night re- blog felt uppropriate. For those not famillar with him. Ricky Gervais is beilfiont Brilish actor. dinctor and writer, hest known for portraying David Reent in his original BBC veries The Office. past hos of the Golden Giloees, the oflen contovariial comedian knows raising the ite Hellywood A-listers with blunt zingers and one-liners,
the director Dr. McPhoe Night at the Museum with Ben Stiller, and later stumbled upen him in The Office on Netlix. quickly became fin and adoned all his sibsoguent pregeers, including Deekk, An Idiot Abread and lifes to Short, all bilarious and poigrant. As I wrilc this, Gervais new movic Special Conespondnts is currently playing on Nerfix, and his feature film David Brent Life on the Road will he relaased theaters this Office
coming August
jealous! love Ricky. diseovered Besides his animal rights activism, love his omspotemess about atheiom in The Invention of ying, he cleverly made fun the origins. of religion. Have a good time somorrow! Thank you, looking forward to it!!
like his
while
catching up with movies.
City located in the following country: - United States
(incorrect*, groundtruth: Finland)
Proponent retrieval #1 (entailing):
a unique xip assigned by and Time Zone is Cental (GMT -06-00); The US ZIP Code basic format comiist of five docimal minerical digits assigned to the Many Ciry, An extended ZIPe4 code, introduced in the 1980s, inclates the five digits of the Many ZIP code, a hypten, and four more digits that determine more specific location within given ZIP code in Many, Louisiana. The Many Zip Code is 71499 you want send mail Many City, you should...
Many city located Louisiana State, Sabine County
Many Code is
Munich Symphony Orchestra originated in the following country: + Germany
(correct, groundtruth: Germany)
Proponent retrieval #1 (non-entailing):
Amstendam Baroque Orchestin, conductor Teo Koopmnss. Label: Challenge Classics. Recording Date: Sepeemhos, 2009.
Cadillac was formed in the city of: + Detroit (correct, groundtruth: Detroit)
Proponent retrieval #1 :oon-entallnng):
RBG: Real Bar Gaming Need Jabher Loves... Comics! What did Happy Gilmore do next! 90' nostnigia and ridiculous theories with bost Cltire Lim, guest host Paul McCallum and special guest Josh Macuga. The podcast available on illunes, Spotify or via the emhhed below.
Alemayehu Shumye died in the city of: + Addis Ababa
(correct, groundtruth: Addis Ababa)
Proponent retricval #1 (non-entailin))
Dr. Alemaychu Marium (A) Mariam) a professor of political science at Califomiia State University. Bemandino (CSUSB), He received his Ph.f) from the University of Minnesots in 1984, and his J.D. from the University of Maryland in 1988. He serves on the board of the Center for the Study of Hate and Extremism at CSUSB. He his given human rights lecmures at the Canr Canter, Harvard Univeritt, and various law schools including the University Minnesota, American University
passage of Demoencc and Human Rights Act) the House Rejersontaives in 2007. For the sevenl years, has writtes woekly web commentury on Ethlopian humna rights and African ismes. Curreatly, be has a weekly blog at the Huffington Post and his articles are featured on Pambazuka News, New American Medis, Eihiopian Review and Medio-Medis. He hos published two volumes on American constitaional low, including American Civil libestie Editor of the Interntional Joumal of Elhiopian Studies, a leading schuhaly joumal Elhiopia.
UCLA als played central advecccy
(Fihiopia
Constituional Structues and Process (1994) and American Constittional
Rights (1998),
Senior
Alemayehu Shumye died in the city of: + Addis Ababa
(correct, groundtruth: Addis Ababa)
Proponent retrieval #2 (non-entailing):
Respunse to "Addis Abuba Bete Alemnyehu Eshate" 1 love this song. Thanks for posting it!!
Carel Steven Adama van Scheltema had the following native language: ++ Dutch (correct, groundtruth: Dutch)
Proponent retricval #1 (non-entailing):
Scheltems was bom Fahnurry 1877 Amsterdam, Norh Hollard, Netherlands Frederik Adama van Scheltema (1846-1 1999 Hendrika Lulofs (1850-1927) and died May 1924 Berges, North Hollund, Netherlands of unspecified camaes. He narried Anna Catharina Kleetitra (1 (1884-1977) 24 Oetober 1907. Ancestors are from the Nethestads.
Carel Sieven
Example passage retrievals from TrackStar, from the 8B model over the full C4 corpus. Some model predictions are marked as incorrect* because they do not match the "ground truth" from the T-REx dataset, but they are still plausibly correct due to either an ambiguous prompt or an alternative correct answer (for example, Ricky Gervais is both an actor and a comedian).
Table A.I:
15

---

Published as a conference paper at ICLR 2025
Top proponent from TrackStar for randomly-sampled queries:
Thompson River is located the following country: - Canada
Proponent retricval #1:
Blue Moon. Original Paiming by Thomesoo River Carl Strmqqiit. Carl deliberate and focused: Deliberate his desiro to point excellanee, in his quest to cam visions will generate arteork which honors ciin of life and the richness, depth and teachings of Canadian Firs Nations. Being Self taught Carl has spent much time sundying the works of renoweed Native artisams as well as the timeless iruths embudied in the olders' storios and legende. From this journey for tuih has come his desire to portrny the bolance and hamonyy that exists in the Circle of Life.
Carl's influenced by profound love and respect for
Rachael Ray is citizen of the following country: + United States Proponent retrieval #1:
Raciuel Ray Show Snack Day!!!! The Rachael Ray Show- of Day!!!! Copyright 2014 KK's Gourmet. All Rights Reserved.
Lars Johansson speaks the following language: + Swedish Proponent retricval #1:
Lars Johun been Fchnuary 15, 1849. He maniol Sofia Johansooter. was been March 15, 1847 in Fogels.
lean-Baptiste Forest was born in the city of: + Paris
Proponent refrieval #1:
Creative Jean-Haydiste Le Divelae fun turning the entire monie by Stanley Kuhrick, 2001; Space OMyysaey, into animated GIPSs gives 569
Gils he published on the platform Gipbyy through which we can discover the movie second by sccond, for fize. Althoogla the imagee are mute. he look care insegrating the subitles in each GIF.
Disneyland is named after: + Walt Disney
Proponent refrieval #1:
Robert Gerlow Gerlow City is nomed after Robert Clestaws, who 1817 and died 1911, the early 1850s. Gerlaw came to Warren County und married Marry Black. They moved Sec. 34 Spring Grove, which became a jownship in 1854.
Guslow City
Michel Sardou works as: -> actor
Proponent retrieval #1:
In late 201 2, Michel Sardou began foar-moonth "I.ss Grands Moments" tour across France, revisiting his hiegest hits in tnily energetic show. Lighting desigser Jacques Rouveyrolli used just one type spotlight on stage throughout the given was the lightnn fixtures French dhut,
Vipe: Profile somewhat bold choice that delivered all promises. A total of MAC Vipet Profiles, supplaed by the Dushow Group (who cumrailyy have 160 total), were installed, 56 equally spaced over three gantrics and 37 on the floor spread across four levels. No aradtional sporlights were used mre occurience for such show. The simple lighting plan acted the sole stage design element and provided the performance area with
effering wide variety of effects. As front house, MAC' UI Perfommaces desk for extra punch. The MAC Vipers enabled Iaoques Romveynllla design sevenal different lighting scenariss sccompany the varions songs. Using a single type of light source..
installed behand the control
Ryan Miller a citizen of the following country: - United States Proponent retrieval #1:
Posted Tuesday. January 22nd, 2019 at 8:50 pm. Ryan Miller the Vexina last and therefor is the ranked gualic. Ile also lost went perfonn one collapses in sports history. I not saying he shouldn be number one, just saying those are things to think about come drait day., About a year ago, www.nllenneyschhapeee.eo
the gold medal game questionable and got bounced the first round by the Bruis
opic
escuped from the overheated frenzy of the Summer Solsice Farade and drove up the canyon, where bezeze whisgered through the sycamores eucalyuuses lining Taro Casyon Creek down toward Modemist oulpost and its suroundinn Mediterancea landscape was nothing to knuck on besides ginlers. Some of the fins moves leamned by martial arts students is how to punch and block with the arms. Injuries often occur when lick not blocked propecty. including bruses, hartline tractures, and even broken boses. To reduce the risk of these types of injuries, student will sparting gear almost never leaves home, even to pursue the detective work that allows
grace Myers trees and oaks, Forgut about the parade. Nobody was visible, there
below. Amid the
for purtner to protect herself well. He his office wholesale jerseys his house,
expensive lifestyle, Insiead...
My Bloody Valentine 3D originated in the following country: + United States Proponent retrieval #1:
No revieus for My Bloody Valentine Complete Editiion [Limited Edition] yet. there is new movic directed by Parrick Lussier. the price for My Bloody Valentine Compleir Edition [Limited Edition] drops. there new images links available for My Bloody Valentine Complete Edition imiteed Editicn].
Happily Divorced had cast members named: + Fran Drescher Proponent retricval #1:
Happily Divoroed sta Frin Drescher is now hupply marriel! The 56 ye-yerr-old actress suid do" to Dr. Shiva Ayyadurai, 50, their beach home front a small group of friends and family, People confinss. Drescher shared the news on Twitter on Sunday with sweet soltie. Ayyadursi, who holds the putent inventing emai. Drescher Hullington Fost, "Every day as celebeation
hosted by Despak Churra. "Fran heard talk and we fell love, and Fran. Every day almooi romantic hangout with her. We're he always laghing, always enjoying ourselves." This the sccond marriage for The Nanny actress. She was married to Peter Marc Jacobson for 21 years, but diverced him in 1999, Jucobson, who is also Drescher's is producing partner, liter came out
been together since that talk."
Stadium Beyond the Stars is of the following genre: + science fiction Proponent retrieval #1:
4y; Do Androids Dieam of Electri Sheep? Reyond the posts where discuss differest aspects the genre and the many tropes and plot lines associated with Today, I talk about Kunners source material, Do Androds Tream of Flactric Shoep?? a series of posts foensed on the Blade Runner Universe. Reyond Contimu reading Reyond The Stare: Blade Kunner Overview (4 of Do Androids Dream of Electic Sheep?
Beyand The Blade Rumner Overview
of Science Fiction
Top passage retrieval from TrackStar for randomly-sampled queries with ground truth targets, from the 8B model over the full C4 corpus.
Table A.2:
16

---

Published as a conference paper at ICLR 2025
A.2 METHOD DETAILS
A.2.1 OUTPUT FUNCTION ABLATIONS
The first step in TrackStar ($3 is to compute the example loss gradient VoLoss(z, f). However, the approximation of the effect training example 2m on query 2n loss can also be approximated using a different output function f (Park al. 2023). Specifically, TRAK estimates the effect of 2m on bn terms of output function J. then converts this effect to an effect on loss by multiplying the projected and corrected gradient vectors by Q = Loos 8f
Ideally, f is chosen to be maximally linear
with respect to model parameters.
In the case of TRAK specifically, the multi-class margin function log( - corresponds to Q p-p-1 1. However, TRAK applies Q the example level, multiplying influence scores by p. where p is the mean probability over tokens in an example (using 1 rather than 1 so as to flip the sign of dot products). TRAK only applies Q the candidate examples 2m when applying Q at the example level, there is no effect on rankings if Q is applied to the query example 2q. Crucially, we find that applying Q at the example level rather than token level significantly hurts performance. MRR scores on the closed set T-REx evaluation (as in $5) drop from 0.122 to 0.001, and recall scores drop from 0. 194 to 0.001.
Assuming Q is applied at the token level, we can then evaluate different output functions f in TrackStar. To do this, we (1) take gradients for f in place of loss in Equation [ and (2) multiply each Go (z) ()) in Equation 2 by Q asss We consider three possible output functions:
Loss f = - log(p): This the output function we use in the final version of TrackStar. this case, no term is required.
Margin (multi-class margin function) f = log( p This is the output function recommended by Park et al. (2023) and Engstrom et al. (002)), based on its extension from logistic regression. In this case,
Logit f Logits,, the logit for each target token w. This tests the hypothesis that the target logit is a more linear (with respect to model parameters) and less noisy proxy for loss. In this case, =
We test these output functions with optimizer second moment correction, coon-task-specinn Hessian approximation, and unit normalization, i.e. analogous to Experiment 5 in Table [ Results are reported in Table A.3 We find that varying the output function has fairly little effect, but loss performs slightly better than other output functions.
Optim.
R
Unit norm
MRR
Recall @ 10
Tail-patch
Loss
0.295
0.406
+0.87%
Margin Logit
/
0.293
0.403
+0.87% %
0.288
0.388
+0.71%
Table A.3: Results as Table but evaluating different output functions f take gradients (Park
et al. 2023).
A.2.2 RANDOM PROJECTION
To efficiently project gradients into lower dimensionality (33), we use the two-sided random projec- tion from Pruthi et al. (2020) and equivalent to the low-rank projections in Choe al. (2024). For gradient matrix E gunxn projected into dimensionality rather than use naive projection the flattened gradients Pa Rdxma, we use two projection matrices Pa, E Rvăxm and Pa, E IR vdxn Projection matrix entries are sampled i.i.d. from (O, ta)) The resulting projection is:
Pa, WP I vavavaa
For our models. concatenate gradients into eight layer blocks. For example, the 8B 3B-parameter model has 32 layers, so we concatenate gradients for every four consecutive layers. We concatenate attention and ML.P matrices separately. We then project each concatenated matrix into d 4096
17

---

Published as a conference paper at ICLR 2025
64 x 64 dimensions using Equation 4 This results in total of 2 (attention vs. MLP) x 8 (laver
blocks) x 4096 dimensions, or 2 total dimensions. We ablate over dimensionality in 55.2 by decreasing the dimensionality per layer block.
A.2.3 TASK-SPECIFIC HESSIAN APPROXIMATION
Our Hessian approximation in $ follows the Gauss-Newtor approximation discussed by SSauun et al. (0188) and Park et al. 2023) which is based on the autocorrelation matrix of the projected gradient vectors. This is computed as = T o. where rows of 4 E RR are projected gradient vectors for individual examples. For efficiency and following previous work approximating block diagonal Hessian (Grosse et al.] 2023), we compute the autocorrelation per layer block (recall from $A.2.2 that gradient vectors are projected into 4096 dimensions per layer block). In other words, rather than R x2 our Hessian approximations are in R R²*x2 "x2"" Instead of applying R-1 directly in the inner product, we compute R. (this is casily computed from the SVD of R) and apply it separately to both the train and query vectors (Eq. 2); this allows to express retrieval as symmetric dot product as in Equation|
For the task-specifie Hessian approximation ($3), aim to downweight gradient components that are common (often high magnitude) for a given target task. For example, many tasks include a natural language template shared across all examples. To downweight these gradient components, we compute a Hessian approximation Reval computed from the target task query vectors. However, applying Reval entirely in place of Raain results in excessive upweighting of components that are rare (often low magnitude) for the target task. If there are relatively few target task examples, Reval may even be invertible. Thus, we dampen Reval as in Equation repeated here for convenience:
R AReal (1 R
We select À such that roughly the top 1000 out of 2 65536 task-specific components downweighted. Concretely, we select À such that the 1000th largest singular values of XRRval and Ramin are roughly equal. Put another way, we scale such that the (monotonic) singular value curves of XfReval and (1 (1-A) Ramin intersect at the 1000th component. We note that this requires larger when Raain computed over C4 rather than T-REx, because C4 examples have overall larger gra- dients (and thus larger Ruain singular values) due to longer sequence lengths. We set À = 0.90 for the T-REx experiments in $ and x 0.99 for the C4 experiments in $6 We also find that these values work well empirically, selecting from {0.50, 0.90, 0.99, 0.999}. Determining the optimal for a given task (i.e. the number of task-specific components to downweight) an interesting direction for future investigation.
We also note that this task-specific approach defines semantic "dimensions" that should be down- weighted (in this case, downweighting examples that support the overall task structure rather than an individual fact). When there is no a priori known task, the non-task-ppecific Hessian can be used, at slight cost to performance (Experiment 5 in Table [T. More generally, the determination of semantic dimensions that should be downweighted somewhat subjective; the eval data itself is not necessarily required for Hessian approximation, but rather any set of examples to define what semantics are less relevant.
A.3 MODEL. DETAILS
As described in $4.1 we pretrain a 154M- IB-, and 8B- 8B-parameter decoder-only language model on two epochs of English C4 (Raffel et al. 2020). All of models are implemented using T5X (Roberts et al. 2022). For all model sizes, we use the same SentencePiece tokenizer (Kudo & Richardson, 2018) trained on C4 data with vocabulary size 32K. All model sizes trained on the same shuffle of the pretraining data. We pretrain with batch size 1024 and sequence length 2048 for epochs (187K steps). The 154M, IB, and 8B models reach eval losses (log-perplexities) of 2.34, 1.99. and 1.77 respectively. Specific hyperparameters are in Table A.4
A.4 DATASET DETAILS
Our T-REx dataset is merged from KILT (2.3M facts, a further processed version of T-REx; Petroni et al. 2021) and the original T-REx dataset (11,1M facts; Elsahar al.] 2018). These
18

---

Published as a conference paper at ICLR 2025
Hyperparameter
154M
1B
8B
Layers
8
16
Embedding size
1024
2048
4096
Hidden size
1024
2048
4096
MLP hidden size
4096
8192
16384
Attention heads
4
8
16
Attention head size
256
256
256
Optimizer
Adafactor
Learning rate
0.01
Vocabulary size
32K
Batch size
1024
Sequence length
2048
Activation function
SwiGLU
Attention type
Multi-query
Position embedding
RoPE
Learning rate decay
Inverse square root
Warmup steps
10K
Dropout
0.0
Table A.4: Language model pretraining hyperparameters, following Chowdhery et (2023).
datasets consist of fact triples (input entity, relation, target entity) scraped from Wikipedia. We start from the KILT dataset because it has useful surface form aliases (different possible surface strings) for each entity, for more robust scoring open-ended text generation. We exclude en- tity aliases with less than three characters. However, the KILT dataset does not directly contain entity URIs and entailing sentences from Wikipedia abstracts. Thus we match the KILT dataset back with the original T-REx dataset using entity string matching, keeping only facts that are un- ambiguously matched between the two datasets. The original T-REx dataset contains entity URIs and secdiine-annttted entailing Wikipedia abstracts. We remove ambiguous facts that have mul- tiple correct target URIs (e.g. "France", "shares border with", "Spain, Germany, elc"), along with nine manually-identified ambiguous or noisy relation types: facet_of, is_a_l st_of, instance_of. locatediing innt the
part_of, subclass of, has part, main s sub ject, residence. The resulting dataset has 1.2M fact triples covering 97 relation types.
To obtain entailing Wikipedia sentences from the annotated abstracts for each fact, we use the sen- tence boundaries, entity spans, and relation spans annotated in T-REx. Each abstract in T-REx con- tains annotated fact triples, with corresponding input entity, target entity, and relation spans. Thus, within entailing abstract for a fact, mark individual sentence as entailing if the input, target, and relation spans all fall within the sentence boundaries. We do not include entity spans that match stop word (e.g. "she""), because these cases are generally ambiguous pronoun coreferences previous sentence, and thus the sentence in isolation does not entail the fact. However, because ab- stracts T-REx often have un-annotated entity spans, we also mark a sentence entailing within an entailing abstract if at least one surface form of each entity appears within the sentence boundaries, based on lowercased string matching; while this slightly biases our entailment annotations towards lexical matching, we observe that the sentence annotations have a large number of false negatives if we omit this step. We use our annotations of entailing sentences for fact tracing evaluations in $5
For C4 frequency counting, we use lowercased string matching. marking C4 example as relevant to a fact matches at least one surface form alias for both entities in the fact (Kandpal et al. Frequencies range from zero roughly 10 out of 365M examples in C4.
2023
Finally, because standard existing prompt templates are designed for masked rather than autoregres- sive language models (Petroni et al. 2019), manually write natural language template for each relation type. Templates all end with a colon, which we find better constrains the model to answer the fact rather than continue the sentence in another way during open-ended generation. For exam-
19

---

Published as a conference paper at ICLR 2025
ple, the template for countzy is "[entity0] is located in the following country:".. Results for other templates are reported in |
For all reported experiments, Specifically, we separate facts into six frequency buckets: 1 to 10, 10 to 102, 102 to 10", 10 to 10 10a to 10 and 10 to 10 occurrences in C4, with frequency annotations described above. We randomly sample up to IK facts from each frequency bucket. Per bucket, we restrict facts with given relation and target entity (e.g. "country". "USA") 25 examples, and restrict each target and relation overall to 100 examples. We also restrict samples that are incorrect for all three model sizes to 100 per bucket. The first five frequency buckets successfully sample IK facts each. The highest frequency bucket only samples 415 facts, because there are overall fewer of those high- frequency facts.
the same subsample of 5.4K facts balanced for fact frequency.
A.5 ADDITIONAL RESULTS
A.5.1 RESULTS FOR DIFFERENT PROMPT TEMPLATES
To verify that our results are not entirely reliant on the templates used for factual prompts, we write two additional prompts for each factual relation type. As described in $A.4 our original templates are designed to constrain model generations to produce factual completion, e.g. "/entityO] was born in the city of.". Our second set of templates removes the colon and words designed only to constrain the generation, e.g. the template "[entityO] was born in"", The last set of templates rewords the prompts such that the input entity is not at the beginning of the prompt, e.g. the template "The birthplace of [entity0] is".
Results for different templates are reported in Table A.5 In line with the results in the main text, for all templates, BM25 and Gecko perform better than TrackStar for attribution (MRR and recall), but StackStar performs better for influence (tail-patch scores over 2x higher).
Template
Method
MRR
Recall @ 10
Tail-patch
Original
T-REx gold
Gold
Gold
+0.52%
BM25
0.592
0.773
+0.41%
Gecko
0.620
0.794
+0.31%
TrackStar
0.365
0.496
+0.90%
Variation #1
T-REx gold
Gold
Gold
+0.55%
BM25
0.617
0.797
+0.57%
Gecko
0.593
0.766
+0.30%
TrackStar
0.331
0.460
+1.16%
Variation #2
T-REx gold
Gold
Gold
+0.39%
BM25
0.603
0.791
+0.22%
Gecko
0.579
0.760
+0.01%
TrackStar
0.299
0.424
+0.79%
Table A.5: Results as in Table Dutt using different templates for factual prompts.
A.5.2 TAIL-PATCH RESULTS FOR TOP-A PROPONENTS
In Table I report the tail-patch score (target probability increase from train step single proponent) averaged over the top i = 10 proponents for each fact. In Table we report tail- tail-patch scores when averaging over the top 1,3. 5, and 10 proponents per fact. As expected, lower leads to higher tail-patch scores, because only higher-ranked proponents are considered when is lower. For all k. TrackStar outperforms BM25, Gecko, and the "ground truth" entailing sentences from T-REx, in line with the results in the main text.
A.5.3 RESULTS SPLIT BY MODEL CORRECTNESS
In 15.1 our results based proponents retrieved for the ground truth target for each factual prompt. However, the model only correctly predicts subset of these facts. Intuitively, we might expect results differ based whether proponents retrieved for facts that the model predicts
20

---

Published as a conference paper at ICLR 2025
Method
==1
3=
==5
T-REX gold
+0.89%
+0.71%
+0.61%
+0.52%
BM25
+0.81%
+0.59%
+0.50%
+0.41%
Gecko
+0.94%
+0.59%
+0.45%
+0.31%
TrackStar
+1.90%
+1.40%
+1.18%
+0.90%
Table A.6: Tail-patch scores as Table but averaged over the top h proponents for different k. Results in the main text use k 10.
correctly incorrectly, even though the proponents are retrieved for the ground truth correct answer in both cases. We split these two conditions in Table A.7
We find that whether the model predicts fact correctly does not substantially affect MRR and recall, although incorrectly-predice facts tend to have slightly higher scores. This may be be- cause correctly-pret prreeccy-peee facts are more likely to have diminished or saturated gradients, because the model already assigns high probability to the correct output (Pruthi et al. [020). Indeed, incorrectly- predicted facts are much easier to tail-patch to higher probabilities (tail-patch scores; Table AA.7) than facts that the model already predicts correctly. The model- methods (BM25 and Gecko) also exhibit this effect, suggesting that our proponents are not necessarily better for incorrectly -peedictd facts; for those facts, it is simply casier push the model probability towards the correct prediction because the model does not already have high correct probability. In both cases, trends between methods are consistent.
Exp.
Optim.
Unit
MRR
Recall @ 10
Tail-patch
T-REx gold
Gold
Gold
+0.32% +1.10%
BM25
0.591 0.593
0.780 / 0.756
+0.30%
+0.66%
Gecko
0.611 0.641
0.789 0.806
+0.21% +
TRAK
0.001 0.001
0.001 0.001
0.02%
7...2%%
0.055 0.086
0.099 0.144
+0.24%
/1.611
2
0.258 0.286
0.350 0.376
+0.48%
++1.07%
3
0.282 0.312
0.387 0.428
+0.64% +11.36%
4
0.291 0.321
0.405 0.432
+0.53%
+1..16%
5
0.286 0.316
0.397 0.428
+0.65%
11.411
TrackStar
Mixed
0.358 0.379
0.489 0.515
+0.69% +1.42%
Table A.7: Results from Table 1 split into facts that the model predicts correctly vs. incorrectly (left / right). Trends between methods similar, but tail-patch scores are higher for facts that the model predicts incorrectly. Intuitively, it is easier to push the model probability towards the correct prediction when the model is not already correct.
A.6 ADDITIONAL TASKS
To facilitate further TDA research, we provide retrievals from TrackStar on the SB-parameter LLM for several additional tasks, including factual predictions, factual errors, commonsense reasoning, arithmetic, and open-ended generation. Specifically:
T-REx ground truth (Elsahar al. 2018 Petroni al. 2021): This is the task used for the main results in this paper. We use the T-REx dataset of fact triples and manually write templates, as described in $A.4 The 8B model accuracy for this task 32.4% using open-ended generation (chance close to 0.0%) TDA queries use the factual prompt as input text and the ground truth answer as target text. As described in $A.4 we use a sample of 5415 factual queries.
T-REx incorrect predictions: This uses the same set of prompts as the T-REx ground truth queries. We filter to facts that the 8B model predicts incorrectly (1593 facts), and manually remove facts that are clearly ambiguous (e.g. as in Table AI or where the model prediction does repeating sentence rather than responding
not follow the task template (e.g. just producing to the factual prompt). This results in a set of 966 coooreetly-preddcce facts. TDA queries use the factual prompt input text and the incorrect model prediction target text.
21

---

Published as a conference paper at ICLR 2025
COPA (Roemmele al. 2011): We use the 400 commonsense reasoning completions in the COPA train set. The 8B model accuracy for this task is 80.3% (assigning higher probability to the correct completion; chance 50.0%). TDA queries use the imput prompt imput text and the ground truth commonsense completion as target text.
PIQA (Bisk et al. 2020): We use the 1838 physical commonsense reasoning completions in the PIQA validation set. The 8B model accuracy for this task is 78.3% (assigning higher probability to the correct completion; chance 50.0%). TDA queries use the input prompt as input text and the ground truth commonsense completion as target text.
Arithmetic word problems (Roy & Roth 2018): We use the 1492 arithmetic word problems from Roy & Roth (2018). We note that because models are only trained on C4 (a smaller and less curated dataset than most modern LLMs), the 8B model performs quite poorly on the arithmetic word problems (1.6% accuracy using open-ended generation). However, its top pro- ponents for ground truth answers generally consist of examples with mathematical operations and similar word problems. TDA queries the input question with "Answer:" appended the input text and the ground truth answer as target text.
Simple arithmetic: We use templates to generate simple arithmetic prompts (e.g. "5+5="). We generate 500 queries each for integer addition, subtraction, multiplication, and division. We sample the first operand from [I, 100] and the second operand from [1, 10]. As with arithmetic word problems, the 8B model performs quite poorly simple arithmetic prompts (9.7% accuracy using open-ended generation), but its top proponents for ground truth answers generally consist of examples containing arithmetic.. TDA queries use the input prompt input text and the ground truth answer as target text.
Story generation: We use templates to generate story prompts for 50 manually-compiln story genres (e.g. "fantasy", "scifi", and "horror"). The templates are of the form "Genre: [genre]. [Book/Film] summary:". We generate four story summaries per genre: two book summaries (sampling temperatures 0.3 and 0.7) and two film summaries (sampling temperatures 0.3 and 0.7). We find that the resulting 200 story summaries are fairly coherent, usually consisting of a short paragraph describing generic story of that genre. TDA queries use the input prompt as input text and the generated story as target text.
For each query, we retrieve the top proponents from C4 using TrackStar. For the added tasks (i.e. the non-T-REx tasks), use the non-task-specine Hessian approximation R (Experiment 5 in Table because the number dimensions to downweight is both somewhat subjective and task-dependent ($.22.3). Our data and the results browser view identified influential pretraining examples are at wtttttp://// hub. cote
22


Reviews:
[{"review_id": "lUKi4jRIyk", "reviewer": "Reviewer_kPtE", "summary": "The authors propose TrackStar, a gradient-based influence method which works effectively on a 8B LM pretrained on 160B tokens, without the need for filtering or subsampling. The proposed method performs well at identifying instances which influence the model prediction, but model agnostic methods suc as BM25 still outperform it at finding passages which contain explicitly relevant facts. The paper delves into the problem of misalignment between factual attribution and causal influence, which they show becomes a smaller issue when scaling models.", "questions": "Tail-patch probability increase works worse with model scale (Fig 2). Do you believe this to be caused by a single instance mattering less for a larger model, or due to the inherent better performance of the underlying model? It would be interesting to distribute these scores for cases where the underlying model was right/wrong in its prediction.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Well written and motivated, easy to follow, very comprehensive exprimental evaluation with ablations and subsequent error analyses\n- Scales up TDA to large models, corpora and number of targets\n- Identifies and demonstrates a misalignment between attribution and influence\n- Empirically shows that influence becomes closer to attribution when scaling models", "weaknesses": "The only drawback I see is not delving deeper into the mismatch between infuelnce and attribution, however this is understandable as it is worthy of a separate paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mdwIOdsg1A", "reviewer": "Reviewer_wust", "summary": "This paper introduces a method for training data attribution that works effectively at scale. The proposed method relies on gradient-based estimations, where the authors combine several techniques such as optimizer state correction and normalized encodings to improve its attribution performance. Experiments on an 8B model pretrained on a corpus of over 160 billion tokens show the effectiveness of the proposed method and its different design choices. There are also some other interesting findings such as the misalignment between examples that can fully entail the target and examples that have actual high influence on model prediction.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Training data attribution is very important for understanding the connection between training data and the model's learned behaviors, and scalability is one key criterion especially nowadays when the model/data are large. The method this paper introduces is shown to be both more scalable and more effective than previous methods.\n\n- The ablation experiments are thorough and clearly show the effectiveness of the various design components of the proposed method.\n\n- The findings/analysis on the misalignment between factual attribution and causal influence are very interesting and could be a good inspiration for the community to understand LLM's behaviors around factual prediction.", "weaknesses": "- The different design components of the proposed method are mostly adaptations of previously proposed ones, whereas the proposed method is a combination of them. Hence, the technical novelty of this work is a bit limited.\n\n- Overall, the proposed designs are quite heuristical and not very solidly justified. For example, why *can* we enforce a block-diagonal structure for the auto-correlation matrix or perform unit normalization on the input vectors? What are the assumptions for doing these and are there some theoretical basis for them?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "985amGmrTv", "reviewer": "Reviewer_wLd8", "summary": "The paper studies extending influence functions to LLMs by modifying key components of TRAK (Park et al., 2023) such gradient computation, correction, and input processing (e.g., vector normalization). The authors applied their approach, dubbed TRAKSTAR, to a decoder-only model with 3 different sizes pretrained on the C4 dataset and compare to TRAK and other retrieval based approach. They found their approach to perform better than influence-based method, but to perform much worse compared to retrieval-based approaches. \n\nThis leads them to investigate the mist-alignment between attribution and influence (Section 7, the part that I find most interesting in the paper) and show that LLMs rely on priors for predictions---not just entailment. They perform an analysis on error cases.", "questions": "* You mention that you manually write natural language templates for all 97 relation types. I wonder how much will the results rely on these templates? Do you have experiments with different template variations? \n* Could you provide a discussion on real-world scenarios where scaling influence functions to LLM would be helpful/practical?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "* I like how the paper discusses difference/discrepancy between attribution and influence. I think this discussion could be useful to future work. \n* I found the error analysis in section 7 to be interesting.\n* Different design choices are well-ablated. \n* The proposed approach outperforms influence-based baselines at retrieving influential examples.", "weaknesses": "* The tail-patch metric used to evaluate influence will probably depend on the templates used. It is unclear how reliable such metric is.\n* If I understood correctly, the hessian approximation proposed requires access to an evaluation dataset. The authors should explain how this would work in the case of no have access to the full evaluation dataset. \n* It seems to me that lots of \"hacks\" are needed to get the approach to work (e.g., gradient correction, normalization, gradient compression via projection etc.) But I suppose this is a limitation of influence functions rather than of the proposed technique.\n* It is unclear to me that scaling influence functions to LLM is practical/reliable at all, as it seems that spurious correlations between the input fact and the pretraining examples to play a big role. For instance, looking at the errors provided in the Appendix, many of the retrieved proponents seem to completely unrelated to the fact-- some don't even have any word overlap such as example ID 862.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q0VrdzOy8b", "reviewer": "Reviewer_D1Kg", "summary": "The paper proposes TrackStar, a scalable, gradient-based influence tracing method for large language model (LLM) pretraining. It aims to identify influential training data examples for model predictions, enhancing model transparency. To achieve TrackStar, the authors compute the influence between two examples as the dot product between the projected and\ncorrected model gradients for those examples. TrackStar outperforms prior influence methods in identifying examples that entail a fact (attribution) and influence predictions (causal impact), though classical retrieval methods like BM25 still excel at attribution. They further analyze the effectiveness of TrackStar using extensive experiments.", "questions": "Please see the weaknesses. I have no more questions. I will improve my rate if you can address my concerns", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper addresses a critical gap in training data attribution (TDA) by scaling to pretraining setups, which may be further used for efficient pre-training or transparent training.\n- Experiments are well-designed, covering a range of models (up to 8B parameters) and extensive datasets, including T-REx and C4, to validate TrackStar’s effectiveness.\n- The authors offer a detailed analysis of the distinction between influence and attribution, adding clarity to how models use training data.", "weaknesses": "- ''For each fact, there are an average of 2.5 “ground-truth” sentences out of 19.6M sentences total''. Does this mean any training or testing examples contain 2.5 relations? The statement needs further clarity.\n- In the experiments, although the pass rate of records extracted by StarTrack is lower than that of BM25, training a model on StarTrack's examples leads to greater improvement. I wonder if limiting the selection to only the top 3 or top 5 examples—or even just the top example—would reverse this outcome.\n- For C4 evaluation, an NLI model is used. However, the accuracy cannot be fully guaranteed since C4 consists of longer documents. Human annotation of a subset is needed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)