Enhancing Training Data Attribution with Representational Optimization

We thank the reviewer for engaging with our work and for recognizing the importance of TDA for model transparency, as well as AirRep’s strong performance. We will address each of your points in detail below, with particular focus on clarifying our contributions to AirRep’s novelty, scalability, and performance as a principled and robust framework for TDA.

---

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].

Thank you for pointing this out! While LDS has been a widely used benchmark in recent TDA research, we agree that evaluating attribution quality with alternative metrics adds valuable perspective. To this end, we conducted a Leave-One-Out (LOO) analysis, where we removed individual training examples and measured the Pearson correlation between the predicted influence scores and the actual change in test loss. The results on a FLAN subset are summarized below:

AirRep achieves higher correlation than baseline methods, suggesting that its optimization-based approach provides more robust and reliable individual-level attribution. We will include this analysis and discussion in the revised version of the paper.

---

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.

Thank you for raising this concern. The clipping thresholds were set to 0.1 (low) and 5.0 (high), as described in the AirRep training details in Section 4. These values were not extensively tuned and were selected based on simple, generalizable observations:

• The low threshold (0.1) corresponds to approximately one standard deviations of the loss difference observed from repeated fine-tuning of language models on the same data, which we treat as typical variation due to training stochasticity.
• The high threshold (5.0) was empirically chosen to cap rare but extreme differences, which we consider outliers and clip to prevent them from disproportionately influencing the learning process.

Importantly, we applied the same thresholds and training hyperparameters across all experiments, and observed consistently stable and performant training. We will clarify these points in the revised version to ensure transparency and reproducibility.

---

There are loads of papers cited in their arxiv versions, not their published versions.

Thank you for pointing out this issue! We will correct these citations in our final version.

---

The paper would benefit from presenting a qualitative analysis, accompanied by concrete examples.

Thank you for the suggestion. We agree that including qualitative examples can provide deeper intuition into how AirRep assigns attribution scores. In the revised version, we will include detailed case studies and concrete examples to illustrate attribution behavior. Additionally, we will open-source our data and models to support transparency and enable further exploration by the community.

---

AirRep uses BERT for the encoder. Why BERT?

Thank you for this question. We initialize AirRep with a BERT-style encoder (specifically, GTE-small) due to its strong performance in text embedding tasks and ease of training. GTE-small is well-suited for pairwise scoring due to its retrieval-focused pretraining. That said, AirRep is not tied to any specific encoder architecture. We expect that other text embedding models could also be integrated successfully.

---

What exactly is the "re-training outcome" in line 83

The “re-training outcome” mentioned refers to r(x, S), defined in line 79 as the loss of the model on example x after training on dataset S. We recognize that this phrasing could cause confusion and will clarify this in our final version.

---

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?

In our evaluation, we did not retrain AirRep for each downstream task. Once trained, AirRep learns model-aligned representations for predicting influence scores via loss change estimation and can be applied to various downstream TDA tasks, such as data selection and classification. While we observe promising generalization to unseen data, we acknowledge that performance may vary in domains with substantially different distributions—such as complex reasoning tasks where long COT models generate significantly longer thinking content. Exploring domain-adaptive training strategies is a promising direction for future work.

Regarding training cost, Section 9 provides a detailed comparison with baseline TDA methods. We show that although AirRep incurs a one-time training cost, it becomes cost-effective after approximately 5.6 GPU hours, and achieves up to 80× faster inference than gradient-based approaches. Since the trained model can be reused across tasks, this cost can be amortized in multi-task settings, offering both principled modeling and practical scalability.

---

This paper only considers fine-tuning, where he knowledge bias present in the base model before fine-tuning is ignored, Is this understanding correct?

While our evaluation focuses on fine-tuned models, we do not ignore the influence of the pretrained model's biases. AirRep’s supervision signal is derived from actual training outcomes, meaning that any biases inherited from the base model and expressed during fine-tuning are naturally reflected in the measured loss changes. Thus, our attribution framework implicitly captures the combined effects of both pretraining and fine-tuning.

---

Is it possible to use KSD [3] as an evaluation metric?

Thank you for highlighting this interesting work. After reviewing [3], we agree that KSD presents a promising alternative for evaluating attribution methods. However, adapting it to our current LLM training setup would require non-trivial adjustments, particularly since the original implementation was designed for image classification.

That said, because [3] also evaluates methods like Influence Functions and TracIn, we believe KSD could potentially be extended to evaluate AirRep as well. We consider this an exciting direction for future work and will explore it in follow-up studies.

---

Thank you again for your thoughtful feedback—we hope these clarifications will resolve your concerns.

We sincerely thank the reviewer for the thoughtful and encouraging feedback. We appreciate the recognition of AirRep’s originality in formulating training data attribution as an optimization problem, its strong empirical performance and generalizability, and the connection to higher-order influence functions via attention-based pooling. We respond to each of your comments and suggestions in detail below.

---

It is unclear whether training an AirRep model requires specific tuning - how easy will it be to implement this?

