# TRAK: Attributing Model Behavior at Scale

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

## Key facts

- Cited by 283 papers in the map (data through 14 Oct 2026).
- Builds on 80 of its 111 references that are in the map.
- First posted on arXiv 24 Mar 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [A Unified Approach to Interpreting Model Predictions](https://mlsanity.com/paper/1705.07874/) (2017)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Language Models as Knowledge Bases?](https://mlsanity.com/paper/1909.01066/) (2019)
- [mT5: A massively multilingual pre-trained text-to-text transformer](https://mlsanity.com/paper/2010.11934/) (2020)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness](https://mlsanity.com/paper/1811.12231/) (2018)
- [LaMDA: Language Models for Dialog Applications](https://mlsanity.com/paper/2201.08239/) (2022)
- [Reproducible scaling laws for contrastive language-image learning](https://mlsanity.com/paper/2212.07143/) (2022)
- [BadNets: Identifying Vulnerabilities in the Machine Learning Model Supply Chain](https://mlsanity.com/paper/1708.06733/) (2017)
- [Adversarial Examples Are Not Bugs, They Are Features](https://mlsanity.com/paper/1905.02175/) (2019)
- [Quantifying Memorization Across Neural Language Models](https://mlsanity.com/paper/2202.07646/) (2022)
- [Deduplicating Training Data Makes Language Models Better](https://mlsanity.com/paper/2107.06499/) (2021)
- [Data Shapley: Equitable Valuation of Data for Machine Learning](https://mlsanity.com/paper/1904.02868/) (2019)
- [Estimating Training Data Influence by Tracing Gradient Descent](https://mlsanity.com/paper/2002.08484/) (2020)
- [Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks](https://mlsanity.com/paper/1901.08584/) (2019)
- [What Neural Networks Memorize and Why: Discovering the Long Tail via Influence Estimation](https://mlsanity.com/paper/2008.03703/) (2020)
- [Underspecification Presents Challenges for Credibility in Modern Machine Learning](https://mlsanity.com/paper/2011.03395/) (2020)

## Built on by (most cited first, 30 of 283)

- [SalUn: Empowering Machine Unlearning via Gradient-based Weight Saliency in Both Image Classification and Generation](https://mlsanity.com/paper/2310.12508/) (2023)
- [DataInf: Efficiently Estimating Data Influence in LoRA-tuned LLMs and Diffusion Models](https://mlsanity.com/paper/2310.00902/) (2023)
- [DsDm: Model-Aware Dataset Selection with Datamodels](https://mlsanity.com/paper/2401.12926/) (2024)
- [MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models](https://mlsanity.com/paper/2406.06046/) (2024)
- [ContextCite: Attributing Model Generation to Context](https://mlsanity.com/paper/2409.00729/) (2024)
- [What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions](https://mlsanity.com/paper/2405.13954/) (2024)
- [A Primer on the Inner Workings of Transformer-based Language Models](https://mlsanity.com/paper/2405.00208/) (2024)
- [Cognitive Mirage: A Review of Hallucinations in Large Language Models](https://mlsanity.com/paper/2309.06794/) (2023)
- [Data Shapley in One Training Run](https://mlsanity.com/paper/2406.11011/) (2024)
- [A Survey of Multimodal Large Language Model from A Data-centric Perspective](https://mlsanity.com/paper/2405.16640/) (2024)
- [Describing Differences in Image Sets with Natural Language](https://mlsanity.com/paper/2312.02974/) (2023)
- [Intriguing Properties of Data Attribution on Diffusion Models](https://mlsanity.com/paper/2311.00500/) (2023)
- [Evaluating Data Attribution for Text-to-Image Models](https://mlsanity.com/paper/2306.09345/) (2023)
- [JiuZhang3.0: Efficiently Improving Mathematical Reasoning by Training Small Data Synthesis Models](https://mlsanity.com/paper/2405.14365/) (2024)
- [The Journey, Not the Destination: How Data Guides Diffusion Models](https://mlsanity.com/paper/2312.06205/) (2023)
- ["According to ...": Prompting Language Models Improves Quoting from Pre-Training Data](https://mlsanity.com/paper/2305.13252/) (2023)
- [Scalable Influence and Fact Tracing for Large Language Model Pretraining](https://mlsanity.com/paper/2410.17413/) (2024)
- [Feature-Learning Networks Are Consistent Across Widths At Realistic Scales](https://mlsanity.com/paper/2305.18411/) (2023)
- [Prismatic Synthesis: Gradient-based Data Diversification Boosts Generalization in LLM Reasoning](https://mlsanity.com/paper/2505.20161/) (2025)
- [SoK: Memorization in General-Purpose Large Language Models](https://mlsanity.com/paper/2310.18362/) (2023)
- [LLMs Meet Multimodal Generation and Editing: A Survey](https://mlsanity.com/paper/2405.19334/) (2024)
- [Training Data Attribution via Approximate Unrolled Differentiation](https://mlsanity.com/paper/2405.12186/) (2024)
- [A Survey on Mechanistic Interpretability for Multi-Modal Foundation Models](https://mlsanity.com/paper/2502.17516/) (2025)
- [Influence Functions for Scalable Data Attribution in Diffusion Models](https://mlsanity.com/paper/2410.13850/) (2024)
- [Give Me the Facts! A Survey on Factual Knowledge Probing in Pre-trained Language Models](https://mlsanity.com/paper/2310.16570/) (2023)
- [Most Influential Subset Selection: Challenges, Promises, and Beyond](https://mlsanity.com/paper/2409.18153/) (2024)
- [TAGCOS: Task-agnostic Gradient Clustered Coreset Selection for Instruction Tuning Data](https://mlsanity.com/paper/2407.15235/) (2024)
- [Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget](https://mlsanity.com/paper/2407.15811/) (2024)
- [Unleashing the Power of Data Tsunami: A Comprehensive Survey on Data Assessment and Selection for Instruction Tuning of Language Models](https://mlsanity.com/paper/2408.02085/) (2024)
- [There Will Be a Scientific Theory of Deep Learning](https://mlsanity.com/paper/2604.21691/) (2026)

## Abstract (arXiv)

The goal of data attribution is to trace model predictions back to training data. Despite a long line of work towards this goal, existing approaches to data attribution tend to force users to choose between computational tractability and efficacy. That is, computationally tractable methods can struggle with accurately attributing model predictions in non-convex settings (e.g., in the context of deep neural networks), while methods that are effective in such regimes require training thousands of models, which makes them impractical for large models or datasets. In this work, we introduce TRAK (Tracing with the Randomly-projected After Kernel), a data attribution method that is both effective and computationally tractable for large-scale, differentiable models. In particular, by leveraging only a handful of trained models, TRAK can match the performance of attribution methods that require training thousands of models. We demonstrate the utility of TRAK across various modalities and scales: image classifiers trained on ImageNet, vision-language models (CLIP), and language models (BERT and mT5). We provide code for using TRAK (and reproducing our work) at https://github.com/MadryLab/trak .