Enhancing Training Data Attribution with Representational Optimization
Authors: Weiwei Sun
Haokun Liu
Nikhil Kandpal
Colin Raffel
Yiming Yang
Venue: NeurIPS 2025 spotlight ()
Keywords: Training Data Attribution
Primary area: deep_learning
TL;DR: AirRep is a text representation model optimized for TDA, offering performance comparable to gradient-based methods while being significantly more efficient.

Abstract: Training data attribution (TDA) methods aim to measure how training data impacts a model's predictions. 
While gradient-based attribution methods, such as influence functions, offer theoretical grounding, their computational costs make them impractical for large-scale applications. 
Representation-based approaches are far more scalable, but typically rely on heuristic embeddings that are not optimized for attribution, limiting their fidelity.
To address these challenges, we propose AirRep, 
a scalable, representation-based approach that closes this gap by learning task-specific and model-aligned representations optimized explicitly for TDA.
AirRep introduces two key innovations: a trainable encoder tuned for attribution quality, and an attention-based pooling mechanism that enables accurate estimation of group-wise influence.
We train AirRep using a ranking objective over automatically constructed training subsets labeled by their empirical effect on target predictions.
Experiments on instruction-tuned LLMs
demonstrate that AirRep achieves performance on par with state-of-the-art gradient-based approaches while being nearly two orders of magnitude more efficient at inference time.
Further analysis highlights its robustness 
and generalization across tasks and models.
Our code is available at https://github.com/sunnweiwei/AirRep.

Enhancing Training Data Attribution with Representational Optimization
Weiwei Sun
Haokun Liu²
Nikhil Kandpal
Colin Raffel Yiming Yang'
CCamegie Mellon University University of Toronto & Vector Institute
(sunnveiwei, hawiinahokkuueee nkandpa2. craffel]@gmail. com
Abstract
Training data attribution (TDA) methods aim to measure how training data im- pacts model's predictions. While gradient-based attribution methods, such as influence functions, offer theoretical grounding. their computational costs make them impractical for large-scale applications. Representation based approaches are far more scalable, but typically rely on heuristic embeddings that are not op- timized for attribution, limiting their fidelity. To address these challenges, we propose AirRep, a scalable, representation- -baeed approach that closes this gap by learning task-specific and model-aligned representations optimized explicitly for TDA. AirRep introduces two key innovations: trainable encoder tuned for attribution quality, and an attention-based pooling mechanism that enables accurate estimation of group-wise influence. We train AirRep using ranking objective over automatically constructed training subsets labeled by their empirical effect tar- get predictions. Experiments on instruction-tun tuned LLMs demonstrate that AirRep achieves performance on par with state-of-the-art gradient-based approaches while being nearly two orders of magnitude more efficient at inference time. Further analysis highlights its robustness and generalization across tasks and models. Our code is available at https: /ttt::////tttttt github.. com/ sunnveiveii www../unneeiieiiiiiiieoeee wwaeeeiiii/AApeppppper /irrRep.
1 Introduction
The remarkable success large language models (LLMs) has been demonstrated across a wide range of tasks. However, fundamental question remains open in machine learning: how does the behavior of LLMs depend on their training data? More specifically, what training examples cause models to generalize well-or underperform-ont specific inputs tasks? Training Data Attribution (TDA), the process of measuring the impact of specific parts of the training data on the predictions of machine learning models, an important step answering this question [1]. TDA crucial for ensuring transparency and accountability in AI systems by revealing how data influences model outputs.
Existing approaches to TDA broadly be categorized either gradient -ased representation- based. Gradient-based approaches, rooted in influence estimation [2, 3], aim to quantify the impact of individual training points model predictions. Their core idea is to
gradients and the inverse Hessian of the loss function to make first-order approximation of how a model's predictions would change if certain example theoretically well-motivated, they are typically computationally expensive and rely on assumptions of loss convexity and model optimality both which are violated in modern neural networks [4, 5].
removed from the training [3]. While these methods
In contrast, representation-bast based methods estimate influence based on similarity in the representation space. The core assumption is that training examples whose representations close to a test input are likely to have influenced its prediction [6]. Recent works have explored various notions of similarity
based model hidden states [7, 6], n-gram features [8], and text embeddings [9, 10]. Compared to gradient- based TDA, representation based approaches are more computationally efficient and
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Linear Datamodel Score (LDS)
Gradient
Representation
3000
Gradient
Representation
Speed (Examples/Soconc)
26.2
2500
2500
2500
2000
16.2
13.8
11.3
10-
7.7
1000
5-
3,9
500
278
31.2
5.3
30.2
29.5
0
0
Tracin
LESS
LoGra (9x)
LoGra (48x)
TF-IDF
RD5
GTE- GTE-Small
AirRep
Tracin
LESS
LoGra (3x)
LoGra (48x)
TF. IFF
RDS
GTE-Small
AirRep
Figure I: Performance comparison of gradient- and representation- based training data attribution (TDA) approaches. Left: Average linear data model score (LDS) [1] on unseen datasets (FLAN, Alpaca, Tulu, SafcRLHF). AirRep outperforms state-of-the-an pradient-based methods such as LoGra (which uses 48x more storage). Right: Inference speed (encoded examples per second on a single GPU), AirRep is nearly two orders of magnitude more efficient than gradient-based methods during inference.
scalable, making them well-suited for large-scale applications, such as curating pre-training data for LLMs [8, 11]. However, the quality of representation-based method depends heavily the selected feature space. This problem is compounded by the fact that existing work utilizes heuristically designed representations that are tailored to the particular target task model [12].
Finally, a persistent challenge both method families is how these methods are extended to estimate the collective influence of a group of training examples. Most past work extends single- example attribution methods to the group setting by simply summing over individual attribution scores within the group. However, this additive assumption fails to capture groupwise interactions and may lead to inaccurate influence estimation [13].
The advantages and shortcomings of existing methods raise natural question: Can we develop method that provides the best of both worlds?
We answer this question by introducing the Attentive Influence Ranking Representation (AirRep), a new representation-base approach that incorporates two novel mechanisms to improve performance over prior representation-basde methods. First uses a trainable encoder that produces task- and model-specifie text representations, ensuring that the feature space optimized for data attribution estimation. Second, it introduces an attention-based pooling mechanism that effectively aggregates multiple examples into single representation for more accurate group influence estimation. AirRep is optimized to produce scores that accurately reflect the influence of particular groups of training examples on a model's predictions. To achieve this, we device a automatic pipeline to construct pairwise comparison dataset where data subsets are ranked based on their contributions to model's predictions on specific target examples. We optimize AirRep on the dataset with a weighted pairwise ranking objective so that the learned representations and aggregation accurately represent the influence of data on the model's predictions.
Building on the setup of datamodels [1], we apply our approach on FLAN [14), an instruction-tuning dataset, and UltraCaat [15], a large-scale dialogue generation dataset, and evaluate the model on five unseen iatructtiontuuning test set (FLAN, Alapca, Tulu, SafcRLHF) that do not appear in the training data. Figure shows the main evaluation results. Notably, in the standard Linear Datamodeling Score (LDS) evaluation [1], AirRep significantly outperforms existing gradient- approach [16], despite being ~80x more computationally efficient and ~50x more storage efficient at inference time. Our evaluation additionally shows consistent patterns across various downstream TDA tasks, including data selection, data source identification, and task classification. AirRep also exhibits strong generalizabilty to new tasks and models, which further underscoring the effectiveness of AirRep (Section 5). Finally, we show that the inference efficiency of AirRep can effectively amortizes the cost of AirRep training [17] (Section 9).
2

---