Thank you for this question. In our experience, AirRep is relatively easy to train and implement. The training data is generated automatically by sampling random subsets from training corpus, without requiring any specialized design or manual curation. Training follows a standard contrastive learning setup and does not rely on extensive hyperparameter tuning.

Specifically, we used a fixed learning rate of 1e-4 and a batch size of 32, selected to fit typical GPU memory constraints. These settings yielded stable training across all experiments. We have also shared our data generation and training code in the Supplementary Material, which we hope will make implementation straightforward for future users.

---

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?

Thank you for this important question. In Section 9, we provide a quantitative analysis of how training set size impacts AirRep’s performance. In our experiments, training with several hundred random subsets over approximately 5 GPU hours was sufficient for AirRep to outperform LoGra across multiple tasks.

We found that the required number of subsets depends more on the diversity of the data distribution than on the model or dataset size. For example, more diverse datasets like Tulu typically benefit from more training subsets, whereas less diverse datasets such as Alpaca require fewer. Notably, we did not observe saturation in performance when increasing the number of training subsets in our current experiments. This suggests that the training cost can be flexibly traded off against performance and inference efficiency, depending on available resources.

---

How well does AirRep work for smaller/larger subset sizes than the ones reported in the paper?

Thank you for this thoughtful question. In Figure 6, we report results on the FLAN dataset across a range of subset sizes. Although AirRep was trained only on subsets of size 1000, it consistently outperforms LoGra across both smaller and larger subset sizes at test time. AirRep’s advantage becomes more pronounced with larger subsets, likely due to its attention-based pooling mechanism, which is designed to explicitly model group-level influence.

---

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?

Thank you for this insightful question. Capturing attribution at the individual-sample level is indeed challenging. Since AirRep is trained on many sampled subsets, this may have an averaging effect similar to repeated training, helping to stabilize the outcomes. Additionally, our importance reweighting objective is designed to reduce the influence of training noise.

In addition, to further assess AirRep's performance on single sample attribution, we evaluated it using a Leave-One-Out (LOO) metric. where we remove a single training example and measure the actual change in test loss, then compare it to the predicted influence score. The results on a FLAN subset are summarized below:

AirRep outperforms the baselines in this setting, suggesting that its group-level optimization framework can yield stable attribution at the individual-example level. We will include a more detailed discussion of this analysis in the revised paper.

---

Typos …

Thank you for pointing out the typos! We will correct these in the revised version of the paper.

---

Thank you again for your thoughtful feedback—we hope these clarifications will resolve your concerns.

Thank you for your thoughtful and constructive review—your feedback highlights the strengths of our method (e.g., attention-based pooling, confidence gating, and ablation clarity) while also providing valuable suggestions on evaluation direction, baseline completeness (e.g., TRAK and TracIn), and broader applicability, which we will carefully address to improve the paper. We respond to each of your points in detail below.

---

why the direction of evaluation has been done in reverse (adding influential training examples rather than removing training examples) compared to other papers.

Thank you for this insightful observation. We chose to focus on adding influential examples—rather than removing them—primarily due to computational efficiency. In removal-based evaluation (e.g., Leave-One-Out LOO), each example must be removed from the training set and the model’s performance must be re-evaluated, often requiring a full retraining per removal. In contrast, additive evaluation can be performed efficiently by training on a selected subset, which is far more efficient in practice.

However, following your suggestion, we conducted a LOO evaluation, in which individual examples were removed and we measured the correlation between predicted and actual influence scores. The results on a FLAN subset are summarized below:

We find that AirRep outperforms baseline methods in this setting as well. We will include this new analysis and discussion in the revised version of the paper.

For context, our evaluation setup otherwise follows established precedents: for LDS, we follow [Datamodels] to evaluate on randomly sampled subsets (see Figure 6 for different subset size), and for data selection and classification tasks, our setup mirrors that of LESS and Hanawa et al.

---

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.

Thank you for this thoughtful suggestion. We agree that fact tracing is a compelling and high-impact application of training data attribution. While time constraints during the rebuttal period may prevented us from conducting additional evaluations on fact-tracing dataset, we see this as a promising future direction. We will add a discussion of fact-tracing applications in the future work section of the revised paper and plan to explore corresponding evaluations in follow-up work.

---

One of the more fundamental issues is lack of some important works such as TRAK for comparison. Why wasn't TRAK evaluated as a baseline?

Thank you for raising this important point. In our experiments, we include DsDm [REF], an improved variant of TRAK specifically optimized for LLM-scale settings. In our evaluation, we found that DsDm performs comparably to LoGra and is consistently outperformed by AirRep across tasks. We will clarify the relationship between TRAK and DsDm in the revised version of the paper to avoid confusion and ensure proper attribution.

[REF] DsDm: Model-Aware Dataset Selection with Datamodels

---

It is not clear why each experimental section contains only a subset of the mentioned baselines. If there is a good reason for dropping an approach from evaluation, it should be explicitly mentioned. Why some baselines (such as TracIn) are not evaluated in some sections?

