# Enhancing Training Data Attribution with Representational Optimization

arXiv: https://arxiv.org/abs/2505.18513 · HTML: https://mlsanity.com/paper/2505.18513/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 53 of its 67 references that are in the map.
- First posted on arXiv 24 May 2025.

## Builds on (most cited first, 30 of 53)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Enhancing Chat Language Models by Scaling High-quality Instructional Conversations](https://mlsanity.com/paper/2305.14233/) (2023)
- [Towards General Text Embeddings with Multi-stage Contrastive Learning](https://mlsanity.com/paper/2308.03281/) (2023)
- [Safe RLHF: Safe Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2310.12773/) (2023)
- [Learning To Retrieve Prompts for In-Context Learning](https://mlsanity.com/paper/2112.08633/) (2021)
- [Data Shapley: Equitable Valuation of Data for Machine Learning](https://mlsanity.com/paper/1904.02868/) (2019)
- [Learning from Noisy Labels with Deep Neural Networks: A Survey](https://mlsanity.com/paper/2007.08199/) (2020)
- [Estimating Training Data Influence by Tracing Gradient Descent](https://mlsanity.com/paper/2002.08484/) (2020)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [What Neural Networks Memorize and Why: Discovering the Long Tail via Influence Estimation](https://mlsanity.com/paper/2008.03703/) (2020)
- [How Far Can Camels Go? Exploring the State of Instruction Tuning on Open Resources](https://mlsanity.com/paper/2306.04751/) (2023)
- [DataComp-LM: In search of the next generation of training sets for language models](https://mlsanity.com/paper/2406.11794/) (2024)
- [Data Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (2023)
- [Studying Large Language Model Generalization with Influence Functions](https://mlsanity.com/paper/2308.03296/) (2023)
- [TRAK: Attributing Model Behavior at Scale](https://mlsanity.com/paper/2303.14186/) (2023)
- [Towards Efficient Data Valuation Based on the Shapley Value](https://mlsanity.com/paper/1902.10275/) (2019)
- [Influence Functions in Deep Learning Are Fragile](https://mlsanity.com/paper/2006.14651/) (2020)
- [A Survey on Data Selection for Language Models](https://mlsanity.com/paper/2402.16827/) (2024)
- [Datamodels: Predicting Predictions from Training Data](https://mlsanity.com/paper/2202.00622/) (2022)
- [On the Accuracy of Influence Functions for Measuring Group Effects](https://mlsanity.com/paper/1905.13289/) (2019)
- [If Influence Functions are the Answer, Then What is the Question?](https://mlsanity.com/paper/2209.05364/) (2022)
- [Unified Demonstration Retriever for In-Context Learning](https://mlsanity.com/paper/2305.04320/) (2023)
- [Scaling Up Influence Functions](https://mlsanity.com/paper/2112.03052/) (2021)

## Built on by (most cited first, 8 of 8)

- [GradAlign: Gradient-Aligned Data Selection for LLM Reinforcement Learning](https://mlsanity.com/paper/2602.21492/) (2026)
- [NanoKnow: How to Know What Your Language Model Knows](https://mlsanity.com/paper/2602.20122/) (2026)
- [A Bayesian Information-Theoretic Approach to Data Attribution](https://mlsanity.com/paper/2604.03858/) (2026)
- [Training data attribution in diffusion models via mirrored unlearning and noise-consistent skew](https://mlsanity.com/paper/2605.17938/) (2026)
- [On the Fragility of Data Attribution When Learning Is Distributed](https://mlsanity.com/paper/2605.15520/) (2026)
- [idSCD: Identifying Training Datasets through Semantic Correlation Descriptors](https://mlsanity.com/paper/2605.30462/) (2026)
- [Influcoder: Distilling Decoders' Gradient Influence Rankings into an Encoder for Data Attribution](https://mlsanity.com/paper/2606.13668/) (2026)
- [Effective Synthetic Data Curation Requires Group-Level Signals](https://mlsanity.com/paper/2610.00779/) (2026)

## Abstract (arXiv)

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