Enhancing Training Data Attribution with Representational Optimization

Yiming Yang (University College London, University of London) · Weiwei Sun (Carnegie Mellon University) · Haokun Liu (Department of Computer Science, University of Toronto) · Nikhil Kandpal (Department of Computer Science) · Colin Raffel (University of Toronto, Vector Institute and Hugging Face)

attention-based poolingattribution qualitycomputational efficiencyempirical effectgeneralization across tasks.group-wise influenceheuristic embeddingsinfluence functionsinstruction-tuned llmsmodel-aligned representationsranking objectiverepresentation-based approachesscalable methodstask-specific learningtraining data attribution

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.