LM Data Attribution Design Space: Negative Log-Odds Gradients Boost Recall

Language model data attribution methods are typically evaluated via downstream tasks such as fact retrieval, but such evaluations vary multiple dataset properties at once, making it difficult to pinpoint the specific causes of attribution algorithm failures. A new study systematically dissects the data attribution design space by constructing a controlled benchmark and proposes two overlooked design choices: the choice of loss function and the aggregation of token scores.

The Dilemma of Attribution Evaluation: Confounding Factors Are Hard to Isolate

The pain point of existing data attribution evaluation lies in "variable coupling." A single evaluation set may simultaneously contain redundant documents, lexical distractors, texts of varying lengths, and more. When attribution performs poorly, researchers cannot determine whether redundant data diluted the signal or distractors misled the gradients. This confounding leaves algorithm improvements without a clear direction.

To address this problem, the researchers propose the Controlled Attribution Benchmark (CATT), focusing on the fact-tracing task. Its core idea is to fine-tune language models on entirely fictional knowledge, thereby providing unambiguous ground truth for attribution. Since the knowledge is fictional, any fact output by the model can be precisely traced to a specific document in the training data, eliminating the ambiguities common in real-world corpora.

Data Attribution aims to quantify the contribution of each training sample to a given model prediction. Common methods include influence functions, gradient similarity (e.g., TracIn), and leave-one-out approaches. Its core assumption is that model outputs can be decomposed as a sum of contributions from training samples, enabling the localization of key data in debugging, auditing, and copyright tracing. However, attribution results are highly sensitive to training dynamics, loss functions, and aggregation methods. When multiple dataset properties vary simultaneously during evaluation, it becomes difficult to determine whether performance degradation stems from algorithmic deficiencies or changes in data configuration—this is the origin of the "variable coupling" dilemma described in the original text.

Redundancy and Distractors: Persistent Killers of Attribution Performance

CATT's synthetic setup allows researchers to change one dataset configuration at a time, such as document redundancy and lexical distractors. This configurable evaluation reveals two stable patterns: redundancy and distractors consistently degrade attribution performance.

Redundant documents mean the same fact appears repeatedly across multiple documents, making it difficult for attribution methods to distinguish which document is the "true" source; lexical distractors introduce superficially similar but semantically irrelevant text that interferes with gradient signals. Both factors are ubiquitous in real corpora, but prior evaluations could not quantify their effects individually. CATT's value lies in separating the two, making failure causes diagnosable.

Lexical distractors refer to words or phrases inserted into text that are superficially similar to the target fact but semantically irrelevant—for example, the irrelevant city name in the rewrite "Paris is the capital of France, but Berlin is the capital of Germany." Such distractors do not change the fact itself but alter the model's internal representations and gradient signals, causing attribution methods to misallocate attention to distractor words. In real corpora, paraphrasing, entity substitution, and templated sentences all produce similar effects. CATT therefore studies them as an independent controllable variable, helping distinguish between two types of noise with different natures: "semantic redundancy" and "surface-level distraction."

Loss Function and Aggregation: Two Underestimated Design Dimensions

Rather than proposing a new estimator, the researchers examine two underexplored design choices: the loss function used to compute gradients, and how per-token attribution scores are aggregated into document-level scores.

Causal Evidence Is Highly Localized

Analysis reveals that the true causal evidence in training documents is highly localized—only a few key tokens carry the factual association. But standard methods average attribution scores across the entire document, diluting the key signal into noise. This "average aggregation" appears robust but actually obscures the most important evidence.

Negative Log-Odds Gradients Highlight Key Tokens

The researchers find that replacing standard cross-entropy loss gradients with negative log-odds gradients can highlight key tokens from a large number of uninformative tokens. On this basis, aggregating only the top token scores achieves a 35 to 72 percentage point improvement in ground truth recall over multiple baseline methods on the diagnostic benchmark.

This improvement does not rely on a new attribution estimator but rather corrects the loss function and aggregation strategy in the existing pipeline, giving it good generality.

Negative log-odds is a loss form commonly used in classification tasks, defined as the negative of the difference between the log probability of the correct class and the log probability of the incorrect class. Compared with standard cross-entropy, it is more sensitive to the "margin between the correct class and the strongest incorrect class," and gradients concentrate more on tokens that truly change the prediction tendency. In language model fact tracing, this means gradient signals are not diluted by large numbers of high-probability but uninformative tokens (such as function words and punctuation), allowing key entities or relation words to receive higher attribution scores. This property is the mathematical basis for "highlighting key tokens" described in the original text.

Generalization Validation: Counterfactual Retraining on FTRACE-TREX

The generalization ability of the negative log-odds correction is validated on the FTRACE-TREX benchmark: for every method tested, counterfactual retraining performance improves, by up to 4.9×. Counterfactual retraining measures the change in model predictions after removing the key documents identified by attribution, and is a metric closer to real-world application than recall.

This result shows that the loss function correction not only improves recall on the diagnostic benchmark but also transfers to more realistic evaluation scenarios.

Counterfactual retraining is a downstream metric for evaluating data attribution quality: an attribution method first selects the training documents most important to a given prediction, then those documents are removed from the training set and the model is retrained to observe how much the target prediction changes. If attribution is accurate, removing key documents should cause significant prediction changes; if attribution is wrong, removing irrelevant documents should leave predictions nearly unchanged. This metric is closer to real-world application than recall because it directly measures the causal impact of attribution results on model behavior, but it is computationally expensive and typically used only on smaller-scale benchmarks.

Implications for Data Attribution Practice

This study sends a clear signal: the performance bottleneck of data attribution may not lie in the estimator itself but may hide in "peripheral" designs such as loss functions and aggregation strategies. When standard cross-entropy gradients drown causal evidence, switching to negative log-odds gradients and focusing on top tokens can yield significant gains at low cost.

CATT's synthetic approach is also worth learning from—constructing unambiguous ground truth with fictional knowledge and then adjusting dataset properties one by one provides a reproducible diagnostic tool for attribution research. For engineering scenarios that require auditing training data influence and tracing fact sources, this "diagnose first, then improve" path is more actionable than blindly swapping estimators.

One-line Summary

A new study builds the controlled benchmark CATT, revealing that loss functions and aggregation methods are overlooked key design dimensions in data attribution.

Paragraph Summary

Existing language model data attribution evaluations vary multiple dataset properties simultaneously, making it difficult to pinpoint failure causes. A new study proposes the controlled attribution benchmark CATT, which fine-tunes models on fictional knowledge to provide unambiguous ground truth and adjusts configurations such as redundancy and lexical distractors one at a time. Experiments reveal that redundancy and distractors consistently degrade attribution performance; they also find that true causal evidence is highly localized and that standard average aggregation dilutes key signals. Switching to negative log-odds gradients and aggregating only top token scores yields 35 to 72 percentage point recall improvements on the diagnostic benchmark and up to 4.9× improvement in counterfactual retraining performance on FTRACE-TREX.