2 Preliminaries
We consider framework that is applicable to different applications of TDA but focuses the training of language models (LMs) as a concrete example. Let S = (21, 22,.... kn) represent a training set of n data points, where 2, is training example that includes both an input and output. In the context of language models, we refer to the input as a "prompt" and the output as a response. Furthermore, let denote test example. The goal of TDA is to understand how the training examples in S contribute to the model's 's prediction on E.
Formally, let 0 represent the parameters of the LM. The model is fine-tuned on S, aiming to find
0" arg min
Σ E(zz:)))
2,E8
where ((-) denotes the cross-entrope token prediction loss. As language model training is non-convex, we use 0" to denote the optimized parameters after fine-tuning, regardless of whether a true minimum has been found. We define the model's "prediction oulcome" on as the cross-entropy loss computed on x, given by r(x, S) = E(a;; 0*), ) i.e. the evaluation loss on the test example æ for the model fine-tuned on the dataset S. Here, we use the cross-entropy loss instead of task-specific metrics like the accuracy, FI score, or ROUGE [18] because the loss is a generic metric that is defined for all NLP tasks and has been found to be correlated with task-specific metrics [19].
data attribution model, denoted by fxx 5), aims to estimate the actual retraining outcome r(z, -e. the loss of the model on example after training on dataset S. Existing methods can be generally categorized into two groups: gradient- based methods and representation-based based methods.
Gradient-based methods Gradient-basse methods largely stem from influence functions [2, 3], which approximate how 0 would change in response to infinitesimal perturbations in the weighting of training instances. Specifically, influence functions quantify the influence from a single training example 2; to evaluation example with closed form expression:
fir(s, z;) = - oob(ax; 0)
()((=;
(2)
where H is the inverse Hessian of the training loss wrt model parameters. The derivation of Eq. 2 given by first-order Taylor approximation of E(1; 8) around 6* [20] (see Appendix A for details).
The size of model neural networks makes the computation and inversion of the Hessian intractable. Consequently, practical influence function-based methods often use an approximation of the Hes- sian [12, 21J. Recent techniques have also considered techniques like gradient projection [16] and model ensembling [22. 23]. As representative approach in LLM application, LoGra [16] define the influence of a training example on is computed as:
faD 2,) = <(x) T (2;),,
(3)
where (02) is the projected, Hessian-corrected, and unit-normalized gradient for the input example x, given model parameters 8:
(b=) = norm H 1 [=;
where norm [
| a unit normalization operation [24]; 0 denotes the trainable weights of a
LoGra module [16] with PCA initialization for gradient projection; Vg((2; f) is the gradient of the loss with respect to 0; and Hg à is computed on å and is approximated based on the Kronecker- Factored Approximate Curvature (KFAC) algorithm [21]. Further, under first-order approximation, the group influence fow(2,, S) is estimated as the summation of individual influences.
Representation- based methods Representation- based methods encode examples as embedding vectors, eliminating the need for gradient computation. The influence score of a training example 22 to r can then be computed as:
(x, 24) Fnc(x)] T Fac((z))
where Enc(2) denotes a continuous embedding of input 2 encoded by encoder model Enc. Group
influence is usually defined as summation of individual influence under linear assumption [1].
In term of the encoding function Enc, existing work in the text domain has investigated different heuristic designs including (last-layer) hidden states of a language model [7, 6). n-gram representa- tions [8], TF-IDF representations [25, 26), and text embeddings [11, 27J. See Appendix for further comparison of the two types of TDA methods.
3

---

x
Encoder
Pairwise Comparison
$
Encoder
Pooler
f(x,,,,)
Attention
$2
Encoder
Pooler
f(x,S2) 0.2
Training Subsets
Individual Embeddings
Group Embeddings
Model Architecture and Optimization. The test example x and the training subsets S1 and S2 are encoded by an encoder with pooler obtain embeddings. The score is computed as the inner product of the embeddings. The overall model is trained based on pairwise comparisons to distinguish the usefulness of different subsets with respect the example œ.
Figure 2:
3 AirRep
To bridge the gap between gradient- and representation- data attribution, we introduce AirRep. novel data-driven approach to enhance representational influence models. At a high level, AirRep consists of trainable encoder Enc for embedding input examples and pooling layer Agg that aggregates the representation of individual examples produce a group representation. The influence score between training set S and a target example is then computed as
(6) Specifically, AirRep incorporates two novel mechanisms improve performance over standard representation-biae based methods: First, we introduce an effective attention-base pooling mechanism that better reflects group influence effect (Section 3.1). Second, we optimize the model on auto- generated data so that AirRep's 's score better reflects the underlying models predictions (Section 3.2).
(z. S) Enc(r) T. Agg(Enc(z.)
S)
3.1 Attention-based Pooling
Recall that data attribution methods typically include mechanism for aggregating the contributions of elements of 2, E S when performing group influence. For this purpose, AirRep uses a pooling layer to aggregate the representations Enc(2; into single embedding. The design of this pooling layer is crucial for modeling the relationships between group of data samples.
In AirRep, propose simple yet effective pooling techniques: attention-baset influence pool- ing. Specifically, AirRep's 's aggregation method incorporates an attention- operation to capture interactions between elements of and S:
fairRe (2,, S) Fnn(x)
Enc(2;),
exp(1 Enc(2) Eno(2;))
where Ok Σyen, exp() Ene(z) T Enc(z.)))
where Enc(x) denotes the sentence embedding predicted by BERT-based sentence encoder [28, 29], Og denotes the attention score of example 2;, and denotes the absolute value operation.
Remark Empirically, observe that influence scores are sparse, where each test example relies on only a few training points, while others add noise. This supports the need for selective pooling, consistent with prior findings [1, 12, 30). Additionally, the proposed attention-based influence pooling can also related to high-order group influence functions. Specifically, Basu et al. [13] showed that second-order terms capture additional relationship between samples. Extending this, show that high-order group influence introduces sample-wise weights, akin to our attention-base based pooling (see Appendix B.1 for more discussion).
3.2 Optimization
Both the Enc and Agg components of AirRep are trainable modules that together take as input and S and output scalar score f(x, S). Ultimately, the objective of data attribution is that f S)

---

is correlated to r(z, S), i.c., the actual prediction on I of model trained on S. Since AirRep is overall trainable, we therefore formulate training objective that aims to make fx, reflect r(x, S). Given ground-truth collection of r(x,, S) scores for different datasets S, we could in principle train f(r,, 5) to match these scores as a regression problem. However, in data attribution we care less about matching the exact values of r(x, 5) than we do reflecting the ranking of different example groups (datasets), since ultimately we aim to answer questions like "which group had the largest influence this example?", We therefore formulate training as a pairwise ranking problem that is, given two different data subsets S and S2, whether the score f(x,, Si) and f(x, S2) reflect the true data preference r(x, S1) and r(z, S2). In this section, will first introduce the data generation pipeline we used for collecting ground-trum hounnd-rrutha r(z, S) scores and then describe the specific training objective we formulated.
Data Generation To construct data for training AirRep, we assume access to a large corpus of example datapoints. Then, we generate woss-walidationn-styl data and compute ground-truth attribution scores. Specifically, we first sample N, examples as a validation split and N₂ examples as a training split. Then, from the training split, randomly sample M subsets of data with replacement, denoted as = {S1, S2...., SM}, ). where each subset S; contains n training examples. For each training subset S₁, we finetune a language model on it as in Eq. 1 to obtain model checkpoint O,- We then evaluate the trained model on each example a in the validation subset to calculate the loss: [(x; 0.). After training M models and calculating the loss, we compute the negative normalized loss for each r in the validation set as
E(x; - Mean ((e(2; 0,) E [M]]) Var ((e(x; 03) 1 E [M])
f(x, Si)
(8)
where the original loss {<;; 0.) is normalized by the mean and variance across models to stabilize the training signal, ensuring that the distribution of scores for each target o is on similar scale.
Training Objective Given target example from our validation set, M training subsets S = (,..,, ,,, }. and the corresponding normalized loss values (Fxx,,)); E [M]] AirRep estimates the normalized loss values for each subset fæz For brevity, we use r; and to represent these scores, respectively. Our proposed training objective is pairwise ranking loss that encourages the difference between the model-predicted scores, ʃi ʃj. to better align with the difference in precomputed scores, r T---
In practice, the label ri is noisy due the stochastic nature of LM training. To mitigate the impact of noisy labels, we adopt a weighted pairwise ranking loss, inspired by importance reweighting [31, 32], which assigns lower weights to uncertain labels while prioritizing more reliable ones:
L(x,8) =- Σ Trary Way log olf fz),
AJEM
(9)
0, (min(1r - TyL. Tmax), if Tmin S Ir: ryl-
if Ir - ryy <. Tmin;
where
is sigmoid function, way is weighting function that depends on the absolute difference in ground truth scores, Ir, - ryl. and is clipped using thresholds Tmin and Tmax to avoid training on i, pairs with incorrect ordering and mitigate the impact of outliers.
4 Experimental Setup
Model Our experiments focus on LM finetuning, using the Qwen2.5 model family [36] as our base LMs. During training, we start with the base LM and fine-tune it using batch size of 32 and the AdamW optimizer [37] with learning rate of 2e-5 for two epochs.
Data Generation For AirRep training, we utilize two datasets, each representing distinct scenar- ios: instruction tuning on standard NLP tasks and training to improve conversational abilities: (i) FLAN [14), an instruction-tuning dataset for language models, and (ii) UltraChat [15], a large-scale dataset of instructional conversations covering wide range of topics. To generate training signal, we
set N, = 10a and N, = 10a. The training subsets number is M 100, with each subsets containing
1.000 samples. We construct 100 cross- -rlidaialii instances. Thus, in total, the data includes
5

---

Table 1: LDS Evaluation Results of Qwen2.5-0.5B four Datasets. Aug denotes the average score.
Dim refers to the dimensionality of the embeddings (which corresponds to storage size).
Method
Dim
Avg
FLAN
Alpaca
Tulu
SaF
TracIn [22]
18432 (48x)
11.33
14.75
9.21
10.75
10.60
LESS [25]
8196 (21x)
16.16
16.40
9.59
13.02
25.63
LoGra [16]
1152 (3x)
13.78
13.32
6.87
10.16
24.76
LoGra [16]
18432 (48x)
18.45
19.75
12.38
14.88
26.82
Dsdm [33]
18432 (48x)
18.02
19.67
12.15
14.31
25.94
TF-IDF [34]
9.98
2.52
7.24
5.24
24.94
DSIR [8]
-0.02
0.49
2.01
-0.49
-2.10
RDS [6]
896 (2.3x)
3.86
0.74
0.87
1.89
11.94
GTE-Small [35]
384 (1x)
7.65
0.92
1.74
1.14
26.80
AirRep (Ours)
384 (1x)
26.23
21.11
22.58
15.14
46.08
10 unique training subsets and 10? training examples. The Qwen2.5-00 5-0.BB LM is then fine-tuned on these training subsets and evaluated on the corresponding validation subset to obtain the label r(..).
AirRep Training Details We initialize AirRep using the GTE-Small 30M parameter embedding model [35], and apply randomly initialized projection matrix on top. AirRep is trained separately on FLAN and UltraChat. To construct the data for each training step of AirRep, we randomly select one cross- valsdvaaiiotion instance, then sample 1, 1,000 examples from its validation subset and 32 training subsets. Distributed training is employed to maximize GPU memory utilization. The clipping thresholds, Tiin and Taxx are set to 0.1 and 5.0, respectively. The model is optimized for up to 2,000 steps using the Adam W optimizer with a learning rate of 1 x 10-4.
Evaluation Datasets We use the following datasets for evaluation: (i) FLAN [14]: The evaluation data of FLAN contains 66 NLP tasks spanning diverse categories. (ii) Alpaca [38]: An instruction- tuning dataset generated by OpenAl's text-davinci-000 model. (iii) Tulu [39]: An instruction-tuning dataset comprising diverse data sources. (iv) SafeRLHF [40]: A dataset for safety alignment of large language models, where each response is labeled as either safe or unsafe. Note that each evaluation data contain a test set and a training set. We ensure that all evaluation data remain excluded from AirRep's s optimization data to guarantee that the evaluation results reflect AirRep's generalization to unseen data. Appendix D provides additional dataset details.
Baselines We compare AirRep with representative gradient- based and representation- based meth- ods. The trudient-based baselines we consider are as follows: (i) LoGra [16], an optimized version of the influence function as described in Eq.4, based on its implementation library. (ii) TracIn [22]. which leverages the dot product of gradient vectors. (iii) LESS [25], which computes projected gradients on LoRA weights and adjusts gradients using the Adam W states. (iv) Dsdm [33], an efficient implementation of Trak [12] for LLM tasks.
the Logix software
The representation-baan
baselines we consider are as follows: (i) TF-IDF [34], (ii) DSIR [8], hashed n-gram features, (iii) RDS [6], which utilizes the last-layer hidden states of LMs, and (iv) GTE [35], text embedding model trained on text similarity and retrieval tasks which was also used as the base pre-trained model for AirRep. See Appendix E for more implementation details.
Evaluation of Data Attribution
Setup To evaluate these data attribution methods, we follow [1] and use the linear datamodeling score (LDS) as our evaluation metric. For each evaluation dataset, we sample 100 different random subsets (S1. S100) of its training set, each containing n training examples, and train the target LM on each of these subsets. For each example in the test set, we approximate the expectation of the model prediction r(x. S.) as the evaluation cross-entropy loss. Given the score f(x, predicted by the data attribution model, the LDS is computed as the Spearman rank correlation between the true and estimated influence scores: p( (r(a, S:) : E [100]], (f S,) E [100]]). Finally, the LDS on the dataset is averaged across all test examples.

---

LoGra
20
(91x)
16
(166x)
12
(262x)
(262x)
Dim
LDS Comrelatin
15
12
(200))
10
4
304(1н)
Qwee2.5-0.58
Qwen2.5-1158
Qwwn2.5-38
Qwen2.5-76
Computation time to GTE- GTE-Smail (in Log10 Scale)
Figure 3: LDS correlation scores on FLAN vs. inference- time cost and storage for various TDA methods across Qwen2.5 models (0.5B-7B). Computation time (Logl0 scale) measured relative to GTE-Small on the same machine; marker size reflects storage (smaller more efficient). Each method has multiple points for different model/dimension settings.
Results Table I presents the LDS evaluation results for attributions computed on the Qwen2.5-0.55 model's predictions across four evaluation datasets. AirRep achieves the best performance across all evaluation datasets. Notably, it outperforms the competitive gradient- despite being 80x more computationally efficient and 50x more storage efficient. Additionally, previous representation-basee methods exhibit low correlation scores in this evaluation, underscoring the importance of task-aware optimization of representations.
approach LoGra [16),
Figure 3 illustrates the LDS correlation when applying TDA methods to different target language models, ranging from Qwen2.5-0.5B to Qwen2.5-7B. For gradient-based methods such as LoGra and LESS, the proxy model can be adjusted -i.., smaller LM can be used as a proxy to compute influence scores for larger target LMs [25]. Additionally, the projection dimensionality can be modified to balance computation, storage, and performance.
cases, we use the same AirRep model trained exclusively on data generated by Qwen2.5-0055. Notably, AirRep demonstrates strong performance when applied to larger target LMs. Specifically, it consistently achieves results comparable to the best gradient-basee methods, which require significantly more resources (e.g.. 84x to 262x x more computation time and 50x to 80x more storage). These results highlight AirRep's potential for efficient and effective data attribution. AirRep also demonstrates strong generalization across different data sizes and target LM types (Figure 8 and Figure 9).
Evaluation of Data Selection
Setup One use of data attribution is data se- lection, where we pick out the highest value training data and train only on these samples. To evaluate each attribution method's utility in this setting, we use the FLAN collection of datasets. For each task in FLAN, use each
DEIR
GTE
LoGrn
58
54
TDA method to select a high-value subset us- ing greedy selection strategy. First, given the scores predicted by TDA models, we rank all training examples based their highest score with respect to the test set samples and retain the top 1,000 examples. We then train LMs the selected subset and evaluate their performance on the test set. The evaluation metrie is the F1 score between the L.-generated outputs and the ground-truth answers.
46
Owen-o.31 3B
Owven-t 5B
Qeenn-B
Figure 4: Evaluation Results of Data Selection. We report the average Fl score of 66 tasks in FLAN, obtained by training Qwen2.5 LM of different sizes on the top- 1000 selected examples for each task.
Results Figure& presents the results of using six TDA methods for data selection, along with random data selection baseline, evaluated in terms of Fl score and selection cost. Key findings
7

---

include: (i) Across all model scales, all methods (except DSIR) outperform random selection. (ii) Gradient-based approaches (LESS and LoGra) achieve better performance than representation-baseee methods (TF-IDF and GTE); however, they are computationally expensive- -tee time required to select subset exceeds the time needed to train the model on the full data, (iii) AirRep performs comparably to gradient-based methods and full-data training, while significantly outperforming other representation-base methods in effectiveness, all while maintaining efficiency.
7
Evaluation of Data Classification
Setup TDA methods have been applied to measuring data similarity for the purpose of model explainability [6]. To evaluate the effectiveness of such explanations, Hanawa et al. [6] introduced the "Identical Class Test" for image classification, which assesses whether the most influential training example identified by TDA methods belongs to the same class as the test instance.
We extend this test to LM fine-tuning and formulate the "data classification" task, which measures the accuracy of TDA methods by evaluating whether the top-retrieved training examples share the same label as the test example. We consider the following sources of labels for classification: (i) FLAN and Tulu: FLAN comprises 66 NLP tasks, while Tulu is collection of 12 source datasets (e.g., ShareGPT, COT, Code-Alpaca). We use the task type or source dataset as label, expecting that training data from the same task source as the test example will be more valuable. (ii) SafeRLHF: SafeRLHF contains both safe and unsafe data annotated with safety labels. We use the safety label, expecting that harmful training data are more likely to lead harmful generations at test time.
Results Table 2 presents the evaluation results of LoGra, GTE-Small, and AirRep. The results indicate that AirRep significantly improves data classification accuracy on FLAN compared to UTEE-mmlll (inereasing from 50.59 to 86.41) and also outperforms LoGra. Notably, since AirRep is trained without using any data labels, this suggests that the model can learn to represent task-related information an unsupervised manner. similar trend is observed for Tulu, where AirRep demon- strates superior performance in identifying the data source. For SafeRLHF, AirRep outperforms LoGra but lags behind GTE-Small, likely because AirRep's 's training data, generated from UltraChat, does not contain harmful content and thus lacks the necessary learning signal.
Table 2: Accuracy of Data Classification.
Method
FLAN
Tulu
SafeRLHF
LoGra (Dim 1152)
71.61
76.60
78.00
LoGra (Dim 18432)
85.44
86.00
83.20
GTE-Small AirRep
50.59
76.60
90.60
86.41
88.20
87.20
Ablation Study
We conduct ablation study evaluate each component's contribution. Figure 5 presents the results as the average LDS score across the four datasets we consider. Starting with GTE (Small), which gets score of 7.65, we first optimize the encoder while keeping mean pooling instead of adding the attention pooling layer. This results in AirRep w/o Attention, which achieves a score of 19.82-an an improvement of 12.17 demon- strating the effectiveness of encoder optimiza- tion. Next, when we introduce the attention pooling layer and jointly optimize the model, we observe a further improvement of 6.41. Additionally, we examine whether simply adding a softmax attention pooling to baseline to re-weight individual influences can enhance performance. From the results of GTE + Attention and LoGra + Attention, we observe only marginal improvements. This suggests that while weighting data is beneficial, optimizing the weighting distribution is also necessary to achieve significant gains. Finally, we observe that adding attention pooling not only improves group influence estimation but also leads to better underlying representations.
GTE
7.65
+11.117
GTE +Attention
8.16
AirRep Attention
+6.41
AirRep
28.23
LoGra
18.45
LoGra +Attention
18.89
14
27
Figure 5: Ablation Study results: average LDS score FLAN, Alpaca, Tulu, and SafeRLHE.
'erre focus subset of most competitive directly relevant baselines, see Appendix H.2 for more.
8

---

Data Altribution (LDS)
Data Classificatio (Accuracy)
Processed Data Size
Rig
Arfie
100M
40
(5439
10M
30
LoGr
Artep (Aig Soom
Lodro (1162
GfE-3mal
015
20
Locre
70
100K
65
10
60
10"
10
10"
10
10"
10"
10
10"
8
18
24
AirRap training size (log scale)
AirRep training size (log scale)
GPU hours
Figure 6: Left: AirRep training data size vs. LDS. Middle: AirRep training data size vs. Data Classification Accuracy. Right: GPU hours number of data processed. AirRep incurs higher training costs due to data generation (yellow region) and model training (orange region), but is more inference-effiiiett.
Amortizing Training Cost
Compared with gradient-based methods, AirRep more efficient at inference but requires training auto-gettt ceito-genrraedd data, which incurs additional cost. Our evaluation above shows that FLAN- or UltraChat-pretraince AirRep generalizes well on new TDA data and tasks. We now consider another setting where AirRep is re-trained from scratch for each TDA task, and analyze how its inference efficiency can amortize its training cost [17].
Figure 6 (left and middle) shows AirRep's average attribution and classification scores with different training data sizes (measured by number of datapoints used in training). Light curves indicate per-test-set scores. As the training size increases, AirRep achieves better performance. In particular, AirRep outperforms LoGra (1122)which has comparable embedding size-when trained with around 10 examples. We refer to this AirRep model as the "crossover checkpoin!".
Then, Figure 6 (right) compares the number of training examples that can be processed under different A100 GPU hour budgets by AirRep (crossover checkpoint) and LoGra, accounting for both data generation and model training costs. LoGra completes model training earlier, but AirRep despite longer training time time-achieves nearly two orders of magnitude more throughput once model is achieves
trained. Notably, turning point appears around 475K examples, where both methods take 5.6 GPU hours. Beyond this, even with retraining, AirRep is faster than LoGra. For example, with 24 GPU hours, AirRep can attribute over 100M examples, whereas LoGra can process only around 2M. This highlights the advantage of AirRep for large-scale TDA scenarios. We additionally emphasize that, as shown in Section AirRep can be transferred from small model to larger models and still attain stronger performance than all baselines, so the "re-training from scratch" results in this section represent worst-case results for AirRep.
182
Amortizing Across Models We evaluate Air- Rep, trained with Qwen-O.5B- -.5B-genertted super- vision, using different target LMs (Qwen-0.5B. Qwen-7B, Alama-1B, Qwen3-0.6th Qwen3- 0.6B (Thinking), TinyLlama--IB, GPT-2)
ThyLlamm-18
Onerá-0 es
233
Iwpet LM
(Thenking))
11.2
11.
FLAN. Results in Figure show that across all targets, AirRep consistently surpasses Lo- Gra, even when LoGra computes gradients on the target LMs. We also observe that LoGra is sensitive to the proxy used for gradient calcula- tion when the target model differs significantly (e.g., GPT-2), its performance drops notably, whereas AirRep maintains strong performance. This indicates the robust generalization of Air- Rep and that training cost can be amortized across different target language models.
Lamo 18
19.0
Gwan-78
Owen se
Locra
LeGra
Airio
Figure 7: Evaluate with different target LMs. Ar- Rep shows strong generalization.

---

10 Related Work
Most existing work in TDA focuses on gradient-basee methods, such as Influence Functions [3] and their subsequent optimizations, including group influence estimation [41, 13), Hessian approx- imation [42, 21, 43], gradient normalization [24, 25], gradient projection [22, 16], model ensem- bling [12, 33, 44], and distillation of the proxy influence models [45-47]. However, studies have found that the derivation of gradient- based methods may rely on incorrect assumptions large neural
networks [48, 4]. Additionally, madient-bbsed methods suffer from the computational expense of gradient calculations and Hessian inverses [21, 33].
Alternatively, representation-bsse approaches for TDA measure data influence in the representation space [7, 49, 6, 9]. offering computational and storage efficiency. However, current methods rely on heuristically designed representations without task-aware and model-specific optimization, often leading to suboptimal results [12]. Some research on TDA explores simulation-base methods such datamodels 150-52, 11, which are usually not tractable for LLM applications due to their computational expense.
TDA methods have been applied to various purposes in machine learning, such as identifying mislabeled training data points [3, 53], detecting data poisoning attacks [54, 55], and guiding data augmentation [56, 57]. TDA methods have also been used for understanding L.M generalization [21] and learning processes [30, 58]. Another notable line of related work is data selection for L.LMs, which focuses on curating data to optimize test performance [59]. TDA-like functions for measuring relationships between data have been widely adopted in data selection pipelines to enhance the quality and diversity of the selected data. Both gradient-based [60, 25, 47] and representation- based [8, 61, 62, 10] TDA approaches have been applied. Meanwhile, studies have explored using text embedding models to select in-context examples for LLM prompting [63 [63-65).
11 Conclusion
This paper introduces new approach training data attribution called AirRep, which combines the advantages of task-driven optimization in gradient-based approaches with the efficiency of representation--aase methods. AirRep represents training data embeddings and scores their influence using inner-product similarity in the representation space. AirRep's novel attention- based pooling mechanism captures the group effects, and the entire model optimized using a novel proposed weighted pairwise ranking objective over automatically generated influence signals. Evaluation on LLM fine-tuning demonstrates the effectiveness and generalization of AirRep.
Limitations of this work include the added cost of model training and evaluation limited to the LLM fine-tuning stage. While our current evaluation focuses on language tasks, the design of AirRep is modality-agnostic-in relies only on the availability of suitable encoders and the ability to estimate loss differences. In principle, this framework could be applied other domains such vision or multimodal settings. For future work, we plan to extend our method to additional scenarios, such as the pre-training stages of LLMs and multi-modal model training.
References
[1] Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclerc, and Aleksander Madry. Datamodels: Predicting predictions from training data. ArXiv, abs/2202.00622. 2022.
[2] Frank Hampel. The influence curve and its role in robust estimation. Journal of the American Statistical Association, 69:383-393, 1974.
[3] Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions.
In International Conference on Machine Learning, 2017.
[4] Samyadeep Basu, Phillip Pope, and Soheil Feizi. Influence functions in deep learning are fragile. ArXiv, abs/2006. 14651, 2020.
[5] Andrea Schioppa, Katja Filippova, Ivan Titov, and Polina Zablotskaia. Theoretical and practical perspectives what influence functions do. ArXiv, abs/23055166971 2023.
10

---

[6] Kazuaki Hanawa, Sho Yokoi. Satoshi Hara, and Kentaro Inui. Evaluation of similarity- explanations. In International Conference on Learning Representations, 2021.
[7] Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, and Oliver Wang. The unrea- sonable effectiveness of deep features as a perceptual metric. 2018 IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 586-595, 2018.
[8] Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. ArXiv, abs/2302.0316$2 2023.
[9] Ekin Akyürek, Tolga Bolukbasi, Frederick Liu, Binbin Xiong, Ian Tenney, Jacob Andreas, and Kelvin Guu. Towards tracing factual knowledge in language models back to the training data. 2022.
[10] Devleena Das and Vivek Khetan, Deft-ucs: Data efficient fine-tuning for pre-trained language models via unsupervised core-set selection for text-editing. In Conference on Empirical Methods in Natural Language Processing, 2023.
[11] Guilherme Penedo, Hynek Kydl'icek, Loubna Ben Allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. ArXiv, abs/2406. 17557, 2024.
[12] Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclere, and Aleksander Madry. Trak: Attributing model behavior at scale. In International Conference on Machine Learning, 2023.
[13] Samyadeep Basu, Xuchen You, and Soheil Feizi. On second-order group influence functions for black-box predictions. In International Conference Machine Learning, 2019.
[14] Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoe V. Le. Finetuned language models are zero-shot learners. ArXiv, abs/2109.01655 2021.
[15] Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. Enhancing chat language models by scaling high-quality instructional conversations. ArXiv, abs/2305. 14233, 2023.
[16] 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 P. Xing. What is your data worth to gpt? Ilm-scale data valuation with influence functions. ArXiv, abs/2405511954. 2024.
[17] Ian Covert, Chanwoo Kim, Su-In Lee, James Zou, and Tatsunori Hashimoto. Stochastic amorti- zation: unified approach accelerate feature and data attribution. ArXiv, abs/2401.15866. 2024.
[18] Chin-Yee Yew Lin. Rouge: A package for automatic evaluation of summaries. Annual Meeting of the Association for Computational Linguistics, 2004.
[19] Mengzhou Xia, Mikel Artetxe, Chunting Zhou, Xi Victoria Lin, Ramakanth Pasunuru, Danqi Chen, Luke Zettlemoyer, and Ves Stoyanov. Training trajectories of language models across scales. ArXiv, abs/2212.098033 2022.
[20] Steven G. Krantz and Harold R. Parks. The implicit function theorem. 2002.
[21] Roger Baker Grosse, Juhan Bae, Cem Anil. Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamil.e Lukovsiut.e, Karina Nguyen, Nicholas Joseph, Sam McCandlish, Jared Kaplan, and Sam Bowman. Studying large language model generalization with influence functions. ArXiv, fis/2308.03290 2023.
[22] Garima Pruthi, Frederick Liu, Mukund Sundararajan, and Satyen Kale. Estimating training data influence by tracking gradient descent. ArXiv, abs/2002.08484, 2020.

---

[23] Sung
Yub Yub Kim, Kyungsu Kim, and Eunho Yang. Gex: A flexible method for approximating influence via geometric ensemble. A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 9888-5911 Curran Associates, Inc., 2023.
[24] Elnaz Barshan, Marc-Etienne Brunet, and Gintare Karolina Dziugaite. Relatif: Identifying explanatory training samples via relative influence. In International Conference on Artificial Intelligence and Statistics, 2020.
[25] Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. Less: Selecting influential data for targeted instruction tuning. ArXiv, abs/2402.00333. 2024.
[26] Stephen E. Robertson and Hugo Zaragoza. The probabilistic relevance framework: Bm25 and beyond. Found. Trends Inf. Retr., 3:333-389, 2009.
[27] Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Yitzhak Gadre, Hritik Bansal, Etash Kumar Guha, Sedrick Scott Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel, Jean- Pierre Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bitton, Marianna Nezhurina, Amro Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham M. Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri. Sewoong Oh, Luke S. Zettlemoyer, Kyle Lo, Alaacldin El-Nouby, Hadi Pouransari, Alexander Toshev, Stephanie Wang. Dirk Groeneveld, Luca Soldani, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-Im: In search of the next generation of training sets for language models. ArXiv, abs/2406. 1994, 2024.
[28] Jacob Devlin, Ming- Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In North American Chapter of the Association for Computational Linguistics, 2019.
[29] Nils Reimers and Iryna Gurevych. Sentence-bert, Sentence embeddings using siamese bert- networks. In Conference on Empirical Methods in Natural Language Processing, 2019.
[30] Tyler A. Chang, Dheeraj Rajagopal, Tolga Bolukbasi, Lucas Dixon, and Ian Tenney. Scalable influence and fact tracing for large language model pretraining. ArXiv, abs/2410.174133, 2024.
[31] Tongliang Liu and Dacheng Tao. Classification with noisy labels by importance reweighting. IEEE Transactions on pattern analysis and machine intelligence, 38(3):447-461, 2015.
[32] Hwanjun Song, Minseok Kim, Dongmin Park, Yooju Shin, and Jae- Gil Lee. Learning from noisy labels with deep neural networks: A survey. IEEE Transactions on Neural Networks and Learning Systems, 34:8135-811 8153, 2020.
[33] Logan Engstrom, Axel Feldmann, and Aleksander Madry. Dsdm: Model-aware dataset selection with datamodels. ArXiv, abs/2401.12926, 2024.
[34] Karen Spärck Jones. A statistical interpretation of term specificity and its application in retrieval.
J. Documentation, 60:493-502, 2021.
[35] Zehan Li, Xin Zhang, Yanzhao Zhang, Dingkun Long, Pengjun Xie, and Meishan Zhang.
Towards general text embeddings with multi-stage contrastive learning. ArXiv, 0at/2300.0328// 2023.
[36] An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang. Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Junyang Lin, Kai Dang, Keming Lu, YanY Chen, Kexin Yang. Mei Li, Min Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao. Runji Lin, Shijie Wang, Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Yang Fan, Yang Yao, Yichang Zhang, Yunyang Wan, Yunfei Chu, Zeyu Cui, Zhenru Zhang, and Zhi- Wei Fan. Qwen2 technical report. ArXiv, abs/2407.106711, 2024.
12

---

[37] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2017.
[38] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto. Stanford alpaca: An instruction-following llama model, 2023.
[39] Yizhong Wang, Hamish Ivison, Pradeep Dasigi, Jack Hessel, Tushar Khot, Khyathi Raghavi Chandu, David Wadden, Kelsey MacMillan, Noah A. Smith, Iz Beltagy, and Hannaneh Ha- jishirzi. How far can camels go? exploring the state of instruction tuning on open resources. ArXiv, fbss2300.0075.. 2023.
[40] Josef Dai, Xuchai Pan, Ruiyang Sun, Jiaming Ji, Xinbo Xu, Mickel Liu, Yizhou Wang, and Yaodong Yang. Safe rlhf: Safe reinforcement learning from human feedback. ArXiv, 0as/2310.12773.1 2023.
[41] Pang Wei Koh, Kai-Siang Ang, Hubert Hua Kian Teo, and Perey Liang. On the accuracy of influence functions for measuring group effects. In Neural Information Processing Systems, 2019.
[42] Andrea Schioppa, Polina Zablotskaia, David Vilar, and Artem Sokolov. Scaling up influence functions. In AAAI Conference on Artificial Intelligence, 2021.
[43] Yongchan Kwon, Eric Wu, Kevin Wu, and James Zou. Datainf: Efficiently estimating data influence in lora-tuned Iims and diffusion models. ArXiv, abs/2310.00902 2023.
[44] Jiachen T. Wang, Dawn Xiaodong Song, James Zou, Prateek Mittal, and Ruoxi Jia. Capturing the temporal dependence of training data influence. ArXiv, abs/2412.095388, 2024.
[45] Faisal Ladhak, Esin Durmus, and Tatsunori Hashimoto. Contrastive error attribution for finetuned language models. In Annual Meeting of the Association for Computational Linguistics, 2022.
[46] Qingyi Liu, Yekun Chai, Shuohuan Wang, Yu Sun, Keze Wang, and Hua Wu. On training data influence of gpt models. In Conference Empirical Methods in Natural Language Processing, 2024.
[47] Zichun Yu, Spandan Das, and Chenyan Xiong. Mates: Model-aware data selection for efficient pretraining with data influence models. ArXiv, abs/2406.000466, 2024.
[48] Juhan Bae, Nathan Ng, Alston Lo, Marzyeh Ghassemi, and Roger Baker Grosse. If influence functions are the answer, then what is the question? ArXiv, has//200.053644, 2022.
[49] Nazneen Rajani, Ben Krause, Wen Yin, Tong Niu, Richard Socher, and Caiming Xiong. Explaining and improving model behavior with k nearest neighbor representations. ArXiv, abs/2010.09030, 2020.
[50] Vitaly Feldman and Chiyuan Zhang. What neural networks memorize and why: Discovering the long tail via influence estimation. ArXiv, hass/200..07703 2020.
[51] Amirata Ghorbani and James Y. Zou. Data shapley: Equitable valuation of data for machine learning. ArXiv, abs/1904 (b0//9994.0288.. 2019.
[52] R. Jia, David Dao, Boxin Wang. Frances Ann Hubis, Nicholas Hynes, Nezihe Merve G" Guurel, Bo Li, Ce Zhang, Dawn Xiaodong Song, and Costas J. Spanos. Towards efficient data valuation based on the shapley value. ArXiv, abs/1902. 10275, 2019.
[53] Shuming Kong, Yanyan Shen, and Limpeng Huang. Resolving training biases via influence- based data relabeling. In International Conference on Learning Representations, 2022.
[54] Minghong Fang, Neil Zhenqiang Gong, and Jia Liu. Influence function based data poisoning attacks to top-n recommender systems. Proceedings of The Web Conference 2020, 2020.
13

---

[55] Matthew Jagielski, Giorgio Severi, Niklas Pousette Harger, and Alina Oprea. Subpopulation data poisoning attacks. Proceedings of the 2021 ACM SIGSAC Conference on Computer and Communications Security, 2020.
[56] Donghoon Lee, Hyunsin Park, Trung Xuan Pham, and Chang Dong Yoo. Learning augmentation network via influence functions. 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages
2020.
[57] Sejoon Oh, Sungchul Kim, Ryan A. Rossi, and Srijan Kumar. Influence-guiie gunnnee-guieedddd data augmenta- tion for neural tensor completion. Proceedings of the 30th ACM International Conference on Information & Knowledge Management, 2021.
[58] Laura Ruis, Maximilian Mozes, Juhan Bae, Siddhartha Rao Kamalakara, Dwarak Talupuru, Acyr F. Locatelli, Robert Kirk, Tim Rocktaschel, Edward Grefenstette, and Max Bartolo. Procedural knowledge in pretraining drives reasoning in large language models. ArXiv, 2024.
[59] Alon Albalak, Yanai Elazar, Sang Michael Xie, Shayne Longpre, Nathan Lambert, Xinyi Wang. Niklas Muennighoff, Bairu Hou, Liangming Pan, Haewon Jeong, Colin Raffel, Shiyu Chang, Tatsunori Hashimoto, and William Yang Wang. A survey on data selection for language models.
ArXiv, abs/a402. 16827, 2024.
[60] Megh Thakkar, Tolga Bolukbasi, Sriram Ganapathy, Shikhar Vashishth, Sarath Chandar, and Partha Pratim Talukdar. Self-influence guided data reweighting for language model pre-training. ArXiv, abs/2311.000 1.00913, 2023.
[61] Ritik Sachin Parkar, Jachyung Kim, Jong Inn Park, and Dongyeop Kang, Selectlim: Can 1lms select important instructions to annotate? ArXiv, abs/2401. 16553, 2024.
[62] Yunfan Shao, Linyang Li, Zhaoye Fei, Hang Yan, Dahua Lin, and Xipeng Qiu. Balanced data sampling for language model training with clustering. ArXiv, Ahs/2402.145266, 2024.
[63] Xiaonan Li, Kai Lv, Hang Yan, Tianya Lin, Wei Zhu, Yuan Ni, Guo Tong Xie, Xiaoling Wang, and Xipeng Qiu. Unified demonstration retriever for in-context learning. ArXiv, Abs/2305.04322) 2023.
[64] Liang Wang, Nan Yang, and Furu Wei. Learning to retrieve in-context examples for large language models. In Conference of the European Chapter of the Association for Computational Linguistics, 2023.
[65] Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in-context learning. ArXiv, abs/2112.08653 2021.
[66] Jeff Johnson, Matthijs Douze, and Herv'e e J'egou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, 7:535-547, 2017.
14

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificcton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In abstract and introduction section, we have clearly claim the contributions and scope of our paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
15

---

Justification:
We have discussed the limitations of this paper in the conclusion section.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We have procide full set of assumptions and complete and detailed proof in Appendix section.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: implementation details section, we have fully disclose the hyperparameters to reproduce the main results. We have provided additional information about the baselines implementation and evaluation in Appendix section. We will open source our model and code for reproducibility of our work.
16

---

Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We have uploaded our code and data as supplemental material with sufficient instructions.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / /tips.ccc/ppppppppp wttpccccppccccpeeecececcccc pwwpuuuguuuuuuss.eepepppppaeaaee gwwp///ieeeeeeee//ee ootee/emmmiiisiiiooooooo..oo.ooo ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
17

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have specified all important training and test details in datasets and implemntation details sections.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
We have included variance information of scores from repeated experiment in
Figure.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have included detailed information on the training and inference computa- tion resources in our evaluation section. Appendix provides additional details.
18

---

Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //tpe://////ss///// cc/ ac/pubblc/Ethi EthicsCuudellnes
Answer: [Yes]
Justification:
Our research conform, every respect, with the NeurIPS Code of Ethies.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We have discussed the broader impacts of out paper in the conclusion section.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
19

---

Answer: [Yes]
Justification:
We have discussed related considerations in the conclusion section.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have cited all the assets used in our research, and checked the licenses to ensure justified usage.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URI
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswithcon thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-pan se-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We didn't release any new assets but conduct all our experiment existing open-source datasets and models.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
20

---

14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
Our research does involve crowdsourcing / human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification; Our research does not involve LLMs any important, original, or con-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy (https: hnttps:/neuurps. http://uuurpps.c/ooeeeeeeeeeeeee for what should or should not be described.
21

---

A Deriving the Influence Function
Let l denote the minimizer of the empirical risk:
0" arg min L(0),
0
where L(0) = - ( Assume that l twice differentiable and strongly convex in 0. Thus, the Hessian
H = ( 1
v)((;00)
exists and is positive definite, ensuring the existence of H- The perturbed parameters 87,, are defined as:
0r,= arg min (I(0) + cl(=,)))
Define the parameter change due to the perturbation as A. : 0* Since 0 does depend on the quantity we seek can be written in terms of A.:
do
db
de de
Because 07,2 minimizer, its first-order optimality condition is:
VL(022) + +VV((2,
For small using the Taylor expansion around 0":
0 VL((") + + (F( A..
Neglecting higher-order terms, this reduces to:
[VL(0") A, R [+ (V ((((,
Since 0* minimizes I(()), it satisfies VT.(0") 0. Dropping higher-order terms in yields;
A H V
Thus, the derivative of 8% with respect to c at c = is:
d0 c.
-
--H''V((,00 V((2,0").
de
le=0
To compute the influence of upweighting 2 on the loss at a test point x, we use the chain rule:
dl(x, 0%
d
Vol(x,0")
de
de ===
c=0
Vollz; 0* H vol(z;0').
The influence function is therefore given by:
fır(2,))= = -vol(<;; " Vo((z;0))
22

---

B Deriving the Group Influence Function
For group influence function, we interest in removing a group of training examples.
Let l" denote the minimizer of the empirical risk:
0 arg min L(0),
where L(0) = - ((2;;)) 0).
By up-weighting subset examples = (21, Zm), we define the new objective:
(0) - n1
Σ (1 8)e(2i, 0) 0)+ + E)++))(2 2:48
2.88
we get the original loss function L(0) (where none of training samples are removed) and if = 1, we get the loss function where samples in are removed from training.
n-m where ē - e. can see
The perturbed parameters 0r be written as:
g = arg min (0))
1
Define the parameter change as A. = 0- 0 Since 0r does not depend on e, the quantity we seek compute can be written in terms of A
dl t.s dA.
de
B.1 Group Influence Function
Extended from Basu et al. [13] results, we show that the structure of the high-order group influence function is given by:
Theorem B.1 (Group Influence Function). Let 6(z) Hg 2 (2 and assume (2 for k is negligible. Then the k-order group influence function given by:
11 (x, ((x)
Σ
c
(k) au (1)))) d(zs),
(10)
2-ES
where: (0) are constants at the 1-th order, dependent the meta- meta-inforration of the distribution of and (2(2) is defined as:
ae(2) [1.
ift= 1, ift 2 2.
(()))
From Eq. 10, we observe that the interaction data samples exhibits an attention structure, with (2,) defining the weights of each data sample contributing to the final output. Intuitively, when k > 1, examples 2; that are closer to the group effect are assigned higher weights. We find the theoretical analysis of high-order group influence derived similar notion of "attention- based influence pooling" in AirRep, though AirRep's 's design can viewed as a simplified version of B.I. These analysis also motivate us to explore more complicated influence pooling mechanism in our future work.
The following section provide the derivation of B.
B.2 Deriving the First-Order Group Influence Function
Since 0* minimizes (f), the first-order optimality condition is:
0 VoLs(0),s).
23

---

For small expand around 0 = 8" and c 0:
(0L)(0)) + (VLL((0)) e + v0") A..
de
r=
At c = 0, the perturbed objective reduces to the original loss, so:
VoL(0*) 0. VoLs((0)) ==
The Hessian at c = is:
(
= ViL((*) H.
The first-order change in the gradient with respect to at A = 8 is:
i VoLs(0')
(ol(2),0").
de
c=0
- m
Substituting into the expansion, we have:
0 0 +
Σ
Vol(e,,0")
c + HA.
Solving for A::
A. - H
Σ
Vol(z,,
2,
Thus:
d0
H
Σ
Vol(z,
de
c==
n
2,ES
Applying the chain rule to measure the influence of the group on the loss at a test point x:
d (1(x, 02,s)
d0
Vol(x,
de
de
c=
c==
Substituting the expression for d0".s de we obtain the first-order group influence function:
S) - (2) Vol(x,
"
Σ
Vol(zs,
where (())
B.3 Deriving the Second-Order Group Influence Function
Since 0r minimizes Ls. its gradient vanishes:
0 = VoLs(0
We perform a Taylor expansion around e = 0 and 0 = 0", keeping terms up to e², Let:
A. A Ag + O
First-Order e At = O, (0) reduces to L(0), so:
VoLs 1 VoL(8") = 0.
For the linear term (6²), have:
0 Rd
VeLs(0',
+ (0
de
c==0
Hence:
A
Vol(zi, 0")
m
This result matches the first-order group influence function.

---

Second- e For the quadratic term (c²), expanding the gradient VoLs and collecting all c? contributions yields:
1 8²
i Vo
0
2 de VoLs(0
A
de
c=0
+ H +
A [
2
Solving for A2:
2
H
8 VoLs (f*) de
+ & ViLL 0 1 A
+ I V2L(0))
Since the weights of the training examples are linear in e, the second derivative with respect of these weights is zero:
02 VoLs(0*)
2 de
=0
1==0 Assuming higher-order terms (=; for k 2 3 negligible, the (OT((") term vanishes. Thus:
A -
(OL)((0)
A1.
de
From the definition of Ls, we have:
1
A
H
Vol(ze 0*)
n
2,68
m
Vil(zi, 0+)) م
Rearranging:
[I I-- H Hss A1
A2
m
where Hs 1 Ex,es Vie(e,
Influence on the Test Point For test point consider f(x; 0-, 0r,s). We aim to expand it as a Taylor series in 0 around 0* Following previous studies, we perform first-order expansion rather than second-order one. This choice justified because the dominant second-order effects can typically be captured through the parameter shift A.. Thus, we obtain:
fxx (0. N ((x,0)) Vol(x,0") AiE
+ Vol(x,0)) Ag 2.
The term multiplying corresponds to the first-order group influence, while the term multiplying ?² the second-order correction. In closed form:
First-Order Group Influence:
A () =V T A
Second-Order Influence:
APP (2, = vol(2,00) D Vol(x,00)) A2
where:
A2 n m
I--H 'H A1.
4 H m Vollzi,0")
25

---

Simplifying Terms
Now clean up the expressions ((2) (x, First, define:
<(z) H Vol(z, 8*),
where ((2) represents the preconditioned gradient of the loss.
Rearranging f (z, S), we get:
f (x, 0z) A m (z, T A
m
Term A
m f( H
Term
where terms A and B emerge.
Substituting 6(2) in A, we have:
A
((x) Σ ((()). 71
2,68
For term B, have:
B Vol(( 0²) n- m
H H
Approximating the subset Hessian Hss using the Fisher Information Matrix (FIM):
Hs Vil(z, 0*)
Σ
TT)
2,88
Substituting this approximation into B, we have:
B ~ nn m)²
((x)
Σ o(2)) ((2))(22)
Σ p(2s)
4:88
4E8
For further simplification, we define:
S) = x(z))
Σ [(2)),
RJES
which measures the alignment of ((2;) with the aggregated gradients in S. Then:
B R2 (n- m)=))202)
Σ
[a(2, S))(2.)
2
Putting together, we get:
(2)
- A+
= =c
Σ ((2)) + (2)
Σ
Ja(z, S))(z))
FIES
21ES
= ((x) Σ
(60)) + (2) azi, ((z))
2,E3
((x)
Σ
Σ
(00 S)
where c( and (22) are the constants of the first- and second-order terms, defined as:
(1)
7
(2)
(n m2
26

---

Additionally,
a1 (1 5) 1, a2(zi, S) (((z)) T Σ 6(zz)
2,ES
are the first- and second-order alignment functions.
B.4 Deriving the Third-Order Group Influence Function
For many practical purposes, second-order corrections already capture the leading "beyond first-order" effects. However, order to have more complete perception the structure of the influence function (and also to address our curiosity), we will continue to derive the third-order group influence function.
Since t minimizes Ls. its gradient vanishes:
VoLs((0,s)
We perform a Taylor expansion around e and 0", keeping terms up e², Let:
A. = A + A2 2 + As e O(c*).
The first-order and cocond-order terms, A₁ and A2, have been derived in the previous section. Now focus on the third-order term, Ag.
Considering the e² term and collecting all contributions of order e², we obtain:
0
VOL) de
VoLs(0,, 0) A
c=0
(i)
8
V L
(0) 0)
A
de
c=0
i
ViLs(0,, 0) 8e
A
(i)
+
VoL (s (0" 0) A
(v)
+ 1/2 A
v Ls (0" 0)
A
de
(v))
Now analyze these terms in turn:
Term (i) vanishes since the per-sample weight depend linearly on
VoLs(O*)
-.
aea
c=
Term (ii) can be written as:
V2Ls(00, 0) A3
H H
Term (iii) vanishes because the per-sample weight is linear e:
0²
(iii) = Se
(2Ls(00000) 0)
A
27

---

Term (iv) is:
(iv) 0
Vi (0, 0)
2
de
Similar the previous derivation, this can be rearranged as:
(iv) m
H H A2.
where H = = E Vil(z,, 0-).
Term (v) involves the third derivative with respect to 0. In most cufluence-function deriva- tions, one often assumes that higher-order derivatives (2 are negligible. Following this convention, we consider:
(v) A V/Ls(0 0) A 0.
Term (vi) is the derivative with respect to of the third derivative with respect to 8. Since we neglect Vge. any derivative of it also vanishes. Thus, we have:
(vi)
A
② de
VOLs(0 0)
A O.
===
Putting everything together, the e stationarity condition simplifies to
HA + n m
H H.s
A2.
Solving for As yields
m (1 -- H-H- )2 Ag
Influence the Test Point For point æ, consider (sr, 0r 07,s). Similar the previous section, we can express the effects of on æ (when considering the et term) as:
JI (x) (x, =vo() (A
The terms related to A and 2₂ have already been simplified in our derivation the second-order group influence function, and we write them as the following simplified expression:
f (x. ((x) Σ o(2:) + (x))
Σ [a(2,, S)o(z)]
(11)
FLES
Term
Term
Therefore, the first two terms and B of ( (x. S) share the same structure, with constants slightly adjusted due to the inclusion of additional correction terms. Moreover, / (x, will have an additional term C for the e² term. Let:
( A+B+C.
Now, lets simplify C:
C
Hg H Σ Vol(z,
2.53
Approximating the subset Hessian Hg using the Fisher Information Matrix (FIM):
Hg Σ 0*) ~ Σ
2,ES
2.ES
Substituting this approximation into C, we obtain:
C (3) ((x) Σ
((2))()))T)
E ((z))).
#E
2,
28

---

For further simplification, we define:
B(z, o(2;)² T Σ [a(z) 5) ((2))
where oz22, S) the same as the previous definition, that is:
al2, ((x)) T
Σ o(2)).
RyES
Then, express C' as:
C ~ (3) o(x) Σ
B(zz, S) o(z)
2:58
Combining this with A and B, we obtain:
f (c, = B+++C
(()))) Σ Φ(2i)
2,ES
+ ( ((x)
Σ
[a(2, S)((())
+ (3) (xx)
Σ
B(zz, S)zz,))((z)) )(x))
We can then rename a, 8, and the constants to obtain:
3
(c, = (bx) Σ
(0) Oz (zz),
where of are constants defined by the distributional properties of S, and
an S) = 1, a2 (2), 5) = 6(2)) Σ o(2)),
2,E9
as(2;, 8) x(2;)²
Σ
02(2,, S))(z))
s,
These are the aligmment functions, which can also be defined recursively as:
(x(2) 1, [t(2) ((5)
if = 1, if 2 2.
This gives the same format Theorem B.I.
B.5 Deriving the Higher-Order Group Influence Function
We extend the previous derivation higher-order terms. Before manually doing this, lets first review the effective terms (i.e., nonzero terms) of the first, second, and third-order terms:
For first order:
a VoLs(000
0 RS de
+ (V)
l==0
For second order:
① 0z VoLs(00)
A + V,Ls(00,
de
c=0
29

---

For third order:
a 0 de (0 .c) A2 + (i
We see that at each order e, the only "new" term that enters the equations (and thus the recursion for A&) is the first derivative with respect to multiplied by lower-order Ay ss, plus the usual expansions in powers of A. (where A. = f-.s- -
To explain this, we start with the stationarity condition at all
VoLs(00 + A.,t).
Define:
c(0,c) :- VoLs((())
Since 8* is the minimizer at O, we know -) =0
Expanding G(0 + De.e) E) using a multi-variable Taylor series and keeping terms up first order in
A, gives:
artrG
(or (o:.0)l121e)
G(0 Ar,t)
p!r! OePO0"
~G 0) + (00, ,0
0P
pl dep ((0",)ee"
Op
(3 ES.
p! dev
p=1
Define:
ô H' V2Ls(0.C) c) de
= VOLs 0),
0=8*
Because each training sample's ss weight depends linearly on e, we have:
2 dep
weights
for p 2 2.
Thus, oigher-order terms in In other words:
0²
VoLs. de3
8 VoLs. etc., vanish, simplifying the stationarity expansions.
Oe VoLs(0 0)
G 0)
= and
= 0.
der
der
p22
p22
Furthermore, collecting the & terms in the stationarity condition G(0 Ac,e) 0 for k 2 gives:
0 = H A + H' AR-
Since Ao 0 by definition, and A is determined by [ G(B, 0) we solve for:
A H H'
By repeated substitution, we obtain closed-form expressions for all Ag:
Order 1:
8 de G 0)
A -H
Order :: >
= -
This leads to the recursively defined alignment functions the final group influence function formula. This completes the proof of Theorem B. for all cases where lo > 3.
30

---

C Comparison between gradient-based and representation-basel methods
From Eq. 3 and Eq. 5, it is evident that gradient-based and representation-basse methods share a similar structure: both rely on the inner product of two vectors and differ only in terms of the calculation of these vectors. Specifically, the former utilizes gradients and and Hessian, while the latter uses embeddings. Additionally, under typical assumptions both classes of methods produce group influence estimates via the same aggregation process (i.e. simply summing per-example influence).
Apart from computational considerations. what factors might influence the choice between a gradient- or representation- based data attribution method? We highlight three considerations:
Accuracy:
whether the attribution results are reliable. Gradient-based methods are designed to simulate the model optimization process, which are theoretically accurate under certain assumptions. However, existing representation-bast based methods rely on heuristic metries that are agnostic to the objective of data attribution, leading to sub-optimal results.
Computation cost: Gradient-based methods generally rely on the computation of gradients and some approximation of the Hessian, making them computationally expensive. In practice, the cost of gradient-based methods is comparable to training the model on the all of the target examples because computing influence involves computing loss gradients for each target example. In contrast, representation based methods are efficient and applicable for large-scale retrieval [66].
Storage cost: Influence functions can be particularly storage-intensive as they can involve storing the full gradient (equal in size the model parameters themselves) for each data sample, which prevents large-scale data attribution. Follow-up work projects gradients with random matrices [22] or low-rank approximations [25] at the cost of reduced accuracy. Representation-be based methods output fixed-size embeddings, which are more manageable in size.
In summary, representation- based data attribution methods are cheaper than gradient-based method but do not reflect the actual process of model optimization.
D Details of Evaluation Datasets
We use four datasets for evaluation:
We randomly sampled 100,000 examples from the original FLAN training set as the TDA evaluation training split. For the test set, we retained the first 100 examples of each task in the test set. Since there are 66 tasks, the test set contains 6,520 examples.
FLAN:
Alpaca: We use the original Alpaca training set (https: / htt///uugggggaccceee ct/ co/datasets/ alpaca) catsu-labe lapappaa) as the training data. For the test data, we design the test set to consist of two parts: (i) seen subset, where we randomly sample 250 instructions from the training set and use Llama model trained on the full Alpaca training set to generate the responses; (ii) an unseen subset, where we sample 250 examples from https: hhtgg///gggggggcce co/ datasets/tatsur hatss//ttstllaaaaaaaeeae caatsu--abbalppcc alpacalpwaaaaaaaaa _eall. Thus, the total test set consists of 500 examples.
Tulu: We use version v1 of Tulu (https: hatg:///huggggaeeeeeee
tulu-vi-se www--vl--sstuuureeeeeeee-eaa--te
The original dataset contains over 490k instances. We ran- domly sample 100,000 examples for our training set and 500 for our test set.
SafeRLHF: SafeRLHF (https:/ httg://ghhggnaccc//e co/ /ddaassssss/PK Alignment/ SrU-SafeRLHF QA) is a dataset for safety alignment in large language models, where each example has a safety label of either "safe" or one of 19 possible harm categories. We use the original training set. For the test set, we sample 500 examples from the original test set, ensuring a 250:250 ratio of safe and harmful instances for label balance.
See Table 3 for statistics of evaluation datasets.
E
Details of Baselines
Our baselines include two categories: gradient-baser methods and representation- based methods.
31

---

Table 3: Evaluation Data for data attribution. # Train and # Test denotes number of examples in training and testing set.
Name
# Train
# Test
FLAN
100,000
6,520
Alpaca
51,760
500
Tulu
100,000
500
SafeRLHF
251,963
500
Gradient- based methods rely on the gradient of data with respect to a warmed-up model. To warm up the model, for FLAN, fine-tune it 100,000 examples randomly sampled from the FLAN training set. For Alpaca, Tulu, and SafeRLHP, we warm up the model on 100,000 examples randomly sampled from UltraChat.
LoGra LoGra is a toolkit for efficient influence function calculation (https: pectps://githbb thub com/ and its core techniques are introduced Choe al. [16]. Its key idea is to use low-rank gradient projection to achieve implicit gradient projection during backpropagation, thereby reducing computational cost. We experimented with different configurations and found that the following variants achieved the best performance in our evaluation (as also described in Eq. 3): (i) applying the LoGra module all MLP layers, (ii) initializing the LoGra weights using PCA, (iii) approximating the Hessian using KFAC, and (iv) performing unit normalization of vectors.
The rank value of LoGra controls the projection dimension. We set 8, and 16, resulting in final dimension ranging from 1,152 to 27,648.
TracIn TracIn [22] is derived from first-order approximation of model weights and essentially computes the dot product between gradients without Hessian correction. For implementation effi- ciency, we use LoGra to perform gradient projection and apply unit normalization to the vectors. The original TracIn keeps multiple checkpoints and ensembles their influence scores, but this approach linearly increases computation and storage costs. Instead, we use only one checkpoint (i.e., the final checkpoint) of the warmed-up LMM
LESS LESS [25] adopts LoRA tuning instead of full fine-tuning and projects the LoRA gradient to a lower dimension using random matrix. Its key idea to use AdamW states (first- and second-order momentum) to correct the gradients. We use the official implementation (https: //gi http///rinnnooonnn.//.ae //l//LSSS) and set the projection dimensions 768 and 8192, which are the default dimensions in the LESS paper.
For LM warm-up, follow their paper by setting the LoRA rank to 128 and alpha to 512, fine- tuning for epochs. Similar to TracIn, the original LESS method ensembles influence scores from checkpoints. For fair comparison, we only retain the final checkpoint.
Another note is that the official implementation does not support batched gradient calculation, requiring a batch size of 1, which slows down execution. To better utilize the GPU, we run multiple processes on single GPU to accelerate computation and report the computation time. However, its speed remains lower than LoGra, which supports batched computation.
DsDm and TRAK TRAK [12] is a representative gradient-bar based method that provides an efficient implementation of the influence function. TRAK approximates the Hessian using the Fisher Informa- tion Matrix (FIM), ensembles influence scores across multiple checkpoints, and adjusts the influence score of each data sample by multiplying it with the model's prediction probability. DsDm [33] later applied TRAK to LLM data selection tasks. Since DsDm has not released its official code (https://gi) thub. com/Ne WaadrLLabDDsDii the Logix framework. where the Hessian is approximated with FIM (also known as raw Hessian approximation in Logix), LoGra is used to project the gradient to different dimensions with random weight initialization, and only the final checkpoint is applied without checkpoint ensembling.
we approximate its implementation (i.e., TRAK) using
The following citissenntttioo-baaeee baselines.
32

---

2231
2131
LoGra
20
2042
09.3
19.75
AirRep
15
10
3112
500
800
1K
5K
10K
Subset Size
Figure 8: LDS performance of LoGra and AirRep under different evaluation subset sizes. AirRep consistently outperforms LoGra, especially on larger data size. Note that AirRep is trained only on data with subset size of IK.
TF-IDF
We sklearn to calculate the cosine similarity of TF-IDF feature of data.
DSIR DSIR utilizes hashed n-gram features, where n-grams mapped onto a fixed number of buckets for efficiency and scalability. The influence estimator is parameterized by a aag-of-wordh generative model the hashed n-grams, learned simply by counting the hash bucket frequencies. We use the official implementation (https:/ //gi thub. com/p- cmbbda/dsir).
RDS RDS was introduced Hanawa et al. [6], which uses the model's hidden states as data representations. Through evaluation of different representation design choices, average pooling over the last-layer hidden states of all tokens as the sentence embedding.
GTE GTE [35] general text embedding model trained on mixture of text retrieval, text classification, text clustering, and text entailment data. It achieves strong performance text embedding leaderboards relative to its model size (e.g.. https: : htt::///ggggaaeeeeeeeee httpagginnnae..ceeppeee/e/eeee/e cc/sspace//mtee hetece////eb/////e/eee leaderboard). For direct evaluation, we assessed various GTE model sizes, including Small, Base, Large, 1.5B, and 7B.
The following are discussion about other TDA approaches which we did considered for empirical comparison.
Datamodels AirRep is inspired by Datamodels [1]. in terms of problem formulation for data attribution and both methods focus on estimating empirical group influence scores. A key difference lies in how embeddings are generated. Datamodels assign fixed embedding vectors for all instances in given data collection, and require the full retraining when new instances are added. AirRep, on the other hand, uses trained encoder produce the embeddings for any new instance on the fly. result, AirRep is computationally much more efficient than Datamodels in handling dynamically changing data. Empirical findings in [12] shows that Datamodels typically incur inference costs around 100x greater than gradient-based methods to achieve comparable performance. AirRep addresses this issue with the trained encoder, and being 80x more efficient when handling new instances. We will add these discussions in our revised version of the paper.
Training Cost
For data generation, we create 10K subsets (each containing 1K data examples) and train the model on each subset for epochs. This results in total of 10M training examples for the Qwen2.5-0.55 LM, requiring about 20 hours on eight A100 GPUs. During the optimization stage, each training step involves sampling
subsets K examples each), totaling 32K examples per gradient descent step. The model is trained for up to 2K steps, completing in about 5 hours on eight-GPU machine. While this training cost exceeds the warm-up training requirements gradient-based TDA methods, it is exchanged for significant efficiency gains during inference. See Section 9 for detailed analysis.
33

---

LoGra (Qwen-0.5B)
16.79
LoGra (Llama-1B)
19.95
AirRep
22.54
5
10
20
LDS on FLAN (Target LM = Llama-1B)
LDS scores of LoGra and AirRep when the target LM is LLaMA-1B. For LoGra, we evaluate both variants: LoGra applied to Qwen-0.5B and LoGra applied to LLaMA-IB. AirRep is trained on Qwen-0.5B data and directly evaluated using LLaMA-IB labels. Despite the model shift, AirRep achieves strong LDS performance, even outperforming LoGra applied to LLaMA- demonstrating its ability to generalize across similar LMs without re-training.
Figure 9:
Table 4: Training Cost of AirRep. Asymptotic complexity and wall-clock runtime (measured as eight A100 GPU hours) for each training stage. n denotes the size of each subset, N, represents the total number of subsets, E is the number of epochs for training on each subset, M is the number of sampled subsets per training step, and T is the total number of training steps for AirRep.
Stage
Data Generation
Optimization
Complexity Actual Value
O(n N₂ E) O(1K 10K 2)
O(n M M.T)
O(KK 32 2K)
G Data Selection Setup
Our data selection evaluation is conducted independently for each task in FLAN, using a greedy search strategy to identify the most relevant training data. For a given task, the test set contains m examples (typically
100 due to our data sampling strategies), denoted as (æ1, .2, .mm) Let training set be (z1, : =n). The TDA method computes an influence score [(x;, zj) for each pair (x,, 2j), where
[m] and j E [n]. Next, determine the rank of each training example relative to a given test example. Specifically, we define rank 2,, xyy) E [n] as the position of xy among the most influential training examples for 24. In other words, if rank( (2,, 1, then xy is the most influential training example for 2;. We then define the score of each test example 2y as the best (lowest) rank it achieves across all training examples: Score(x,) This represents the highest influence rank of x; among all test examples. Finally, we rank the training data based Score 2ty) and select the top 1,000 examples.
min(rank(2;; 22)
[ml).
After training the target L.M on the selected data, we evaluate the model on test set, setting the maximum generation length to 64. We use unigram F1 as the evaluation metric for all tasks, even though some tasks have other conventional metrics, such as accuracy for classification, ROUGE for summarization, and BLEU for translation. Notably, unigram Fl is correlated with these task-specific metrics since the model's generations are typically very short. For example, in sentiment classification task, the model is expected to generate single token (positive negative), making unigram Fl equivalent to accuracy. Finally, the final F1 score averaged over 66 tasks FLAN
For gradient- based methods like LESS and LoGra, which can be configured to balance computational and storage costs, we use a 0.5B LM as the proxy model to compute gradients. We set the projection dimensions to 768 and 1152 for LESS and LoGra, respectively, as this ensures a comparable storage requirement to AirRep (which uses a 384-dimensional representation) while remaining practical configuration for data selection from large-scale corpora
34

---

H Additional Evaluation
H.1 LOO Evaluation
To analyze the effectiveness of our method on individual data influence estimation, we conducted a leave-one-out (LOO) evaluation, in which individual examples were removed and we measured the correlation between predicted and actual influence scores. The results on a FLAN subset are summarized in Table We find that AIRREP outperforms baseline methods in this setting as well.
Table 5: LOO correlation between predicted and actual influence scores FLAN subset.
Method
LO0 Correlation
LoGra (18432)
10.75
GTE-Small
0.72
AirRep
15.36
H.2
Data Classification Results
Table 6: Accuracy of Data Classification.
Method
FLAN
Tulu
SafeRLHF
LoGra (Dim 1152)
71.61
76.60
78.00
LoGra (Dim 18432)
85.44
86.00
83,20
DsDm (Dim 18432)
83.68
85.34
80.25
TracIn (Dim 18432)
77.69
79.43
67.12
GTE-Small (Dim 384)
50.59
76.60
90.60
AirRep (Dim 384)
86.41
88.20
87.20
I Detailed Data Selection Results
See Table 7, 8, 9 for detailed results of data selection.
35

---

Table 7: Data Selection Results for Qwen2.5-0.5B LM.
Data
Random
DSTR
TF-IFF
GTE
LESS
LoGra
AirRep
acsic
16.86
26.61
27.17
25.53
27.2
29.61
ag nessesssbbbt
76.86
67.61
79.41
74.75
82.85
82.97
77.72
anii
28.14
28.02
38.9
36.04
50.44
40.13
44.49
anli r2
34.02
29.7
40.17
32.91
35.01
41
36.58
anli
36.55
24.47
35.55
37.54
36.7
36.17
41.98
are challenge
are eas_
46.38 43.27
43.71
48.45
51.89
48.86
48.8
bool_g
49.0 48.0
54.0
53.0
56.0
55.0
56.0
12.7 31.81
38.73
34.64
56.67
70.6
56.88
cmn, dailymail
19.68 12.65
24.19
23.09
22.83
21.76
26.99
cola
57.23 62.18
63 23
63.18
61.48
63.18
63.43
common
41.69 45.23
45.78
43.2
46.26
43.41
45.1
copa
60.59
61.79
63.94
10.99 6.36
39.21
28.72
37.42
41.98
41,86
cosmos_qu
39.02 40.01
39,92
42.04
43.17
43.15
39.54
dart
62.73 66.65
66.08
68.98
67.42
67.39
67.7
definite_pronnns pronoun
46.07 52,73
51.0
58.0
$7.07
54.67
$4.07
drop
8.08 13.49
19.35
20.0
14.81
e2c_nig
63.48 68.12
67.91
65.83
66.34
68
65.82
fix_pennt
95.08
95.17
94.93
95.2
95.24
gigaword
26.59 26.13
28.46
27.14
28.68
28.06
29.73
glue_mppe
55.0
69.0
51.0
$7.0
66.0
68.0
73.0
glue qup
66.0
72.0
78.0
74.0
78.0
71.0
79.0
hellaswug
29.31
28.89
37,21
38.62
45.81
46.02
43.3
indd_revice
64.85
70.43
70.41
math dataset
3.0
2.0
7.0
6.0
3.0
1.0
4.0
maicched
mnli mitmmtheetehee
77.0
73.0
74.0
74.0
81.0
76.0
78.0
mulii news
15.45
3.38
19.57
18.56
10.98
19:47
19.24
multire:
11.06
0.34
37.26
51.0
36.0
29.0
36.0
natural natural_questinn questions
6.61
8.57
10.41
7.04
7.43
7.79
6.77
openbookga
opinion aatiroutssbtssssssooma idebute
11.13
9.91
17.85
20.44
20.1
18.82
21.31
opinion abbttrctt, rotten
9.9
sara_crawweer
46.02 43.37
48.66
47.28
45,21
46.23
46.03
paws_ _iki
53.0 51.0
61.0
63.0
83.0
91.0
83.0
piqa
74.34 76.46
77.71
77.91
76.49
79.53
78.45
quli
68.92 44.49
68.98
60.51
74.58
78.81
71.17
quac
record
18.38 16.42
18.38
16.45
17.55
16.1
10.8
56.
57.06
28.03 13.63
34.45
34,95
34.29
35.19
36.85
sentiment 140
48.35 40.95
33.68
40.36
49.02
47.8
39.9
soli
77.1 77.94
76.21
78.24
79.94
77.14
squad_vi
28.33 26.42
34.24
38.07
39.19
37.87
40.24
squad_v2
sat2
69.13 68.0
70.23
71.39
69.27
71.61
69.76
story_cloze
57.71
64.47
78.17
81.92
77.51
stsb
25.74 32.24
29.
35.92
34.52
33.26
41.0
tree
48.24 79.0
75.89
63.39
78.34
75.25
72.41
9.05 10.49
10.97
true true_case case
96.64 97.28
98.71
98.07
98.98
99.43
99.22
bibb_nng_ee wic
72.7
55.0 56.0
52.0
44.0
49.0
55.0
40.0
wiki wwlllinnnnngliiiie
11.87
18.56
17.9
19.54
16.91
20.95
wmr14 enft
46.07
37.69
44.19
46.46
46.74
46.78
47.17
16 translate _sen esen
24.58
21.91
24.57
24.75
26.8
25.78
27.68
tra.staaaaeeecoo doen
wmi16 translate fien
22.54
19.79
24.18
22.13
22.93
22.61
24.72
translate
30.65
32.56
35.65
wmnt 16 translate maen
39.96
35.99
40.59
39.13
41.4
40.67
40.11
wni16 @tt.anslltc_ttoc tren
29.4
24.42
26.85
28.88
25.61
27.44
29.5
57.14
57.14
61.43
58.57
58.57
54.29
$7.14
word _ed_metn
71.65
87.1
82.68
80.29
83.79
85.93
91.07
50.0
yelp. Satlppolliity..mmm potaarttyreevveeewss reviews
72.95
63.22
73.24
73.66
73.53
71.81
74.46
36

---

Table 8: Data Selection Results for Qwen2.5-1.5B LM.
Data
Random
DSTR
TF-IF
GTE
LESS
LoGra
AirRep
acsic
21.59
20.66
28.66
28.56
27.0
25
27.05
ag nessesssbbbet
79.47
79.37
77.22
80.21
82.15
80.5
82.27
anii
44.93 44.46
48.75
50.67
51.04
59.91
54.44
anli r2
47.12 45.21
44.7
37.04
42.31
48.8
53.12
anli
39.98
36.8
40.05
43.03
44.8
42.91
43.76
are challenge
are eas_
59.02 61.85
58.46
64.19
60.65
63.54
64.6
bool_g
60.0
48.0
63.0
62.0
74.0
78.0
72.0
30.86
48.2
29.15
71.1
77.04
53.02
75.22
cmn, dailymail
18.45
9.19
28.8
27.6
27.29
27.97
27.69
cola
50.29 72.07
63.32
73.17
64.15
45.32
common
47.56
50.01
50.08
49.44
48.26
$1.34
47.91
copa
73.35
75.61
74.53
73.37
72.45
9.78 4,89
31.64
26.47
36.81
32.24
42.75
cosmos_qu
46.48 44.01
45.56
49.96
50.01
49.81
49.89
dart
65.46
66.73
68.79
69.85
66.71
68.57
68.29
definite_pronnns pronoun
50.89
50.0
53.67
$2.0
59.0
52.0
$7.67
drop
14.24 16.78
17.8
23.23
22.89
20.68
e2c_nig
66.71 67,82
68.02
67.35
68.41
67.05
67.78
fix_pennt
94,98
95.32
95.39
96.25
95.69
gigaword
29.49 27.74
29.92
28.86
30.04
29.36
31.02
glue_mppe
49.0 70.0
76.0
68.0
73.0
58.0
67.0
glue qup
78.0
84.0
78.0
81.0
81.0
78.0
hellaswug
45.13
29.01
49.58
54.93
64.47
58.89
60.07
indd_revice
28.44
71.99
67.19
66.67
66.76
math dataset
3.0
3.0
8.0
3.0
1.0
3.0
4.0
maicched
79.0
mnli miimmtheetehee
84.0
83.0
79.0
79.0
86.0
84.0
83.0
mulii news
12.63
4.33
19.39
19.27
11.65
19.68
19.69
multire:
22.7
1.37
27.45
23.48
46.0
53.0
48.0
natural natural_questnn questions
11.99
14.75
16.47
15.17
11.82
14.43
13.86
openbookga
opinion_ aattious_bbstttss__cm idebute
12.47 13.78
19.48
22.28
20.55
21.71
23.96
opinion aabttrctt, rotten
aaru_crawweene
$1.87 52.19
53.94
52.89
52.9
52.94
53.59
paws_ _iki
62.0 71.0
77.0
82.0
87.0
86.0
92.0
piqa
77.96
81.91
80.38
80.03
81,1
80.32
quli
77.9 46.08
61.87
76.73
78.8
78.08
79.22
quac
record
15.65 14.58
13.45
15.16
15.67
12.53
10.62
62.85
58.81
34.09 21.79
40.03
40.6
38.11
38
39.88
sentiment 140
42.67 46.46
47.87
49.02
48.57
50.72
42.49
soli
83.1 83.46
86.58
83.46
82.94
76.17
squad_vi v1
44.56 43.91
45.59
48.18
48.3
46.95
$1.21
squad_v2
22.24
sat2
70.83
72.22
70.32
71.79
74.05
75.64
74.07
story_cloze
79.2
82.
85.49
86.29
86.
stsb
26.88 43.59
37.17
35.78
29.36
25.66
41.1
tree
62.52 82.03
82.03
77.41
78.89
79.67
76.41
29.45
18.42
21.0
19.
20.09
true case
98.75
98.99
98.88
99.5
99.75
99.73
99.75
bibbnnng_ee wic
76.42
57.0
43.0
47.0
63.0
45.0
46.0
56.0
wiki _iki_lingoa_com tanguunnnglish_om
12.41
19.16
20.22
20.49
16.25
19.87
wmr14 enft
55.94
51.51
55.44
53.22
55.23
56.22
55.78
16 translate esen
35.5
33.26
36.68
35.76
37.13
36.96
39.24
tra.ssaaaaeeccom doen
55.78
wmi16 translate fien
29.23
27.37
30.5
30.86
27.53
30.09
31.11
translate
37.97
wmnt 16 translate meen
47.36
45.48
45.56
48.7
45.91
47.82
49.07
wni16 -tt.nsslltc_tooo tren
38.49
34.42
38.3
40.47
38.33
36.59
38.79
62.86
57.14
67.14
58.57
44.29
50,0
67.14
word _eg_metn
82.06
90.72
90.31
90.78
87.07
89.64
90.36
yelp. Saslppolliity..mmm potaartty_eevieeews reviews
74.94 66.41
77.04
62.37
76.02
76.68
74.79
37

---

Table 9: Data Selection Results for Qwen2.5-3B LM.
Data
Random
DSTR
TF-IF
GTE
LESS
LoGra
AirRep
acsic
22.94
25.58
31.06
26.32
ag nessesssbbbet
82.79
49.79
75.93
78.6
80,28
79.08
79.92
anii
56.4 43.82
57.49
52.64
45.6
62.64
55.25
anli r2
37.08 39.16
52.42
36.01
46.08
44.18
18
anli
38.95
44.1
43.07
39.63
44,86
58 84
48,98
are challenge
are casy
60.46 64.55
66.64
66.44
67.47
64.67
66.84
bool_g
72.0
76.0
66.0
66.0
76.0
76.0
36,95 68.78
73.76
70.97
68.8
67.18
76.89
cmn, dailymail
21.55 8.77
29.1
27.4
15.03
26.56
29.74
cola
63.34
71.0
52.13
43.46
78.58
79.38
common
49.08 47.42
48.97
48.44
47.24
48.72
48.36
copa
68.18
72.25
76.38
77.18
5.75
2.51
31.11
24.45
27.37
43.23
43.33
cosmos_qu
53.96
51.83
54.15
$2.9
54.28
56.99
54.0
dart
65.09 66.37
69.82
70.53
69.55
69.06
70.21
definite_pronnns pronoun
52.0 56.0
52.0
55.0
49.0
53.0
63.0
drop
15.29
20.24
27.99
22.72
22.09
e2c_nig
64.11
66.03
67.87
65.88
66.93
66.71
66.38
fix_penne
94,46
95.88
95.4
95,
95.57
95.76
gigaword
28.77
27.11
31.32
31.37
29.66
29.92
32.01
glue_mppe
70.0
74.0
53.0
32.0
79.0
72.0
50.0
glue qup
79.0
79.0
75.0
78.0
79.0
76.0
80.0
hellaswug
57.54
49.18
61.75
56.11
64.21
68 58
69.16
indb_revicws
67.74
27.28
70.13
71.74
math dataset
5.33
6.0
8.0
S.0
5.0
7.0
5.0
maichee
mnli millmthsmeeheed
87.0
89.0
67.0
87.0
89.0
87.0
85.0
mulii news
6.42 3.29
20.25
20.29
8.2
20.33
20.26
multire:
24.08
2.94
45.53
$2.0
54.0
54.0
46.46
natural_ questions
16.98
19.6
15.28
15.28
18.
17.0
17.78
openbookga
opinion_ aattiout_bbstttss__cm idebute
14.6 10.67
22.65
21.71
19.8
20.59
23.67
opinion abbttrctt, rotten
17.36
mara_crawweer
$4.92 56.77
55.78
55.94
55.08
55.51
55.55
paws_ _iki
69.0 76.0
54.0
81.0
92.0
87.0
89.0
piqa
78.8 74.89
78.67
79.4
81.72
79.95
80.78
quli
71.33 75.33
78.95
66.37
79.61
70.57
77.84
quac
17.9
record
19.74 17.4
19.5
17.01
15.33
14.88
16.39
54.7
67.95
32.46 21.29
40.79
41.16
38.62
38.67
42.4
sentiment 140
47.71 45.72
48.66
45
48.52
50.44
38.4
soli
74.1 59.36
85,
81.94
72.16
83,1
squad_vi v1
40.29 44.99
41.19
48.93
$1.69
45.93
$4,25
squad_v2
59.77
sat2
70.33 65.0
73.63
74.87
73.48
73.65
72.93
story_cloze
76.82
84.51
84.14
86.27
86.43
stsb
37.82 35.56
38.
33.6
29.03
42.23
37.22
tree
64.12 80.41
$7.66
74.58
85.07
79.85
80.41
35.72
31.63
30.28
26.53
true true_case case
97.48 99.22
98.99
99.21
99.45
99.41
99.38
bibbnnng_ee wic
78.76
75.37
56.0 44.0
54.0
44.0
44.0
42.0
44.0
wiki Bai__ingga__nngil
17.6
10.1
19.76
17.7
16.1
19.82
wmt14 enft
55.52.
55.11
57.02
57.87
57.51
58.27
57.11
16 translate
_ecn esen
38.36 41.69
41.63
40.6
41.45
39.79
43.68
tra.saaaaaeeecoo doen
55.14
56.85
$7.45
wmi16 translate fien
31.17 31.53
33.66
34.73
32.27
33.07
32,54
translate
wmnt 16 translate maen
51.89 48.51
52.63
50.32
52.01
52.94
51.71
wni16 @ttrnsslltc_too tren
40.43
42.03
43.06
41.96
42.4
43.57
42.85
58.71
58.57
45.71
65.71
70.0
61.43
47.14
word _ed_met
85.99
90.12
93.99
89.57
92.07
93.04
92.66
36.0
35.0
yelp. Saslppolliity..mmm potarrtttrreeeeeewws reviews
69.79
62.21
76.74
74.89
74.93
74,1
76.97
38


Reviews:
[{"review_id": "dmYzsv7mBW", "reviewer": "Reviewer_VLz2", "summary": "The paper proposes AirRep, a representation-based training-data attribution (TDA) method for large language models (LLMs). AirRep preserves the bilinear test-rep × train-rep structure of classical influence functions but dispenses with gradients and Hessian inverses. Instead, it (i) introduces a soft-attention pooling layer to model higher-order, group-level interactions among training points and (ii) jointly learns both the encoder and the pooler on automatically generated pairwise rankings derived from true re-training losses (Datamodel-style supervision). Once trained, AirRep scores influence with a single forward pass, delivering ~80× faster inference and ~50× smaller storage than LoGra while matching or exceeding gradient baselines on fidelity (Linear Datamodeling Score), data‐selection utility, and data-classification accuracy. A single 0.5 B-model–trained checkpoint generalises to 7 B Qwen models and even to LLaMA-1 B without retraining, amortising the one-off training cost after roughly 5 × 10⁵ attributed examples.", "questions": "None", "limitations": "None", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "## Strengths \n\n- **Principled combination of influence-function theory and Datamodel learning**. AirRep keeps the theoretically grounded bilinear structure of influence functions yet learns the coefficients from ground-truth re-training data, uniting fidelity with efficiency.\n\n- **Attention pooling derived from higher-order influence analysis**. The softmax weighting mirrors coefficients that arise in second- and third-order influence terms, yielding a principled and effective mechanism for group attribution that naïve summation misses.\n\n- **Joint encoder-and-pooler optimisation on ground-truth comparisons**. Auto-generated pairwise rankings align the representation directly with model behaviour, producing task- and model-aware embeddings with no manual labels.\n\n- **Downstream applications/utility**. AirRep  enables data selection (retaining accuracy while pruning the corpus) and data classification (retrieving semantically aligned training points).\n\n- **Thorough empirical validation**.  Experiments cover fidelity, data selection, data classification, ablations, scalability, and cost amortisation across four corpora and five LM sizes, against well-tuned gradient and representation baselines.\n\n- **Strong practical gains and robust generalisation** AirRep achieves ~80× speed and ~50× storage savings over LoGra, with a single checkpoint transferring unchanged to larger Qwen models and to LLaMA-1 B.\n\n## Weaknesses \n\n- **High initial supervision cost.** Building the O(1m) pairwise comparison labels requires training many mini-LMs and forward passes. This can get prohibitively expensive, especially when each run takes a long time to converge. \n\n- **Proxy-model mismatch.** Supervision is gathered on a 0.5 B Qwen model, yet AirRep is evaluated on targets up to 7 B (and on a different architecture). The paper does not quantify how fidelity degrades as the gap between proxy and target widens.\n\n- **Experimental-design limitations.** Baseline gradient methods (TracIn, LESS) are minimally tuned (single checkpoint / LoRA rank), and evaluation stays within instruction-tuning. Stronger baselines and broader domains (code, vision-language, RL) are needed to establish generality.\n\n- **Better iso-FLOP baselines**. To make the cost comparison fully transparent, can maybe include a LoGra (or similar) run whose total compute budget—gradients, Hessian pre-conditioning, plus any pairwise scoring—is matched to AirRep’s cost for generating comparison labels and training the encoder + pooler. This feels like a more apples-to-apples comparison to me.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback! We appreciate that you find our method principled and effective, our experiments thorough, and the performance is strong. We would like to respond to your comments as follows:\n\n---\n\n> ### **High initial supervision cost:**\n\nThank you for raising this concern. In Section 9, we provide a quantitative analysis of AirRep’s initial training cost and demonstrate how it can be amortized during inference. While the upfront cost of training AirRep is indeed higher than that of baseline methods, this is **amortized after approximately 5.6 GPU hours—beyond which AirRep yields an 80× speedup in inference**. Given this substantial gain in downstream efficiency, we believe the one-time offline training cost is a worthwhile investment and does not pose a significant limitation in practical settings where such training time is likely affordable. \n\n---\n\n> ### **Proxy-model mismatch:**\n\nThank you for this insightful point. We agree that proxy-target mismatch may lead to fidelity degradation as the model gap widens. To investigate this issue further, we conducted additional evaluations across a range of target models with varying sizes and architectures. The results on FLAN are summarized below: \n\nMethod | Qwen-0.5B | Qwne-7B | Llama-1B | Qwen3-0.6B | Qwen3-0.6B (Thinking) | TinyLlama-1B | GPT-2 |\n| - | - | - | - | - | - | - | - |\nLoGra (Qwen-0.5B) | 19.75 | 7.96 | 16.79 | 11.42 | 17.20 | 18.60 | 9.85 |\nLoGra (Target-LM) | 19.75 | 8.75 | 19.95 | 19.54 | 21.85 | 20.14 | 17.78 |\nAirRep (Qwen-0.5B) | 21.11 | 9.18 | 22.54 | 21.62 | 23.26 | 21.48 | 18.16 |\n\nThese results show that AirRep consistently outperforms LoGra, even when supervision is derived from a much smaller proxy model (Qwen-0.5B), suggesting strong generalization across both model scale and architecture. While **degradation is observed as the proxy-target gap increases (e.g., Qwen-0.5B → GPT-2), AirRep maintains robust relative gains compared to LoGra**. Due to limited time for the rebuttal, we will include a more detailed discussion and evaluation on this point in the final version.\n\n---\n\n> ### **Minimal tuning of baselines:**\n\nThank you for the thoughtful critique. Regarding baseline tuning, as noted in Appendix E, we carefully optimize both performance and efficiency for LoGra, TracIn and LESS. This included exploring different configurations and applying engineering optimizations on our evaluation data.\n\n- For **TracIn**, we report results using the final checkpoint rather than ensembles, as ensembling across multiple checkpoints linearly increases computation and storage overhead. In our preliminary experiments, we found that ensembling yields modest gains compared to simply increasing the projection dimension, while incurring higher resource costs and latency. We will elaborate on this trade-off more fully in the revised version of the paper.\n\n- For **LESS**, we used the LoRA rank recommended by the original paper. This setting already results in an embedding of dimensionality 8192—over 21× larger than that used by AirRep—making it a resource-intensive baseline.\n\n---\n\n> ### **Better iso-FLOP baselines:**\n\nThank you for this valuable suggestion. In **Section 9**, we currently report a full breakdown of the total computational cost—including data generation, training, and inference—for both AirRep and LoGra. Our analysis shows that AirRep becomes cost-effective after approximately 5.6 GPU hours, beyond which it offers significant inference-time speedups (up to 80× faster).\nWe note that this comparison represents an upper bound on AirRep’s training cost, as the learned representations can often be reused across different target models or tasks without retraining. This enables cost amortization in multi-task or multi-model scenarios, improving overall efficiency in practical deployments.\n\nFollowing your suggestion, we will revise the paper to include a clearer iso-FLOP-style analysis by aligning LoGra’s total budget (including gradient computations, Hessian preconditioning, and pairwise scoring) with that of AirRep’s full training pipeline.\n\n---\n\n**Thank you again for your thoughtful feedback—we hope these clarifications and forthcoming additions will resolve your concerns.**"}, {"review_id": "lRQeXRYYmK", "reviewer": "Reviewer_wX2D", "summary": "This paper introduces a new approach to training data attribution called AirRep. AirRep combines the advantages of gradient-based approaches with the efficiency of representtation-based methods. It encodes the training data with embeddings and scores their influence using inner-product with the provided example using attention based pooling. AirRep’s attention-based pooling mechanism captures group effects. The model is optimized using a weighted pairwise ranking objective over automatically generated influence signals since the expected usage is not pure influence scores but relative ordering of training examples. Evaluation on LLM fine-tuning demonstrates the effectiveness and generalization of AirRep. Authors also note that there is increased cost of training the embedding model but it can be done in advance.", "questions": "- Why some baselines (such as TracIn) are not evaluated in some sections?\n- Why wasn't TRAK evaluated as a baseline?\n- Would certain set of encoding models work better than others?\n\nI am voting borderline accept for now, but if the question about baselines are addressed I am willing to increase my score and vice versa.", "limitations": "One additional point could have been discussed as to whether this is NLP-specific TDA method or applicable to other general domains. \n\nOther than that, the limitations has been adequately addressed.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Quality\nThis paper has the following strengths:\n* It adds a clever attention-based pooling to represent input groups much better. It is also nice of authors to apply the same logic to some of the baselines to measure the impact.\n* Applying sigmoid-based training data confidence gating is also a good idea.\n* The evaluations show good improvement over baselines.\n* Its baselines (will touch on this later) are reasonable.\n* Good sections on Ablation Study and Amortizing Training Cost.\n\nThe paper could have been improved in these areas:\n* While it is understandable that the authors can't possibly evaluate their approach with all the metrics, this reviewer wonders why the direction of evaluation has been done in reverse (adding influential training examples rather than removing training examples) compared to other papers.\n* A good evaluation could have been performed with non-random subsets. An evaluation on fact-tracing, which could be a very popular use case of this work, could go a long way.\n* One of the more fundamental issues is lack of some important works such as TRAK for comparison.\n* It is not clear why each experimental section contains only a subset of the mentioned baselines. If there is a good reason for dropping an approach from evaluation, it should be explicitly mentioned. \n* nit: While it has been used before by other papers, correlation with another metrics is still not fully accepted form of evaluation.\n\n### Clarity\nThe paper is well written in general. A few comments:\n* Introduction and Preliminaries are too long. \n* LDS is introduced in [12] and not in [1].\n* Section 5 is called Data Attribution, but it is measuring more of counterfactual correctness.\n* very nitpicky comment: \"We define the model’s “prediction” on x as the cross-entropy loss computed on x, given by $r(x, S)$ = $ℓ(x; θ_*)$\" -- defining prediction as this causes a mental load for the rest of the paper for readability as one is too used to assume predictions are regular model outputs.\n\n### Significance\nThis paper adds to the plethora of training data attribution methods. Based on the improvement on the baselines (and lack of some), and added computational cost, this paper will be useful for some.\n\n### Originality\nThe paper has few interesting additions on TDA: weighted attention for group selection, confidence-based example selection and ranking-based training objective. None of these approaches are novel in isolation, but this paper demonstrates their clever application.", "was_revised": true, "final_justification": "This paper provides a new, solid approach to training data attribution-- AirRep. I think this paper makes a good addition to interpretability methods. Authors have addressed nearly all my comments leaving some for future work.", "rebuttal": "Thank you for your thoughtful and constructive review—your feedback highlights the strengths of our method (e.g., attention-based pooling, confidence gating, and ablation clarity) while also providing valuable suggestions on evaluation direction, baseline completeness (e.g., TRAK and TracIn), and broader applicability, which we will carefully address to improve the paper. We respond to each of your points in detail below. \n\n---\n\n>  **why the direction of evaluation has been done in reverse (adding influential training examples rather than removing training examples) compared to other papers.**\n\nThank you for this insightful observation. We chose to focus on **adding** influential examples—rather than removing them—primarily due to computational efficiency. In removal-based evaluation (e.g., Leave-One-Out LOO), each example must be removed from the training set and the model’s performance must be re-evaluated, often requiring a full retraining per removal. In contrast, additive evaluation can be performed efficiently by training on a selected subset, which is far more efficient in practice.\n\nHowever, following your suggestion, we conducted a **LOO evaluation**, in which individual examples were removed and we measured the correlation between predicted and actual influence scores. The results on a FLAN subset are summarized below:\n\n|  Method | LOO Correlation |\n| - | - |\n| LoGra (18432) |  10.75  |\n| GTE-Small | 0.72  | \n| AirRep | 15.36 |\n\nWe find that **AirRep outperforms baseline** methods in this setting as well. We will include this new analysis and discussion in the revised version of the paper.\n\n**For context,** our evaluation setup otherwise follows established precedents: for LDS, we follow [Datamodels] to evaluate on randomly sampled subsets (see Figure 6 for different subset size), and for data selection and classification tasks, our setup mirrors that of LESS and Hanawa et al.\n\n---\n\n>  **A good evaluation could have been performed with non-random subsets. An evaluation on fact-tracing, which could be a very popular use case of this work, could go a long way.**\n\nThank you for this thoughtful suggestion. We agree that fact tracing is a compelling and high-impact application of training data attribution. While time constraints during the rebuttal period may prevented us from conducting additional evaluations on fact-tracing dataset, we see this as a promising future direction. We will add a discussion of fact-tracing applications in the future work section of the revised paper and plan to explore corresponding evaluations in follow-up work. \n\n---\n\n>  **One of the more fundamental issues is lack of some important works such as TRAK for comparison. Why wasn't TRAK evaluated as a baseline?**\n\nThank you for raising this important point. In our experiments, we include **DsDm** [REF], an improved variant of **TRAK** specifically optimized for LLM-scale settings. In our evaluation, we found that DsDm performs comparably to LoGra and is consistently outperformed by AirRep across tasks. We will clarify the relationship between TRAK and DsDm in the revised version of the paper to avoid confusion and ensure proper attribution.\n\n[REF] DsDm: Model-Aware Dataset Selection with Datamodels\n\n---\n\n>  **It is not clear why each experimental section contains only a subset of the mentioned baselines. If there is a good reason for dropping an approach from evaluation, it should be explicitly mentioned. Why some baselines (such as TracIn) are not evaluated in some sections?**\n\nThank you for pointing this out. In Table 1, we present a comprehensive comparison across all baselines. In subsequent experimental sections, we focus on a subset of baselines—specifically the most competitive or directly relevant ones (e.g., LoGra, LESS, and GTE-small)—to ensure clarity and conserve space in the main text. Baselines such as TracIn, which consistently underperform relative to LoGra, were omitted from detailed comparisons for conciseness.\n\nThat said, we understand that this selective presentation may have caused confusion. To address this, we will (1) include **complete evaluation** results for all baselines in the Appendix, and (2) **explicitly state our rationale** for baseline selection in each experimental section in the revised version of the paper.\n\n---\n\n>  **The paper is well written in general. A few comments …**\n\nThank you for your detailed and constructive suggestions for improving clarity. In response, we will streamline the Introduction and Preliminaries sections to improve readability and focus. We will also revise the title of Section 5 to more accurately reflect its content, explicitly distinguishing between data attribution and counterfactual correctness. Additionally, we will clarify our use of the term “prediction” to avoid confusion with its conventional meaning in model outputs. \n\n---\n\n>  **Would certain set of encoding models work better than others?**\n\nThank you for this interesting question. In our preliminary experiments, we compared **GTE-small** and **DeBERTa-small**, and observed that GTE-small achieved slightly better performance—likely due to its pre-training on text retrieval tasks, which aligns well with our pairwise comparison objective.\n\nWhile GTE-small provided strong overall results, we agree that broader exploration of encoding models is worthwhile. In particular, models like **BERTScore** or **ColBERT**, which support fine-grained token-to-token matching, may offer further improvements in attribution fidelity. We view this as a promising direction and plan to explore it in future work.\n\n---\n\n>  **One additional point could have been discussed as to whether this is NLP-specific TDA method or applicable to other general domains.**\n\nWhile our current evaluation focuses on language tasks, the design of AirRep is **modality-agnostic**—it relies only on the availability of suitable encoders and the ability to estimate loss differences. In principle, this framework could be applied to other domains such as vision or multimodal settings.\n\nWe will include this point in the discussion of limitations and future work in the revised paper and plan to explore cross-domain generalization in follow-up research.\n\n---\n\n**Thank you again for your thoughtful feedback—we hope these clarifications will resolve your concerns.**"}, {"review_id": "RNKJvwMeF4", "reviewer": "Reviewer_zyXJ", "summary": "The paper proposes an optimisation-based approach to computing training data attribution scores, where a model consisting of an encoding module and an aggregation module is trained to predict the attribution score of a subset of the training data on a specific model output. The model is trained with true attribution scores, which the authors define as the model's loss on a test sample where the model was finetuned on the subset.", "questions": "- How many subsets $r$ and therefore additional training cost is required to get a well-performing AirRep model? Does it depend on model and data size of the model one wishes to investigate with training data attribution? \n- How well does AirRep work for smaller/larger subset sizes than the ones reported in the paper? TDA methods are known to struggle to capture attribution for individual samples due to the training noise [a, b, c], does an optimization-based approach alleviate this?\n\n[a] Basu et al. (2020): Influence Functions in Deep Learning Are Fragile.\n\n[b] Epifano et al. (2023): Revisiting the fragility of influence functions. \n\n[c] Nguyen et al. (2023): A bayesian approach to analysing training data attribution in deep learning.", "limitations": "Yes, shortly in the conclusion section.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- To train a model for predicting attribution scores is an innovative approach to training data attribution that could scale well for larger models and datasets.\n- The evaluation strongly supports AirRep's performance and generalizability across different subsets.\n- The authors present a connection between higher-order influence functions and the attention-based pooling approach of AirRep.\n\nWeaknesses:\n- The paper would improve if the authors added a small paragraph or some sentences of how the related works relate to their work. \n- It is unclear whether training an AirRep model requires specific tuning - how easy will it be to implement this?", "was_revised": true, "final_justification": "I gave the score of accept because this paper presents a promising alternative, scalable approach to training data attribution and presents a convincing evaluation. The authors were responsive and able to answer questions and engage in discussions, demonstrating their fluency in the field. They also provided additional results comparing with a ground truth recorded through counterfactual LOO experiments, showing that their optimization-based approach to TDA is valid.", "rebuttal": "We sincerely thank the reviewer for the thoughtful and encouraging feedback. We appreciate the recognition of AirRep’s originality in formulating training data attribution as an optimization problem, its strong empirical performance and generalizability, and the connection to higher-order influence functions via attention-based pooling. We respond to each of your comments and suggestions in detail below. \n\n---\n\n> **It is unclear whether training an AirRep model requires specific tuning - how easy will it be to implement this?**\n\nThank you for this question. In our experience, **AirRep is relatively easy to train and implement.** The training data is generated automatically by sampling random subsets from training corpus, without requiring any specialized design or manual curation. Training follows a standard contrastive learning setup and does not rely on extensive hyperparameter tuning.\n\nSpecifically, we used a fixed learning rate of 1e-4 and a batch size of 32, selected to fit typical GPU memory constraints. These settings yielded stable training across all experiments. We have also shared our data generation and training code in the Supplementary Material, which we hope will make implementation straightforward for future users.\n\n---\n\n> **How many subsets r and therefore additional training cost is required to get a well-performing AirRep model? Does it depend on model and data size of the model one wishes to investigate with training data attribution?**\n\nThank you for this important question. In **Section 9**, we provide a quantitative analysis of how training set size impacts AirRep’s performance. In our experiments, training with **several hundred random subsets** over approximately 5 GPU hours was sufficient for AirRep to outperform LoGra across multiple tasks.\n\nWe found that the required number of subsets depends more on the diversity of the data distribution than on the model or dataset size. For example, more diverse datasets like Tulu typically benefit from more training subsets, whereas less diverse datasets such as Alpaca require fewer. Notably, we did **not observe saturation** in performance when increasing the number of training subsets in our current experiments. This suggests that the training cost can be flexibly traded off against performance and inference efficiency, depending on available resources.\n\n---\n\n> **How well does AirRep work for smaller/larger subset sizes than the ones reported in the paper?**\n\nThank you for this thoughtful question. In **Figure 6**, we report results on the FLAN dataset across a range of subset sizes. Although **AirRep was trained only on subsets of size 1000**, it consistently outperforms LoGra across both smaller and larger subset sizes at test time. AirRep’s advantage becomes more pronounced with larger subsets, likely due to its attention-based pooling mechanism, which is designed to explicitly model group-level influence.\n\n---\n\n> **TDA methods are known to struggle to capture attribution for individual samples due to the training noise [a, b, c], does an optimization-based approach alleviate this?**\n\nThank you for this insightful question. Capturing attribution at the individual-sample level is indeed challenging. Since AirRep is trained on many sampled subsets, this may have an **averaging effect** similar to repeated training, helping to stabilize the outcomes. Additionally, our *importance reweighting objective* is designed to reduce the influence of training noise.\n\nIn addition, to further assess AirRep's performance on single sample attribution, we evaluated it using a Leave-One-Out (LOO) metric. where we remove a single training example and measure the actual change in test loss, then compare it to the predicted influence score. The results on a FLAN subset are summarized below:\n\n| Method | LOO Correlation |\n| - | - |\n| LoGra (18432) |  10.75  |\n| GTE-Small | 0.72  | \n| AirRep | 15.36 |\n\nAirRep outperforms the baselines in this setting, suggesting that its group-level optimization framework can yield stable attribution at the individual-example level. We will include a more detailed discussion of this analysis in the revised paper.\n\n---\n\n> **Typos …**\n\nThank you for pointing out the typos!  We will correct these in the revised version of the paper.\n\n--- \n\n**Thank you again for your thoughtful feedback—we hope these clarifications will resolve your concerns.**"}, {"review_id": "8xTzs8bQ90", "reviewer": "Reviewer_32FH", "summary": "This paper introduces AirRep, a training data attribution (TDA) method to explain how the training data contribute to the model prediction. Previous TDA methods, such as gradient-based methods Tracin and LoGra, have theoretical grounds, its high computational cost hinders scalability for LLMs. Existing representation-based TDA is more efficient but relies on heuristic design and the representations are not optimized for attribution, limiting accuracy. \nAirRep adopts a trainable encoder and attention-based pooling to learn task-specific and model-aligned representations explicitly for TDA, where the core idea is to accurately group influence estimation. Experiments demonstrate AirRep's performance is comparable to computationally expensive gradient methods.", "questions": "- AirRep uses BERT for the encoder. Why BERT?\n- What exactly is the \"re-training outcome\" in line 83\n- My understanding is that this AirRep approach learns task-specific and model-aligned representations explicitly optimised for TDA. Therefore, it is not efficient that for each downstream task, we need to train again? Although previous TDAs are also task-dependent, but AirRep is quite expensive to train compared to representation-base TDA. Is it correct?\n- This paper only considers fine-tuning, where he knowledge bias present in the base model before fine-tuning is ignored, Is this understanding correct?\n- Is it possible to use KSD [3] as an evaluation metric?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\n- Research on TDA is an important step towards understanding how LLMs learn from their training data, which is fundamental to ensure transparency and accountability in AI systems.\n- AirRep achieves performance on par with the stoa gradient-based approaches. Specifically, it outperforms LoGra in the LDS evaluation.\n\n**Weaknesses**\n\n- The paper only uses the LDS evaluation as the sole metric for comparison. Other metrics like Leave-One-Out or AOPC value are used in previous work [1]. \n- The weighted pairwise ranking loss uses thresholds to clip weights based on the difference in ground-truth scores. The specific values used for these thresholds in the experiments are not explicitly stated in the provided text. Seems like there is some heuristic engineering work. Not sure if the thresholds are generalizable to other models or datasets.\n\n**Others**\n\n- There are loads of papers cited in their arxiv versions, not their published versions. See example below:\n\nMengzhou Xia, Mikel Artetxe, Chunting Zhou, Xi Victoria Lin, Ramakanth Pasunuru, Danqi Chen, Luke Zettlemoyer, and Veselin Stoyanov. 2023. Training Trajectories of Language Models Across Scales. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 13711–13738, Toronto, Canada. Association for Computational Linguistics.\n\nBasu, Samyadeep, Phil Pope, and Soheil Feizi. \"Influence Functions in Deep Learning Are Fragile.\" In International Conference on Learning Representations.\n\n- The paper would benefit from presenting a qualitative analysis, accompanied by concrete examples.\n\n**Missing related work**\n\n[1] NIPS2023, Nguyen, Elisa, Minjoon Seo, and Seong Joon Oh. A bayesian approach to analysing training data attribution in deep learning.\n\n[2] ACL2024, Kangxi Wu, Liang Pang, Huawei Shen, and Xueqi Cheng. 2024. Enhancing Training Data Attribution for Large Language Models with Fitting Error Consideration.\n\n[3] ICLR2025, Mahtab Sarvmaili, Hassan Sajjad, Ga Wu, Data-centric Prediction Explanation via Kernelized Stein Discrepancy", "was_revised": true, "final_justification": "Most of my concerns and misunderstandings have been addressed. Based on the comments and other reviewers' comments, I would like to increase my score.", "rebuttal": "We thank the reviewer for engaging with our work and for recognizing the importance of TDA for model transparency, as well as AirRep’s strong performance. We will address each of your points in detail below, with particular focus on clarifying our contributions to AirRep’s novelty, scalability, and performance as a principled and robust framework for TDA. \n\n---\n\n> **The paper only uses the LDS evaluation as the sole metric for comparison. Other metrics like Leave-One-Out or AOPC value are used in previous work [1].**\n\nThank you for pointing this out! While LDS has been a widely used benchmark in recent TDA research, we agree that evaluating attribution quality with alternative metrics adds valuable perspective. To this end, we conducted a Leave-One-Out (**LOO**) analysis, where we removed individual training examples and measured the Pearson correlation between the predicted influence scores and the actual change in test loss.\nThe results on a FLAN subset are summarized below:\n\n| Method  | LOO Correlation |\n| - | - |\n| LoGra (18432) |  10.75  |\n| GTE-Small | 0.72  | \n| AirRep | 15.36 |\n\nAirRep achieves higher correlation than baseline methods, suggesting that its optimization-based approach provides more robust and reliable individual-level attribution.  We will include this analysis and discussion in the revised version of the paper. \n\n---\n\n> **The weighted pairwise ranking loss uses thresholds to clip weights based on the difference in ground-truth scores. The specific values used for these thresholds in the experiments are not explicitly stated in the provided text. Seems like there is some heuristic engineering work. Not sure if the thresholds are generalizable to other models or datasets.**\n\nThank you for raising this concern. The clipping thresholds were set to 0.1 (low) and 5.0 (high), as described in the AirRep training details in Section 4. These values were not extensively tuned and were selected based on simple, generalizable observations:\n- The **low threshold (0.1)** corresponds to approximately one standard deviations of the loss difference observed from repeated fine-tuning of language models on the same data, which we treat as typical variation due to training stochasticity.\n- The **high threshold (5.0)** was empirically chosen to cap rare but extreme differences, which we consider outliers and clip to prevent them from disproportionately influencing the learning process.\n\nImportantly, we applied the **same thresholds and training hyperparameters across all experiments**, and observed consistently stable and performant training. We will clarify these points in the revised version to ensure transparency and reproducibility.\n\n---\n\n> **There are loads of papers cited in their arxiv versions, not their published versions.**\n\nThank you for pointing out this issue! We will correct these citations in our final version.\n\n---\n\n> **The paper would benefit from presenting a qualitative analysis, accompanied by concrete examples.**\n\nThank you for the suggestion. We agree that including qualitative examples can provide deeper intuition into how AirRep assigns attribution scores. In the revised version, we will include detailed case studies and concrete examples to illustrate attribution behavior. Additionally, we will open-source our data and models to support transparency and enable further exploration by the community. \n\n---\n\n> **AirRep uses BERT for the encoder. Why BERT?**\n\nThank you for this question. We initialize AirRep with a BERT-style encoder (specifically, GTE-small) due to its strong performance in text embedding tasks and ease of training. GTE-small is well-suited for pairwise scoring due to its retrieval-focused pretraining. That said, AirRep is not tied to any specific encoder architecture. We expect that other text embedding models could also be integrated successfully.\n\n---\n\n> **What exactly is the \"re-training outcome\" in line 83**\n\nThe “re-training outcome” mentioned refers to r(x, S), defined in line 79 as the loss of the model on example x after training on dataset S. We recognize that this phrasing could cause confusion and will clarify this in our final version.\n\n---\n\n> **My understanding is that this AirRep approach learns task-specific and model-aligned representations explicitly optimised for TDA. Therefore, it is not efficient that for each downstream task, we need to train again? Although previous TDAs are also task-dependent, but AirRep is quite expensive to train compared to representation-base TDA. Is it correct?**\n\nIn our evaluation, we did not retrain AirRep for each downstream task. Once trained, AirRep learns model-aligned representations for predicting influence scores via loss change estimation and can be applied to various downstream TDA tasks, such as data selection and classification. While we observe promising generalization to unseen data, we acknowledge that performance may vary in domains with substantially different distributions—such as complex reasoning tasks where long COT models generate significantly longer thinking content. Exploring domain-adaptive training strategies is a promising direction for future work.\n\nRegarding training cost, **Section 9** provides a detailed comparison with baseline TDA methods. We show that although AirRep incurs a one-time training cost, it becomes cost-effective after approximately 5.6 GPU hours, and achieves up to 80× faster inference than gradient-based approaches. Since the trained model can be reused across tasks, this cost can be amortized in multi-task settings, offering both principled modeling and practical scalability.\n\n---\n\n> **This paper only considers fine-tuning, where he knowledge bias present in the base model before fine-tuning is ignored, Is this understanding correct?**\n\nWhile our evaluation focuses on fine-tuned models, **we do not ignore the influence of the pretrained model's biases**. AirRep’s supervision signal is derived from **actual training outcomes**, meaning that any biases inherited from the base model and expressed during fine-tuning are naturally reflected in the measured loss changes. Thus, our attribution framework implicitly captures the combined effects of both pretraining and fine-tuning.\n\n---\n\n> **Is it possible to use KSD [3] as an evaluation metric?**\n\nThank you for highlighting this interesting work. After reviewing [3], we agree that KSD presents a promising alternative for evaluating attribution methods. However, adapting it to our current LLM training setup would require non-trivial adjustments, particularly since the original implementation was designed for image classification.\n\nThat said, because [3] also evaluates methods like Influence Functions and TracIn, we believe KSD could potentially be extended to evaluate AirRep as well. We consider this an exciting direction for future work and will explore it in follow-up studies. \n\n---\n\n**Thank you again for your thoughtful feedback—we hope these clarifications will resolve your concerns.**"}]

Author rebuttal:

Meta-review: This submission proposes a method for training data attribution (TDA) called AirRep, where the core idea is to train a representation encoding model for the task of ranking/comparing training examples according to their influence scores.

Strengths identified by the reviewers are as follows:
1. Learning representations for TDA is an innovative idea and "unites fidelity with efficiency" (quoting Reviewer VLz2).
1. AirRep includes an attention-based pooling module to account for the effects of groups of training examples. Reviewer wX2D appreciated that this contribution was separately evaluated in combination with some of the baseline TDA methods.
1. Large efficiency gains are achieved at inference time.
1. A single AirRep checkpoint trained on a Qwen-0.5B model transfers to other LLMs such as Qwen-7B and LLaMa-1B.
1. Thorough empirical evaluation, covering TDA fidelity (measured by Linear Datamodelling Score, LDS), data selection, data classification, an ablation study, and AirRep's supervision and training cost.
1. The pairwise comparison data used to train AirRep is generated automatically.

Weaknesses:
1. AirRep has a high initial cost for generating training data and performing the training. In their rebuttal, the authors emphasized the results in Section 9, which show for example that AirRep requires $\sim 10^5$ training examples to outperform LoGra and amortizes its training cost after 5.6 GPU hours, after which it can process substantially more TDA queries than LoGra. Still, I think that this indicates a requirement for high-volume TDA applications for AirRep to be advantageous.
1. Several reviewers pointed out that leave-one-out (LOO) evaluation (predicting the effect of leaving out one training example at a time) was missing. The author rebuttal reported on a LOO evaluation for a subset of the experiments. This LOO evaluation could be made more complete.
1. Unclear reasons for including only a subset of the baselines in each experimental section (as well as not including TRAK). The rebuttal promised to clarify and include complete comparisons in the appendix.
1. Limitation of the current work to the modality of language could be acknowledged.

The reviewers reached a consensus that the submission should be accepted, with three reviewers rating the work as a clear Accept based on the many strengths listed above. In my view, the weaknesses are either inherent to AirRep and well-discussed in the paper, or very reasonable to address for the camera-ready version. 

I am recommending a spotlight presentation because of strengths 1, 4, 5: The innovativeness of learning representations for TDA, coupled with the thorough evaluation including an analysis of when AirRep's training cost (its main weakness) is amortized, and the transferability between models further mitigating the cost.
Decision: Accept (spotlight)