Training Data Attribution (TDA) seeks to trace model predictions back to influential training examples, enhancing interpretability and safety. We formulate TDA as a Bayesian information-theoretic problem: subsets are scored by the information loss they induce - the entropy increase at a query when removed. This criterion credits examples for resolving predictive uncertainty rather than label noise. To scale to modern networks, we approximate information loss using a Gaussian Process surrogate built from tangent features. We show this aligns with classical influence scores for single-example attribution while promoting diversity for subsets. For even larger-scale retrieval, we relax to an information-gain objective and add a variance correction for scalable attribution in vector databases. Experiments show competitive performance on counterfactual sensitivity, ground-truth retrieval and coreset selection, showing that our method scales to modern architectures while bridging principled measures with practice.
Figures & tables
Figure 1: Bayesian information-theoretic training data attribution (TDA). For a query input x∗ , we quantify the contribution of a training subset S via its information loss : the increase in (posterior) predictive entropy at x∗ when S is withheld from training, attributing credit to examples that resolve epistemic uncertainty rather than label noise. Left: we relate information loss to a submodular information gain relaxation whose leading-order term matches information loss in a high-noise regime, followed by a linear-response variance correction that implements greedy selection with a squared inner-product score, enabling efficient retrieval in a vector database. Right: CIFAR-10 (ResNet-9) examples showing the top-ranked training images under an influence-function estimator versus our information loss criterion.
Figure 2: Information efficiency experiments on binary CIFAR-10. Top: selecting subsets via greedy InfoGain and varying the observation noise. Bottom: fixed observation noise level ( σ2=103 ), using InfoGain and approximate InfoGain.
Figure 3: Our Bayesian information-theoretic methods— InfoLoss , InfoGain , and InfoGain (approx) —recover strong attribution signal in identified subsets, as reflected by higher brittleness. We plot the fraction of previously correct test queries that become misclassified after removing the same-label training examples attributed by each method and retraining. Left: Fashion-MNIST (MLP), where our methods dominate and KronInfluence is the closest baseline. Middle: CIFAR-10 (ResNet-9), where our methods lead across budgets with TRAK most competitive. Right: RTE (BERT), where KronInfluence is strongest at smaller budgets, while InfoLoss / InfoGain catch up and are competitive at larger budgets. Error bars represent standard error across repeated runs.
Figure 4: Examples of backdoored CIFAR-10 images with triggers. The trigger is a 3×3 black/white random pattern placed at the bottom-right corner. This causes the model to misclassify the image to the corrupted class.
Method
Recall@50 ( ↑ )
MRR ( ↑ )
Random
0.011
±
0.003
0.045
±
0.016
GradDot
0.010
±
0.002
0.217
±
0.033
RepSim
0.985
±
0.004
1.000
±
0.000
TRAK
0.009
±
0.002
0.189
±
0.036
KronInfluence
0.058
±
0.011
0.504
±
0.034
InfoLoss
1.000
±
0.000
0.999
±
0.001
Table 1: Our Bayesian information-theoretic method InfoLoss retrieves the ground-truth backdoor examples almost perfectly on CIFAR-10—achieving the highest recall and matching the best baseline on MRR (no significant difference). Moreover, InfoGain and InfoGain (approx) show no significant difference from the best methods on MRR, and their recall is close to the best baseline. We compare standard TDA baselines to our proposed methods (shaded). We denote the best within each block (Baselines / Proposed) with underline and bold denotes the best overall. Reported metrics are accompanied by standard error from repeated runs.
Figure 5: InfoGain and InfoGain (approx) construct substantially better CIFAR-10 coresets than existing TDA baselines which perform worse than random selection. We plot test accuracy after retraining on the selected subset for coreset sizes ranging from 0.2% to 10% of the training set. Error bars show standard error across repeated runs.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: CIFAR-10 coreset selection with class-balanced TDA baselines. Enforcing equal per-class quotas substantially improves TRAK , KronInfluence , GradDot and RepSim , often lifting the strongest baselines above InfoLoss ; however, both InfoGain variants remain strongest overall.
Figure 7: Visualization of training data attribution using our Bayesian information-theoretic methods on the CIFAR-10 dataset. Top: InfoLoss , Middle: InfoGain , Bottom: InfoGain (approx) .
Figure 8: Visualization of training data attribution using baseline influence-based methods on the CIFAR-10 dataset. Top: KronInfluence , Middle: TRAK , Bottom: TracIn .
Figure 9: Visualization of training data attribution using representational similarity ( RepSim ) on the CIFAR-10 dataset.
Carnegie Mellon University 1Department of Statistics & Data Science · Columbia University 2Department of Statistics · Columbia University 3Electrical Engineering Department +2