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.
As large language models are increasingly trained and fine-tuned, practitioners need methods to identify which training data drive specific behaviors, particularly unintended ones. Training Data Attribution (TDA) methods address this by estimating datapoint influence. Existing approaches like influence functions are both computationally expensive and attribute based on single test examples, which can bias results toward syntactic rather than semantic similarity. To address these issues of scalability and influence to abstract behavior, we leverage interpretable structures within the model during the attribution. First, we introduce Concept Influence which attribute model behavior to semantic directions (such as linear probes or sparse autoencoder features) rather than individual test examples. Second, we show that simple probe-based attribution methods are first-order approximations of Concept Influence that achieve comparable performance while being over an order-of-magnitude faster. We empirically validate Concept Influence and approximations across emergent misalignment benchmarks and real post-training datasets, and demonstrate they achieve comparable performance to classical influence functions while being substantially more scalable. More broadly, we show that incorporating interpretable structure within traditional TDA pipelines can enable more scalable, explainable, and better control of model behavior through data.
Data attribution seeks to quantify how individual training examples shape a model's predictions and underpins problems including data valuation, machine unlearning, and model interpretability. Despite having a long line of work, computationally scalable methods often struggle to predict the effect of removing training data in neural networks due to their non-convex nature. To overcome this challenge, metagradient-based methods such as MAGIC (Ilyas and Engstrom, 2025) differentiate each prediction through the entire training run and compute its exact influence with respect to the training data, but require a separate run for every prediction. To reduce this cost, we cast budgeted attribution as estimating a large influence matrix from a small number of measurements. We show that the measurements most appropriate for recovering this matrix differ from those best suited for attribution itself. We then present two algorithms, MAGE and SPELL, suited for reconstruction and attribution respectively, that run on existing metagradient machinery at no extra cost. Empirical studies demonstrate strong performance over existing baselines across training scales and measurement budgets.
Yuxi Chen, Hamza Golubovic, Han Tong +2
Carnegie Mellon University 1Department of Statistics & Data Science · Columbia University 2Department of Statistics · Columbia University 3Electrical Engineering Department +2
Tabular foundation models (TFMs) achieve strong predictive performance by conditioning on labeled demonstrations provided in context, without any parameter update. Yet how individual demonstrations shape a given prediction remains poorly understood. This gap matters in practice: the context is often assembled from whatever labeled data is available, potentially leading to the inclusion of mislabeled, redundant, or low-quality examples that degrade performance. Standard data attribution methods do not transfer to the TFM setting: resampling-based approaches such as DemoShapley require a combinatorial number of forward passes, and gradient-based estimators such as influence functions require computing training point's effect on the model parameters, which in-context learning never updates. We introduce TICDA, a method that measures the influence of every demonstration in the context directly from linear surrogates trained on TFM latent embeddings, in a single forward pass and at negligible cost. We show that TICDA offers the best compromise against competitors across four tasks: detecting labeling errors, curating context to preserve predictive accuracy while lowering inference cost, producing attribution scores that transfer across TFMs, and supporting an acquisition strategy for efficient active learning.
Yacine Benihaddadene, Milan Bhan, Eliot Dugelay +4