Enhancing Training Data Attribution with Representational Optimization
Weiwei Sun, Haokun Liu, Nikhil Kandpal, Colin Raffel, Yiming Yang
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 rigor, their computational costs make them impractical for large-scale applications. Representation-based attribution methods are more efficient, relying on similarity computations between examples in some representation space, but they often lack task-aware and model-specific optimization, limiting their accuracy. To address these challenges, we propose AirRep, a novel representation-based approach that enhances representation quality through task-driven optimization of a representation encoding model.Furthermore, we extend this method beyond single-sample attribution using an attention-based pooling mechanism to effectively estimate the collective influence of groups of samples.Experiments on instruction tuning of large language models demonstrate that AirRep achieves performance on par with state-of-the-art gradient-based approaches while being nearly two orders of magnitude more efficient. Further analysis highlights its robustness, including generalization to new data and new TDA tasks.