On the Feasibility of In-Context Probing for Data Attribution

TL;DR

This paper explores in-context probing as a rapid alternative to gradient-based data attribution methods for identifying influential training data, demonstrating their correlation and similar downstream effects across NLP and synthetic tasks.

Contribution

It introduces in-context probing as a computationally efficient proxy for gradient-based data attribution, validated through empirical NLP experiments and synthetic data analysis.

Findings

ICP correlates well with gradient-based attribution on NLP tasks

Fine-tuning on influential data from both methods yields similar performance

Synthetic data experiments support the connection between ICP and gradient methods

Abstract

Data attribution methods are used to measure the contribution of training data towards model outputs, and have several important applications in areas such as dataset curation and model interpretability. However, many standard data attribution methods, such as influence functions, utilize model gradients and are computationally expensive. In our paper, we show in-context probing (ICP) -- prompting a LLM -- can serve as a fast proxy for gradient-based data attribution for data selection under conditions contingent on data similarity. We study this connection empirically on standard NLP tasks, and show that ICP and gradient-based data attribution are well-correlated in identifying influential training data for tasks that share similar task type and content as the training data. Additionally, fine-tuning models on influential data selected by both methods achieves comparable downstream performance, further emphasizing their similarities. We also examine the connection between ICP and gradient-based data attribution using synthetic data on linear regression tasks. Our synthetic data experiments show similar results with those from NLP tasks, suggesting that this connection can be isolated in simpler settings, which offers a pathway to bridging their differences.