Swapping cross-entropy for negative log-odds gradients in LLM data attribution surfaced the critical training tokens and lifted ground-truth recall by 35 to 72 points
AspireAI community · arXiv (Yu, Wang, Hamidieh, Torralba, Bau)
Data attribution methods for language models are usually evaluated on downstream tasks like fact retrieval, where multiple dataset properties change at once and it is impossible to isolate why an algorithm fails. The Controlled Attribution Benchmark (CATT) fixes that by fine-tuning LLMs on entirely fictitious knowledge, giving unambiguous ground truth, and then altering one dataset configuration at a time - document redundancy, lexical distractors - to diagnose failures. The paper finds redundancy and distractors consistently degrade attribution, and that the true causal evidence inside a training document is highly localized but standard methods dilute it into noise by averaging attribution scores across the whole document. Replacing standard cross-entropy loss gradients with negative log-odds gradients surfaces the critical tokens from the uninformative ones, and aggregating only the top token scores improves ground-truth recall by 35 to 72 percentage points across multiple baselines on the diagnostic benchmark. The correction also generalizes to FTRACE-TREX, improving counterfactual retraining performance for every method tested, by up to 4.9x. The limitation is the synthetic setup - fictitious facts make attribution unambiguous by construction, and the paper's own framing is that this is a diagnostic benchmark rather than a new estimator.
Read the story and explore more from the AspireAI community.
Two sparse-autoencoder analyses have resisted scaling: matching semantically similar features across layers, and compressing large feature circuits into interpretable supernodes. This paper argues both are the same underlying problem, estimating semantic distances between SAE features that sit on different activation manifolds. Its representation treats each feature not as a single decoder vector but as an activation-weighted distribution over the hidden states expressing it, projects those distributions into a shared reference space, and compares them with Wasserstein distance. The representation is proved invariant to activation rescaling and stable under perturbation, and recovers true matches under finite-sample margin conditions. Empirically it outperformed decoder-vector and LLM-based baselines and captured subtle functional distinctions between related features, and it compressed large feature circuits into interpretable supernodes automatically. The limitation is that this is interpretability tooling: the paper demonstrates better feature matching and circuit compression, not that downstream interpretability claims become more accurate.