Language model data tracing works better by focusing on key tokens
Elucidating the Design Space of LM Data Attribution
Computation and Language
Summary
It’s hard to tell which pieces of training data truly influence a language model's answers because past tests changed many things at once. The authors created a special test using totally made-up facts so they know exactly where the correct information is. They found that looking at the most important words rather than averaging all words helps pinpoint what really influences the model. Also, changing the math used to measure influence makes it easier to spot the critical tokens that affect the model’s knowledge.
What this means in practice
- •For ai model developers: Improve diagnostics for which training data influences specific model outputs by focusing on key tokens with adjusted gradient measures.
- •For machine learning engineers: Build tools to better trace model facts back to training documents, improving transparency and debugging of language models.
Authors
Steven Yu, Sheng-Yu Wang, Kimia Hamidieh, Antonio Torralba, David Bau
Abstract
Data attribution methods for Language Models (LMs) are typically evaluated using downstream tasks like fact retrieval. However, existing evaluations vary multiple dataset properties simultaneously, making it impossible to isolate exactly why an attribution algorithm fails. To diagnose these failures, we explore the design space of LLM data attribution, with a primary focus on evaluation. We introduce the Controlled Attribution Benchmark (CATT) for fact tracing, in which we fine-tune LLMs on entirely fictitious knowledge, providing unambiguous ground truth for attribution. This synthetic setup allows us to systematically alter one dataset configuration at a time, such as document redundancy and lexical distractors. With this configurable evaluation, we demonstrate that redundancy and distractors consistently degrade attribution performance. Rather than proposing a new estimator, we examine two design choices that are underexplored: the loss used to compute gradients and how per-token attribution scores are aggregated into a document score. We find that the true causal evidence within a training document is highly localized, but standard methods dilute this evidence into noise by averaging attribution scores across the entire document. We find that replacing standard cross-entropy loss gradients with the negative log-odds gradients surfaces the critical tokens from other uninformative ones. We can then aggregate only the top token scores to improve ground truth recall by 35 to 72 percentage points across multiple baselines on our diagnostic benchmark. The negative log-odds correction also generalizes to FTRACE-TREX, improving counterfactual retraining performance for every method we test, by up to 4.9x.