Thank you for pointing this out. In Table 1, we present a comprehensive comparison across all baselines. In subsequent experimental sections, we focus on a subset of baselines—specifically the most competitive or directly relevant ones (e.g., LoGra, LESS, and GTE-small)—to ensure clarity and conserve space in the main text. Baselines such as TracIn, which consistently underperform relative to LoGra, were omitted from detailed comparisons for conciseness.

That said, we understand that this selective presentation may have caused confusion. To address this, we will (1) include complete evaluation results for all baselines in the Appendix, and (2) explicitly state our rationale for baseline selection in each experimental section in the revised version of the paper.

---

The paper is well written in general. A few comments …

Thank you for your detailed and constructive suggestions for improving clarity. In response, we will streamline the Introduction and Preliminaries sections to improve readability and focus. We will also revise the title of Section 5 to more accurately reflect its content, explicitly distinguishing between data attribution and counterfactual correctness. Additionally, we will clarify our use of the term “prediction” to avoid confusion with its conventional meaning in model outputs.

---

Would certain set of encoding models work better than others?

Thank you for this interesting question. In our preliminary experiments, we compared GTE-small and DeBERTa-small, and observed that GTE-small achieved slightly better performance—likely due to its pre-training on text retrieval tasks, which aligns well with our pairwise comparison objective.

While GTE-small provided strong overall results, we agree that broader exploration of encoding models is worthwhile. In particular, models like BERTScore or ColBERT, which support fine-grained token-to-token matching, may offer further improvements in attribution fidelity. We view this as a promising direction and plan to explore it in future work.

---

One additional point could have been discussed as to whether this is NLP-specific TDA method or applicable to other general domains.

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.

We will include this point in the discussion of limitations and future work in the revised paper and plan to explore cross-domain generalization in follow-up research.

---

Thank you again for your thoughtful feedback—we hope these clarifications will resolve your concerns.

Thank you for your constructive feedback! We appreciate that you find our method principled and effective, our experiments thorough, and the performance is strong. We would like to respond to your comments as follows:

---

High initial supervision cost:

Thank you for raising this concern. In Section 9, we provide a quantitative analysis of AirRep’s initial training cost and demonstrate how it can be amortized during inference. While the upfront cost of training AirRep is indeed higher than that of baseline methods, this is amortized after approximately 5.6 GPU hours—beyond which AirRep yields an 80× speedup in inference. Given this substantial gain in downstream efficiency, we believe the one-time offline training cost is a worthwhile investment and does not pose a significant limitation in practical settings where such training time is likely affordable.

---

Proxy-model mismatch:

Thank you for this insightful point. We agree that proxy-target mismatch may lead to fidelity degradation as the model gap widens. To investigate this issue further, we conducted additional evaluations across a range of target models with varying sizes and architectures. The results on FLAN are summarized below:

These results show that AirRep consistently outperforms LoGra, even when supervision is derived from a much smaller proxy model (Qwen-0.5B), suggesting strong generalization across both model scale and architecture. While degradation is observed as the proxy-target gap increases (e.g., Qwen-0.5B → GPT-2), AirRep maintains robust relative gains compared to LoGra. Due to limited time for the rebuttal, we will include a more detailed discussion and evaluation on this point in the final version.

---

Minimal tuning of baselines:

Thank you for the thoughtful critique. Regarding baseline tuning, as noted in Appendix E, we carefully optimize both performance and efficiency for LoGra, TracIn and LESS. This included exploring different configurations and applying engineering optimizations on our evaluation data.

• For TracIn, we report results using the final checkpoint rather than ensembles, as ensembling across multiple checkpoints linearly increases computation and storage overhead. In our preliminary experiments, we found that ensembling yields modest gains compared to simply increasing the projection dimension, while incurring higher resource costs and latency. We will elaborate on this trade-off more fully in the revised version of the paper.

• For LESS, we used the LoRA rank recommended by the original paper. This setting already results in an embedding of dimensionality 8192—over 21× larger than that used by AirRep—making it a resource-intensive baseline.

---

Better iso-FLOP baselines:

Thank you for this valuable suggestion. In Section 9, we currently report a full breakdown of the total computational cost—including data generation, training, and inference—for both AirRep and LoGra. Our analysis shows that AirRep becomes cost-effective after approximately 5.6 GPU hours, beyond which it offers significant inference-time speedups (up to 80× faster). We note that this comparison represents an upper bound on AirRep’s training cost, as the learned representations can often be reused across different target models or tasks without retraining. This enables cost amortization in multi-task or multi-model scenarios, improving overall efficiency in practical deployments.

Following your suggestion, we will revise the paper to include a clearer iso-FLOP-style analysis by aligning LoGra’s total budget (including gradient computations, Hessian preconditioning, and pairwise scoring) with that of AirRep’s full training pipeline.

---

Thank you again for your thoughtful feedback—we hope these clarifications and forthcoming additions will resolve your concerns.