Title: Enhancing Training Data Attribution with Representational Optimization
Venue: NeurIPS 2025
Accepted: yes
Keywords: Training Data Attribution
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

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 .

1 Introduction

The remarkable success of large language models (LLMs) has been demonstrated across a wide range of tasks. However, a 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-on specific inputs or 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, is an important step to answering this question [1]. TDA is crucial for ensuring transparency and accountability in AI systems by revealing how data influences model outputs.

Existing approaches to TDA can broadly be categorized as either gradient-based or representationbased . Gradient-based approaches, rooted in influence estimation [2, 3], aim to quantify the impact of individual training points on model predictions. Their core idea is to use gradients and the inverse Hessian of the loss function to make a first-order approximation of how a model's predictions would change if a certain example was removed from the training set [3]. While these methods are theoretically well-motivated, they are typically computationally expensive and rely on assumptions of loss convexity and model optimality - both of which are violated in modern neural networks [4, 5].

In contrast, representation-based methods estimate influence based on similarity in the representation space. The core assumption is that training examples whose representations are close to a test input are likely to have influenced its prediction [6]. Recent works have explored various notions of similarity based on 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

30

Gradient

Representation

(Examples/Second)

3000-

Gradient

Representation

26.2

2500

25002500

2500-

Score

25

20

18.4

2000-

16.2

15

13.8

1500-

11.3

10

10.0

7.7

Speed (

1000-

·..................

5-

3.9

500-

278

31.2

5.3

30.2

29.5

-0

0-

T

LESS

RDS

LESS

RDS

Tracln

TF-IDF

AirRep

Tracln

TF-IDF

AirRep

LoGra (3x)

GTE-Small

LoGra (3x)

GTE-Small

LoGra (48x)

LoGra (48x)

scalable, making them well-suited for large-scale applications, such as curating pre-training data for LLMs [8, 11]. However, the quality of a representation-based method depends heavily on the selected feature space. This problem is compounded by the fact that existing work utilizes heuristically designed representations that are not tailored to the particular target task or model [12].

Finally, a persistent challenge in both method families is is how these methods are extended to estimate the collective influence of a group of training examples. Most past work extends singleexample 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 a natural question: Can we develop a method that provides the best of both worlds?

We answer this question by introducing the Attentive Influence Ranking Representation (AirRep), a new representation-based approach that incorporates two novel mechanisms to improve performance over prior representation-based methods. First it uses a trainable encoder that produces task- and model-specific text representations, ensuring that the feature space is optimized for data attribution estimation. Second, it introduces an attention-based pooling mechanism that effectively aggregates multiple examples into a 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 a 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 UltraChat [15], a large-scale dialogue generation dataset, and evaluate the model on five unseen instruction-tuning test set (FLAN, Alapca, Tulu, SafeRLHF) that do not appear in the training data. Figure 1 shows the main evaluation results. Notably, in the standard Linear Datamodeling Score (LDS) evaluation [1], AirRep significantly outperforms existing gradient-based approach [16], despite being ∼ 80 × more computationally efficient and ∼ 50 × 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 Preliminaries

We consider a framework that is applicable to different applications of TDA but focuses on the training of language models (LMs) as a concrete example. Let S = ( z 1 , z 2 , . . . , z n ) represent a training set of n data points, where z i is a training example that includes both an input and an output. In the context of language models, we refer to the input as a 'prompt' and the output as a response. Furthermore, let x denote a test example. The goal of TDA is to understand how the training examples in S contribute to the model's prediction on x .

Formally, let θ represent the parameters of the LM. The model is fine-tuned on S , aiming to find

where ℓ ( · ) denotes the cross-entropy token prediction loss. As language model training is non-convex, we use θ ∗ to denote the optimized parameters after fine-tuning, regardless of whether a true minimum has been found. We define the model's ' prediction outcome ' on x as the cross-entropy loss computed on x , given by r ( x, S ) = ℓ ( x ; θ ∗ ) , i.e. the evaluation loss on the test example x for the model fine-tuned on the dataset S . Here, we use the cross-entropy loss instead of task-specific metrics like the accuracy, F1 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].

A data attribution model, denoted by f ( x, S ) , aims to estimate the actual retraining outcome r ( x, S ) -i.e., the loss of the model on example x after training on dataset S . Existing methods can be generally categorized into two groups: gradient-based methods and representation-based methods .

Gradient-based methods Gradient-based methods largely stem from influence functions [2, 3], which approximate how θ would change in response to infinitesimal perturbations in the weighting of training instances. Specifically, influence functions quantify the influence from a single training example z i to an evaluation example x with a closed form expression:

where H -1 is the inverse Hessian of the training loss wrt model parameters. The derivation of Eq. 2 is given by first-order Taylor approximation of ℓ ( x ; θ ) around θ ∗ [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 Hessian [12, 21]. Recent techniques have also considered techniques like gradient projection [16] and model ensembling [22, 23]. As a representative approach in LLM application, LoGra [16] define the influence of a training example z i on x is computed as:

where ϕ ( z ) is the projected, Hessian-corrected, and unit-normalized gradient for the input example z , given model parameters θ :

where norm [ z ] 2 = z ∥ z ∥ 2 is a unit normalization operation [24]; ˆ θ denotes the trainable weights of a LoGra module [16] with PCA initialization for gradient projection; ∇ ˆ θ ℓ ( z ; θ ) is the gradient of the loss with respect to ˆ θ ; and H ˆ θ -1 2 is computed on ˆ θ and is approximated based on the KroneckerFactored Approximate Curvature (KFAC) algorithm [21]. Further, under first-order approximation, the group influence f GD ( x, 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 z i to x can then be computed as:

where Enc( z ) denotes a continuous embedding of input z encoded by encoder model Enc . Group influence is usually defined as summation of individual influence under a 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 representations [8], TF-IDF representations [25, 26], and text embeddings [11, 27]. See Appendix C for further comparison of the two types of TDA methods.

x

Encoder

S1

Encoder

Pooler

f(x,S1) = 0.5

Attention

S2

Encoder

Pooler

f(x,S2) = 0.2

TrainingSubsets

IndividualEmbeddings

GroupEmbeddings

3 AirRep

To bridge the gap between gradient- and representation-based data attribution, we introduce AirRep, a novel data-driven approach to enhance representational influence models. At a high level, AirRep consists of a trainable encoder Enc for embedding input examples and a pooling layer Agg that aggregates the representation of individual examples to produce a group representation. The influence score between a training set S and a target example x is then computed as

Specifically, AirRep incorporates two novel mechanisms to improve performance over standard representation-based methods: First, we introduce an effective attention-based pooling mechanism that better reflects group influence effect (Section 3.1). Second, we optimize the model on autogenerated data so that AirRep's score better reflects the underlying models predictions (Section 3.2).

3.1 Attention-based Pooling

Recall that data attribution methods typically include a mechanism for aggregating the contributions of elements of z i ∈ S when performing group influence. For this purpose, AirRep uses a pooling layer to aggregate the representations Enc( z i ) into a single embedding. The design of this pooling layer is crucial for modeling the relationships between group of data samples.

In AirRep, we propose a simple yet effective pooling techniques: attention-based influence pooling. Specifically, AirRep's aggregation method incorporates an attention-like operation to capture interactions between elements of x and S :

where Enc( x ) denotes the sentence embedding predicted by a BERT-based sentence encoder [28, 29], α i denotes the attention score of example z i , and | · | denotes the absolute value operation.

Remark Empirically, we 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, we show that high-order group influence introduces sample-wise weights, akin to our attention-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 x and S and output a scalar score f ( x, S ) . Ultimately, the objective of data attribution is that f ( x, S )

is correlated to r ( x, S ) , i.e., the actual prediction on x of model trained on S . Since AirRep is overall trainable, we therefore formulate a training objective that aims to make f ( x, S ) reflect r ( x, S ) . Given a ground-truth collection of r ( x, S ) scores for different datasets S , we could in principle train f ( x, S ) to match these scores as a regression problem. However, in data attribution we care less about matching the exact values of r ( x, S ) 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 on this example? '. We therefore formulate training as a pairwise ranking problem - that is, given two different data subsets S 1 and S 2 , whether the score of f ( x, S 1 ) and f ( x, S 2 ) reflect the true data preference r ( x, S 1 ) and r ( x, S 2 ) . In this section, we will first introduce the data generation pipeline we used for collecting ground-truth r ( x, 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 cross-validation-style data and compute ground-truth attribution scores. Specifically, we first sample N v examples as a validation split and N t examples as a training split. Then, from the training split, we randomly sample M subsets of data with replacement, denoted as S = { S 1 , S 2 , . . . , S M } , where each subset S i contains n training examples. For each training subset S i , we finetune a language model on it as in Eq. 1 to obtain a model checkpoint θ i . We then evaluate the trained model on each example x in the validation subset to calculate the loss: ℓ ( x ; θ i ) . After training M models and calculating the loss, we compute the negative normalized loss for each x in the validation set as

where the original loss ℓ ( x ; θ i ) is normalized by the mean and variance across models to stabilize the training signal, ensuring that the distribution of scores for each target x is on a similar scale.

Training Objective Given a target example x from our validation set, M training subsets S = { S 1 , S 2 , . . . , S M } , and the corresponding normalized loss values { ˆ r ( x, S i ); i ∈ [ M ] } , AirRep estimates the normalized loss values for each subset as f ( x, S i ) . For brevity, we use r i and f i to represent these scores, respectively. Our proposed training objective is a pairwise ranking loss that encourages the difference between the model-predicted scores, f i -f j , to better align with the difference in precomputed scores, r i -r j .

In practice, the label r i is noisy due to 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:

σ is sigmoid function, w i,j is a weighting function that depends on the absolute difference in ground truth scores, | r i -r j | , and is clipped using thresholds T min and T max to avoid training on i, j pairs with incorrect ordering and to 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 a batch size of 32 and the AdamW optimizer [37] with a learning rate of 2e-5 for two epochs.

Data Generation For AirRep training, we utilize two datasets, each representing distinct scenarios: 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 a wide range of topics. To generate training signal, we set N v = 10 4 and N t = 10 5 . The training subsets number is M = 100 , with each subsets containing n = 1 , 000 samples. We construct 100 cross-validation instances. Thus, in total, the data includes

10 4 unique training subsets and 10 7 training examples. The Qwen2.5-0.5B 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, a 30M parameter embedding model [35], and apply a 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-validation instance, then sample 1 , 000 examples from its validation subset and 32 training subsets. Distributed training is employed to maximize GPU memory utilization. The clipping thresholds, T min and T max, are set to 0 . 1 and 5 . 0 , respectively. The model is optimized for up to 2 , 000 steps using the AdamW optimizer with a learning rate of 1 × 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 instructiontuning dataset generated by OpenAI's text-davinci-003 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 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 methods. The gradient-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 in the Logix software 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 AdamW states. (iv) Dsdm [33], an efficient implementation of Trak [12] for LLM tasks.

The representation-based 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], a 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.

5 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 ( S 1 , · · · , S 100 ) of its training set, each containing n training examples, and train the target LM on each of these subsets. For each example x in the test set, we approximate the expectation of the model prediction r ( x, S i ) as the evaluation cross-entropy loss. Given the score f ( x, S i ) predicted by the data attribution model, the LDS is computed as the Spearman rank correlation between the true and estimated influence scores: ρ ( { r ( x, S i ) : i ∈ [100] } , { f ( x, S i ) : i ∈ [100] } ) . Finally, the LDS on the dataset is averaged across all test examples.

AirRep(Ours)

TF-IDF

GTE

LoGra

LESS

20

.

16

(166x)

(262x)

9

Dim

(84x)

12

(262x)

215(85x)

Correlation

15

12

214(43x)

6

O

213(21x)

8

10

8

O

212(11x)

LDS

211(5x)

4

3

210(3×)

5

4

384(1x)

0

0

1

Qwen2.5-0.5B

2

0

0

1

Qwen2.5-1.5B

2

0

0

1

Qwen2.5-3B

2

3

0

0

1

Qwen2.5-7B

2

3

ComputationtimerelativetoGTE-Small(inLog10Scale)

Results Table 1 presents the LDS evaluation results for attributions computed on the Qwen2.5-0.5B model's predictions across four evaluation datasets. AirRep achieves the best performance across all evaluation datasets. Notably, it outperforms the competitive gradient-based approach LoGra [16], despite being 80 × more computationally efficient and 50 × more storage efficient. Additionally, previous representation-based methods exhibit low correlation scores in this evaluation, underscoring the importance of task-aware optimization of representations.

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.e., a 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. In all cases, we use the same AirRep model trained exclusively on data generated by Qwen2.5-0.5B. Notably, AirRep demonstrates strong performance when applied to larger target LMs. Specifically, it consistently achieves results comparable to the best gradient-based methods, which require significantly more resources (e.g., 84 × to 262 × more computation time and 50 × to 80 × 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).

6 Evaluation of Data Selection

Setup One use of data attribution is data selection, 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, we use each TDA method to select a high-value subset using a greedy selection strategy. First, given the scores predicted by TDA models, we rank all training examples based on their highest score with respect to the test set samples and retain the top 1,000 examples. We then train LMs on the selected subset and evaluate their performance on the test set. The evaluation metric is the F1 score between the LM-generated outputs and the ground-truth answers.

Full

DSIR

GTE

LoGra

Random

TF-IDF

LESS

AirRep

58

56

工

54

52

Score

50

F

48

46

44

42

Qwen-0.5B

Qwen-1.5B

Qwen-3B

Results Figure 4 presents the results of using six TDA methods for data selection, along with a random data selection baseline, evaluated in terms of F1 score and selection cost. Key findings

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-based methods (TF-IDF and GTE); however, they are computationally expensive-the time required to select a 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-based 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 a collection of 12 source datasets (e.g., ShareGPT, CoT, Code-Alpaca). We use the task type or source dataset as a label, expecting that training data from the same task or 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 to harmful generations at test time.

Results Table 2 presents the evaluation results of LoGra, GTE-Small, and AirRep. 1 The results indicate that AirRep significantly improves data classification accuracy on FLAN compared to GTE-Small (increasing 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 in an unsupervised manner. A similar trend is observed for Tulu, where AirRep demonstrates superior performance in identifying the data source. For SafeRLHF, AirRep outperforms LoGra but lags behind GTE-Small, likely because AirRep's training data, generated from UltraChat, does not contain harmful content and thus lacks the necessary learning signal.

8 Ablation Study

We conduct an ablation study to 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 a 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 improvement of 12.17-demonstrating the effectiveness of encoder optimization. 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 a better underlying representations.

Data Attribution(LDS)

DataClassification(Accuracy)

ProcessedDataSize

AvgAcc

AirRep

85

100M

40

LoGra（18432)

80

10M

30

75

LoGra (1152)

1M

LoGra

GTE-Small

O(5.6hours,475,588)

20

70

LoGra (18432)

100K

LoGra(1152)

65

10

Traeln

10K

60

10

10°

5

10°

6

10

10

4

10°

5

10°

6

10

0

4

8

12

16

20

24

AirReptrainingsize(logscale)

AirRep training size (log scale)

GPUhours

9 Amortizing Training Cost

Compared with gradient-based methods, AirRep is more efficient at inference but requires training on auto-generated data, which incurs additional cost. Our evaluation above shows that FLAN- or UltraChat-pretrained 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 (1152)-which has a comparable embedding size-when trained with around 10 5 examples. We refer to this AirRep model as the ' crossover checkpoint '.

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-achieves nearly two orders of magnitude more throughput once model is trained. Notably, a 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 5, AirRep can be transferred from a 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.

Amortizing Across Models We evaluate AirRep, trained with Qwen-0.5B-generated supervision, using different target LMs (Qwen-0.5B, Qwen-7B, Llama-1B, Qwen3-0.6B, Qwen30.6B (Thinking), TinyLlama-1B, GPT-2) on FLAN. Results in Figure 7 show that across all targets, AirRep consistently surpasses LoGra, even when LoGra computes gradients on the target LMs. We also observe that LoGra is sensitive to the proxy used for gradient calculation-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 AirRep and that its training cost can be amortized across different target language models.

18.2

GPT-2 -

17.8

9.8

21.5

TinyLlama-1B

20.1

18.6

Target LM

Qwen3-0.6B

(Thinking)

17.2

21.9

23.3

21.6

Qwen3-0.6B

19.5

11.4

22.5

Llama-1B

19.9

16.8

9.2

Qwen-7B

8.8

8.0

21.1

Qwen-0.5B

19.8

19.8

8

10

12

14

16

18

20

22

24

LDS

LoGra

LoGra

AirRep

Qwen-0.5B

Target-LM

Qwen-0.5B

10 Related Work

Most existing work in TDA focuses on gradient-based methods, such as Influence Functions [3] and their subsequent optimizations, including group influence estimation [41, 13], Hessian approximation [42, 21, 43], gradient normalization [24, 25], gradient projection [22, 16], model ensembling [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 in large neural networks [48, 4]. Additionally, gradient-based methods suffer from the computational expense of gradient calculations and Hessian inverses [21, 33].

Alternatively, representation-based 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-based methods such as datamodels [50-52, 1], 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 LLM generalization [21] and learning processes [30, 58]. Another notable line of related work is data selection for LLMs, 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 representationbased [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-65].

11 Conclusion

This paper introduces a new approach to training data attribution called AirRep, which combines the advantages of task-driven optimization in gradient-based approaches with the efficiency of representation-based methods. AirRep represents training data as embeddings and scores their influence using inner-product similarity in the representation space. AirRep's novel attentionbased pooling mechanism captures the group effects, and the entire model is 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-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. 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

Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclerc, and Aleksander Madry. Datamodels: Predicting predictions from training data. ArXiv , abs/2202.00622, 2022.

Frank Hampel. The influence curve and its role in robust estimation. Journal of the American Statistical Association , 69:383-393, 1974.

Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International Conference on Machine Learning , 2017.

Samyadeep Basu, Phillip E. Pope, and Soheil Feizi. Influence functions in deep learning are fragile. ArXiv , abs/2006.14651, 2020.

Andrea Schioppa, Katja Filippova, Ivan Titov, and Polina Zablotskaia. Theoretical and practical perspectives on what influence functions do. ArXiv , abs/2305.16971, 2023.

Kazuaki Hanawa, Sho Yokoi, Satoshi Hara, and Kentaro Inui. Evaluation of similarity-based explanations. In International Conference on Learning Representations , 2021.

Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 586-595, 2018.

Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. ArXiv , abs/2302.03169, 2023.

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.

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.

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.

Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry. Trak: Attributing model behavior at scale. In International Conference on Machine Learning , 2023.

Samyadeep Basu, Xuchen You, and Soheil Feizi. On second-order group influence functions for black-box predictions. In International Conference on Machine Learning , 2019.

Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V. Le. Finetuned language models are zero-shot learners. ArXiv , abs/2109.01652, 2021.

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.

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? llm-scale data valuation with influence functions. ArXiv , abs/2405.13954, 2024.

Ian Covert, Chanwoo Kim, Su-In Lee, James Zou, and Tatsunori Hashimoto. Stochastic amortization: A unified approach to accelerate feature and data attribution. ArXiv , abs/2401.15866, 2024.

Chin-Yew Lin. Rouge: A package for automatic evaluation of summaries. In Annual Meeting of the Association for Computational Linguistics , 2004.

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.09803, 2022.

Steven G. Krantz and Harold R. Parks. The implicit function theorem. 2002.

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 , abs/2308.03296, 2023.

Garima Pruthi, Frederick Liu, Mukund Sundararajan, and Satyen Kale. Estimating training data influence by tracking gradient descent. ArXiv , abs/2002.08484, 2020.

SungYub Kim, Kyungsu Kim, and Eunho Yang. Gex: A flexible method for approximating influence via geometric ensemble. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems , volume 36, pages 5888-5911. Curran Associates, Inc., 2023.

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.

Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. Less: Selecting influential data for targeted instruction tuning. ArXiv , abs/2402.04333, 2024.

Stephen E. Robertson and Hugo Zaragoza. The probabilistic relevance framework: Bm25 and beyond. Found. Trends Inf. Retr. , 3:333-389, 2009.

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, ChengYu 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, Alaaeldin 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-lm: In search of the next generation of training sets for language models. ArXiv , abs/2406.11794, 2024.

Jacob Devlin, Ming-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.

Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bertnetworks. In Conference on Empirical Methods in Natural Language Processing , 2019.

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.17413, 2024.

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.

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-8153, 2020.

Logan Engstrom, Axel Feldmann, and Aleksander Madry. Dsdm: Model-aware dataset selection with datamodels. ArXiv , abs/2401.12926, 2024.

Karen Spärck Jones. A statistical interpretation of term specificity and its application in retrieval. J. Documentation , 60:493-502, 2021.

Zehan Li, Xin Zhang, Yanzhao Zhang, Dingkun Long, Pengjun Xie, and Meishan Zhang. Towards general text embeddings with multi-stage contrastive learning. ArXiv , abs/2308.03281, 2023.

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, Ke-Yang 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.10671, 2024.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2017.

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.

Yizhong Wang, Hamish Ivison, Pradeep Dasigi, Jack Hessel, Tushar Khot, Khyathi Raghavi Chandu, David Wadden, Kelsey MacMillan, Noah A. Smith, Iz Beltagy, and Hannaneh Hajishirzi. How far can camels go? exploring the state of instruction tuning on open resources. ArXiv , abs/2306.04751, 2023.

Josef Dai, Xuehai Pan, Ruiyang Sun, Jiaming Ji, Xinbo Xu, Mickel Liu, Yizhou Wang, and Yaodong Yang. Safe rlhf: Safe reinforcement learning from human feedback. ArXiv , abs/2310.12773, 2023.

Pang Wei Koh, Kai-Siang Ang, Hubert Hua Kian Teo, and Percy Liang. On the accuracy of influence functions for measuring group effects. In Neural Information Processing Systems , 2019.

Andrea Schioppa, Polina Zablotskaia, David Vilar, and Artem Sokolov. Scaling up influence functions. In AAAI Conference on Artificial Intelligence , 2021.

Yongchan Kwon, Eric Wu, Kevin Wu, and James Zou. Datainf: Efficiently estimating data influence in lora-tuned llms and diffusion models. ArXiv , abs/2310.00902, 2023.

Jiachen T. Wang, Dawn Xiaodong Song, James Zou, Prateek Mittal, and Ruoxi Jia. Capturing the temporal dependence of training data influence. ArXiv , abs/2412.09538, 2024.

Faisal Ladhak, Esin Durmus, and Tatsunori Hashimoto. Contrastive error attribution for finetuned language models. In Annual Meeting of the Association for Computational Linguistics , 2022.

Qingyi Liu, Yekun Chai, Shuohuan Wang, Yu Sun, Keze Wang, and Hua Wu. On training data influence of gpt models. In Conference on Empirical Methods in Natural Language Processing , 2024.

Zichun Yu, Spandan Das, and Chenyan Xiong. Mates: Model-aware data selection for efficient pretraining with data influence models. ArXiv , abs/2406.06046, 2024.

Juhan Bae, Nathan Ng, Alston Lo, Marzyeh Ghassemi, and Roger Baker Grosse. If influence functions are the answer, then what is the question? ArXiv , abs/2209.05364, 2022.

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.

Vitaly Feldman and Chiyuan Zhang. What neural networks memorize and why: Discovering the long tail via influence estimation. ArXiv , abs/2008.03703, 2020.

Amirata Ghorbani and James Y. Zou. Data shapley: Equitable valuation of data for machine learning. ArXiv , abs/1904.02868, 2019.

R. Jia, David Dao, Boxin Wang, Frances Ann Hubis, Nicholas Hynes, Nezihe Merve G"urel, 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.

Shuming Kong, Yanyan Shen, and Linpeng Huang. Resolving training biases via influencebased data relabeling. In International Conference on Learning Representations , 2022.

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.

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.

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 10958-10967, 2020.

Sejoon Oh, Sungchul Kim, Ryan A. Rossi, and Srijan Kumar. Influence-guided data augmentation for neural tensor completion. Proceedings of the 30th ACM International Conference on Information & Knowledge Management , 2021.

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 , abs/2411.12580, 2024.

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/2402.16827, 2024.

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.00913, 2023.

Ritik Sachin Parkar, Jaehyung Kim, Jong Inn Park, and Dongyeop Kang. Selectllm: Can llms select important instructions to annotate? ArXiv , abs/2401.16553, 2024.

Yunfan Shao, Linyang Li, Zhaoye Fei, Hang Yan, Dahua Lin, and Xipeng Qiu. Balanced data sampling for language model training with clustering. ArXiv , abs/2402.14526, 2024.

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.04320, 2023.

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.

Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in-context learning. ArXiv , abs/2112.08633, 2021.

Jeff Johnson, Matthijs Douze, and Herv'e J'egou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data , 7:535-547, 2017.

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 a short (1-2 sentence) justification 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 it (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 a 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 a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a 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 .

1. 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 in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in 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 it 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]

Justification: We have discussed the limitations of this paper in the conclusion section. Guidelines:

The answer NA means that the paper has no 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: In 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.

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 a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to 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 to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an 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. In the case of closed-source models, it 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 to have some path to 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 instructions 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/ public/guides/CodeSubmissionPolicy ) 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 a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

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) is 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, hyperparameters, 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 not include experiments.

The experimental setting should be presented in the core of the paper to a 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, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We have included variance information of scores from repeated experiment in Figure.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at 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, or 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 or the standard error of the mean.

It 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 a 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 of range (e.g. negative error rates).

If error bars are reported in tables or 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have included detailed information on the training and inference computation resources in our evaluation section. Appendix provides additional details.

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 of compute required for each of the individual experimental runs as well as 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 it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our research conform, in every respect, with the NeurIPS Code of Ethics.

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 circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal 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 or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to 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 to 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 or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

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 a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: 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?

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 a URL.

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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset'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 on existing open-source datasets and models.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

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 to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our research does not involve crowdsourcing / human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing 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 to 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 an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or 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 as any important, original, or non-standard components.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

A Deriving the Influence Function

Let θ ∗ denote the minimizer of the empirical risk:

Assume that ℓ is twice differentiable and strongly convex in θ . Thus, the Hessian

exists and is positive definite, ensuring the existence of H -1 .

The perturbed parameters θ ∗ ϵ,z are defined as:

Define the parameter change due to the perturbation as ∆ ϵ = θ ∗ ϵ,z -θ ∗ . Since θ ∗ does not depend on ϵ , the quantity we seek can be written in terms of ∆ ϵ :

Because θ ∗ ϵ,z is a minimizer, its first-order optimality condition is:

For small ϵ , using the Taylor expansion around θ ∗ :

Neglecting higher-order terms, this reduces to:

Since θ ∗ minimizes L ( θ ) , it satisfies ∇ L ( θ ∗ ) = 0 . Dropping higher-order terms in ϵ yields:

Thus, the derivative of θ ∗ ϵ,z with respect to ϵ at ϵ = 0 is:

To compute the influence of upweighting z on the loss at a test point x , we use the chain rule:

The influence function is therefore given by:

B Deriving the Group Influence Function

For group influence function, we interest in removing a group of training examples.

Let θ ∗ denote the minimizer of the empirical risk:

By up-weighting a subset of examples S = ( z 1 , · · · , z m ) , we define the new objective:

where ˆ ϵ = m n -m ϵ . We can see if ϵ = 0 we get the original loss function L ( θ ) (where none of the training samples are removed) and if ϵ = -1 , we get the loss function where samples in S are removed from training.

The perturbed parameters θ ∗ ϵ,S can be written as:

Define the parameter change as ∆ ϵ = θ ∗ ϵ,S -θ ∗ . Since θ ∗ does not depend on ϵ , the quantity we seek to compute can be written in terms of ∆ ϵ :

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 ϕ ( z ) = H -1 2 θ ∇ θ ℓ ( z ; θ ) , and assume ∇ k θ ℓ ( z ; θ ) for k ≥ 3 is negligible. Then the k -order group influence function is given by:

where: c ( k ) t are constants at the t -th order, dependent on the meta-information of the distribution of S and α t ( z ) is defined as:

From Eq. 10, we observe that the interaction of data samples exhibits an attention structure, with α t ( z i ) defining the weights of each data sample contributing to the final output. Intuitively, when k > 1 , examples z i that are closer to the group effect are assigned higher weights. We find the theoretical analysis of high-order group influence derived a similar notion of 'attention-based influence pooling' in AirRep, though AirRep's design can be viewed as a simplified version of B.1. These analysis also motivate us to explore more complicated influence pooling mechanism in our future work.

The following section provide the derivation of B.1.

B.2 Deriving the First-Order Group Influence Function

Since θ ∗ ϵ,S minimizes L S ( θ ) , the first-order optimality condition is:

For small ϵ , expand around θ = θ ∗ and ϵ = 0 :

At ϵ = 0 , the perturbed objective reduces to the original loss, so:

The Hessian at ϵ = 0 is:

The first-order change in the gradient with respect to ϵ at θ = θ ∗ is:

Substituting into the expansion, we have:

Solving for ∆ ϵ :

Thus:

Applying the chain rule to measure the influence of the group S on the loss at a test point x :

Substituting the expression for dθ ∗ ϵ,S dϵ , we obtain the first-order group influence function:

where c (1) = 1 n -m .

B.3 Deriving the Second-Order Group Influence Function

Since θ ∗ ϵ,S minimizes L S , its gradient vanishes:

We perform a Taylor expansion around ϵ = 0 and θ = θ ∗ , keeping terms up to ϵ 2 . Let:

For the linear term ( ϵ 1 ), we have:

Hence:

This result matches the first-order group influence function.

Second-Order ϵ For the quadratic term ( ϵ 2 ), expanding the gradient ∇ θ L S and collecting all ϵ 2 contributions yields:

Solving for ∆ 2 :

Since the weights of the training examples are linear in ϵ , the second derivative with respect to ϵ of these weights is zero:

Assuming higher-order terms ∇ k θ ℓ ( z ; θ ) for k ≥ 3 are negligible, the ∇ 3 θ L ( θ ∗ ) term vanishes. Thus:

From the definition of L S , we have:

Rearranging:

Influence on the Test Point For a test point x , consider ℓ ( x, θ ∗ ϵ,S ) . We aim to expand it as a Taylor series in θ around θ ∗ . Following previous studies, we perform a first-order expansion rather than a second-order one. This choice is justified because the dominant second-order effects can typically be captured through the parameter shift ∆ ϵ . Thus, we obtain:

The term multiplying ϵ corresponds to the first-order group influence, while the term multiplying ϵ 2 is the second-order correction. In closed form:

First-Order Group Influence:

Second-Order Influence:

where:

Simplifying Terms Now clean up the expressions f (2) IF ( x, S ) . First, define:

where ϕ ( z ) represents the preconditioned gradient of the loss.

Rearranging f (2) IF ( x, S ) , we get:

where terms A and B emerge.

Substituting ϕ ( z ) in A , we have:

For term B , we have:

Approximating the subset Hessian H S using the Fisher Information Matrix (FIM):

Substituting this approximation into B , we have:

For further simplification, we define:

which measures the alignment of ϕ ( z i ) with the aggregated gradients in S . Then:

Putting together, we get:

where c (1) and c (2) are the constants of the first- and second-order terms, defined as:

Additionally,

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, in order to have a more complete perception of the structure of the influence function (and also to address our curiosity), we will continue to derive the third-order group influence function.

Since θ ∗ ϵ,S minimizes L S , its gradient vanishes:

We perform a Taylor expansion around ϵ = 0 and θ = θ ∗ , keeping terms up to ϵ 3 . Let:

The first-order and second-order terms, ∆ 1 and ∆ 2 , have been derived in the previous section. Now focus on the third-order term, ∆ 3 .

Considering the ϵ 3 term and collecting all contributions of order ϵ 3 , we obtain:

Now analyze these terms in turn:

Term (i) vanishes since the per-sample weight depend linearly on ϵ :

Term (iii) vanishes because the per-sample weight is linear in ϵ :

Term (ii) can be written as:

Term (iv) is:

Solving for ∆ 3 yields

Similar to the previous derivation, this can be rearranged as:

Term (v) involves the third derivative with respect to θ . In most influence-function derivations, one often assumes that higher-order derivatives ∇ 3 θ L ( θ ∗ ) are negligible. Following this convention, we consider:

Term (vi) is the derivative with respect to ϵ of the third derivative with respect to θ . Since we neglect ∇ 3 θ ℓ , any derivative of it also vanishes. Thus, we have:

Putting everything together, the ϵ 3 stationarity condition simplifies to

Influence on the Test Point For a test point x , consider ℓ ( x, θ ∗ ϵ,S ) . Similar to the previous section, we can express the effects of S on x (when considering the ϵ 3 term) as:

The terms related to ∆ 1 and ∆ 2 have already been simplified in our derivation of the second-order group influence function, and we can write them as the following simplified expression:

Therefore, the first two terms A and B of f (3) IF ( x, S ) share the same structure, with constants slightly adjusted due to the inclusion of additional correction terms. Moreover, f (3) IF ( x, S ) will have an additional term C for the ϵ 3 term. Let:

Approximating the subset Hessian H S using the Fisher Information Matrix (FIM):

Substituting this approximation into C , we obtain:

Now, lets simplify C :

For further simplification, we define:

where α ( z i , S ) is the same as the previous definition, that is:

Combining this with A and B , we obtain:

We can then rename α , β , and the constants to obtain:

where c (3) t are constants defined by the distributional properties of S , and

These are the alignment functions, which can also be defined recursively as:

This gives us the same format as Theorem B.1.

B.5 Deriving the Higher-Order Group Influence Function

We extend the previous derivation to 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:

For second order:

Then, we can express C as:

For third order:

Define:

We see that at each order k in ϵ , the only 'new' term that enters the equations (and thus the recursion for ∆ k ) is the first derivative with respect to ϵ multiplied by lower-order ∆ j 's, plus the usual expansions in powers of ∆ ϵ (where ∆ ϵ = θ ∗ ϵ,S -θ ∗ ).

To explain this, we start with the stationarity condition at all ϵ :

Since θ ∗ is the minimizer at ϵ = 0 , we know G ( θ ∗ , 0) = 0 .

Expanding G ( θ ∗ +∆ ϵ , ϵ ) using a multi-variable Taylor series and keeping terms up to first order in ∆ ϵ gives:

Define:

Because each training sample's weight depends linearly on ϵ , we have:

Thus, higher-order terms in ∂ 2 ∂ϵ 2 ∇ θ L S , ∂ 3 ∂ϵ 3 ∇ θ L S , etc., vanish, simplifying the stationarity expansions. In other words:

Furthermore, collecting the ϵ k terms in the stationarity condition G ( θ +∆ ϵ , ϵ ) = 0 for k ≥ 1 gives:

Since ∆ 0 = 0 by definition, and ∆ 1 is determined by [ ∂ ∂ϵ G ( θ, 0) ] , we solve for:

By repeated substitution, we obtain closed-form expressions for all ∆ k :

Order 1:

Order k ≥ 2 :

This leads to the recursively defined alignment functions in the final group influence function formula. This completes the proof of Theorem B.1 for all cases where k > 3 .

C Comparison between gradient-based and representation-based methods

From Eq. 3 and Eq. 5, it is evident that gradient-based and representation-based 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 gradientor 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-based methods rely on heuristic metrics 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 to 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-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:

FLAN: 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.

Alpaca: We use the original Alpaca training set ( https://huggingface.co/datasets/ tatsu-lab/alpaca ) as the training data. For the test data, we design the test set to consist of two parts: (i) a seen subset, where we randomly sample 250 instructions from the training set and use a Llama model trained on the full Alpaca training set to generate the responses; (ii) an unseen subset, where we sample 250 examples from https://huggingface.co/ datasets/tatsu-lab/alpaca_eval . Thus, the total test set consists of 500 examples.

Tulu: We use version v1 of Tulu ( https://huggingface.co/datasets/allenai/ tulu-v1-sft-mixture ). The original dataset contains over 490k instances. We randomly sample 100,000 examples for our training set and 500 for our test set.

SafeRLHF: SafeRLHF ( https://huggingface.co/datasets/PKU-Alignment/ PKU-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-based methods and representation-based methods.

Gradient-based methods rely on the gradient of data with respect to a warmed-up model. To warm up the model, for FLAN, we fine-tune it on 100,000 examples randomly sampled from the FLAN training set. For Alpaca, Tulu, and SafeRLHF, we warm up the model on 100,000 examples randomly sampled from UltraChat.

LoGra LoGra is a toolkit for efficient influence function calculation ( https://github.com/ logix-project/logix ), and its core techniques are introduced in Choe et 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 to 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 it to 4, 8, and 16, resulting in a final dimension ranging from 1,152 to 27,648.

TracIn TracIn [22] is derived from a first-order approximation of model weights and essentially computes the dot product between gradients without Hessian correction. For implementation efficiency, 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 LM.

LESS LESS [25] adopts LoRA tuning instead of full fine-tuning and projects the LoRA gradient to a lower dimension using a random matrix. Its key idea is to use AdamW states (first- and second-order momentum) to correct the gradients. We use the official implementation ( https: //github.com/princeton-nlp/LESS ) and set the projection dimensions to 768 and 8192, which are the default dimensions in the LESS paper.

For LM warm-up, we follow their paper by setting the LoRA rank to 128 and alpha to 512, finetuning for 4 epochs. Similar to TracIn, the original LESS method ensembles influence scores from 4 checkpoints. For a 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 a 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-based method that provides an efficient implementation of the influence function. TRAK approximates the Hessian using the Fisher Information 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://github.com/MadryLab/DsDm ), we approximate its implementation (i.e., TRAK) using 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.

The following are representation-based baselines.

22.91

21.37

21.11

LoGra

20

20.07

19.36

19.75

AirRep

15.84

15

S

13.11

LD

10

9.92

5.46

5

500

800

1K

5K

10K

SubsetSize

TF-IDF We use sklearn to calculate the cosine similarity of TF-IDF feature of data.

DSIR DSIR utilizes hashed n-gram features, where n-grams are mapped onto a fixed number of buckets for efficiency and scalability. The influence estimator is parameterized by a bag-of-words generative model on the hashed n-grams, learned simply by counting the hash bucket frequencies. We use the official implementation ( https://github.com/p-lambda/dsir ).

RDS RDS was introduced in Hanawa et al. [6], which uses the model's hidden states as data representations. Through an evaluation of different representation design choices, we use average pooling over the last-layer hidden states of all tokens as the sentence embedding.

GTE GTE [35] is a general text embedding model trained on a mixture of text retrieval, text classification, text clustering, and text entailment data. It achieves strong performance on text embedding leaderboards relative to its model size (e.g., https://huggingface.co/spaces/mteb/ 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 a given data collection, and require the full retraining when new instances are added. AirRep, on the other hand, uses a trained encoder to produce the embeddings for any new instance on the fly. As a 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.

F Training Cost

For data generation, we create 10K subsets (each containing 1K data examples) and train the model on each subset for 2 epochs. This results in a total of 10M training examples for the Qwen2.5-0.5B LM, requiring about 20 hours on eight A100 GPUs. During the optimization stage, each training step involves sampling 32 subsets (1K examples each), totaling 32K examples per gradient descent step. The model is trained for up to 2K steps, completing in about 5 hours on an eight-GPU machine. While this training cost exceeds the warm-up training requirements of gradient-based TDA methods, it is exchanged for significant efficiency gains during inference. See Section 9 for detailed analysis.

LoGra (Qwen-0.5B)

16.79

LoGra (Llama-1B)

19.95

AirRep

22.54

0

5

10

15

20

25

LDS on FLAN (Target LM = Llama-1B)

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 m = 100 due to our data sampling strategies), denoted as ( x 1 , x 2 , . . . , x m ) . Let the training set be ( z 1 , z 2 , . . . , z n ) . The TDA method computes an influence score f ( x i , z j ) for each pair ( x i , z j ) , where i ∈ [ m ] and j ∈ [ n ] . Next, we determine the rank of each training example relative to a given test example. Specifically, we define rank( z i , x j ) ∈ [ n ] as the position of x j among the most influential training examples for z i . In other words, if rank( z i , x j ) = 1 , then x j is the most influential training example for z i . We then define the score of each test example x j as the best (lowest) rank it achieves across all training examples: Score( x j ) = min(rank( z i , x j ) | i ∈ [ m ]) . This represents the highest influence rank of x j among all test examples. Finally, we rank the training data based on Score( x j ) and select the top 1,000 examples.

After training the target LM on the selected data, we evaluate the model on the 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 F1 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 a single token ( positive or negative ), making unigram F1 equivalent to accuracy. Finally, the final F1 score is averaged over 66 tasks in 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 a practical configuration for data selection from large-scale corpora

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 5. We find that AIRREP outperforms baseline methods in this setting as well.

H.2 Data Classification Results

I Detailed Data Selection Results

See Table 7, 8, 9 for detailed results of data selection.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-05-31"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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?\n\nQUESTIONS:\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- 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.\n\nLIMITATIONS:\nYes, shortly in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nOne 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\nNone\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-04"}