A Bayesian Information-Theoretic Approach to Data Attribution

Abstract: 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.

• The paper introduces a novel method for training data attribution by quantifying the increase in posterior predictive entropy when specific data subsets are withheld.
• It employs a Gaussian Process surrogate model via neural tangent kernel approximations to achieve efficient, scalable computations.
• Empirical evaluations in image and NLP tasks demonstrate superior performance in brittleness, backdoor detection, and coreset selection compared to traditional baselines.

Bayesian Information-Theoretic Training Data Attribution

Problem Formulation and Bayesian Perspective

Training Data Attribution (TDA) aims to trace model decisions to specific training examples, serving concerns of interpretability, fairness, safety, and data valuation. Traditional influence function methods, which locally approximate the effect of upweighting or removing training examples, are limited by strict analytic assumptions and computational bottlenecks related to Hessian inversion and model convergence.

This work redefines TDA as an information-theoretic problem: given a query input $x_*$, it quantifies the contribution of a training subset $S$ as the information loss—the increase in posterior predictive entropy (epistemic uncertainty) at $x_*$ when $S$ is withheld. This approach focuses attribution strictly on behaviorally relevant resolution of uncertainty in model predictions, explicitly separating it from label noise effects. A Gaussian Process (GP) surrogate model, induced via the empirical neural tangent kernel (NTK) built from network Jacobians at the trained parameters, provides closed-form tractable approximations to the entropy and variance terms necessary for scalable computation.

Figure 1: Bayesian information-theoretic TDA compares withheld information loss and influence function estimation, visualizing the divergence in ranked attributions between methods.

From Information Loss to Efficient Algorithms

While information loss provides a principled counterfactual objective, it is neither submodular nor efficiently optimizable in large-scale settings due to the need for repeated conditioning on nearly full data. The authors demonstrate a formal relaxation to an information gain objective, justified in a high-noise regime where both information loss and gain share the same dominant asymptotic expansion. This relaxation renders the selection process submodular, providing approximation guarantees for greedy optimization.

To further scale to modern regimes—enabling ANN retrieval in vector databases—the authors introduce a linear-response variance correction. This reduces greedy step evaluation to squared inner-product searches between a query-dependent residual (given by the Schur complement of sketched Jacobians) and a bank of precomputed Jacobian sketches, all within a low-dimensional projected space. This supports tractable selection, rapid candidate retrieval, and integration with large vector database infrastructures such as FAISS.

Figure 2: Fraction of information recoverable by InfoGain and its approximation relative to InfoLoss, as a function of observation noise and subset size.

Subset Attribution, Diversity, and Baseline Comparisons

Unlike additive influence function-based methods, the information-theoretic framework accounts for interactions and redundancy, directly maximizing uncertainty reduction. With greedy subset selection, this leads to diverse, non-redundant sets of influential training examples.

The visual comparison (see Figure 1, right) between influence-based and entropy-based criteria on real data reveals clear divergences. Influence scores often highlight globally influential or outlier points, while information-theoretic attributions select diverse examples aligned with epistemic resolution.

Empirical Evaluation

Leave-Subset-Out Brittleness

The core empirical protocol follows the counterfactual leave-subset-out approach: for test queries, top-attributed training subsets are removed and the model is retrained; attribution quality is measured by the fraction of queries flipped to incorrect predictions (the "brittleness" metric).

Figure 3: Fraction of target test queries misclassified after removal of top-attributed same-label training data; the information-theoretic methods dominate existing baselines across image (Fashion-MNIST, CIFAR-10) and NLP (RTE) domains.

The Bayesian information-theoretic methods—InfoLoss, InfoGain, and InfoGain (approximate)—consistently outperform influential baselines (KronInfluence, TRAK) in inducing brittleness across both vision and text domains, with the margin widening at larger removal budgets.

Backdoor Detection and Data Poisoning Attribution

Applying the methodology to synthetic backdoor attack scenarios (BadNets on CIFAR-10), InfoLoss and InfoGain methods almost perfectly retrieve the ground-truth backdoored examples, sharply outperforming influence function and nearest-neighbor similarity methods.

Figure 4: Backdoored CIFAR-10 images showing $3\times3$ triggers, which InfoLoss methods consistently retrieve as most influential.

(Table 1)

Method	Recall@50	MRR
Random	0.011	0.045
GradDot	0.010	0.217
RepSim	0.985	1.000
TRAK	0.009	0.189
KronInfluence	0.058	0.504
InfoLoss	1.000	0.999
InfoGain	0.974	0.998
InfoGain (approx)	0.974	0.998

Table 1: InfoLoss achieves exact backdoor recall and ranking, with InfoGain variants closely trailing and outperforming all baselines.

Coreset Construction

For coreset selection (constructing minimal training subsets supporting high test accuracy after retraining), InfoGain and its scalable approximation substantially outperform all influence-based and similarity-based baselines, particularly in low-budget regimes. Baseline methods often underperform random selection unless class balance constraints are enforced.

Figure 5: CIFAR-10 coreset accuracy after retraining on attributed subsets; InfoGain and InfoGain (approx) yield markedly stronger performance, especially at small coreset sizes.

Figure 6: Class-balanced coreset selection elevates baseline performance but InfoGain variants remain superior.

Visualization and Qualitative Insights

The paper provides clear visualizations showing marked qualitative differences between Bayesian information-theoretic methods, influence-based methods, and similarity-based heuristics when retrieving training examples most responsible for a given prediction.

Figure 7: Attributions via InfoLoss, InfoGain, and InfoGain (approx) for CIFAR-10 queries demonstrate consistency and class-alignment.

Figure 8: Influence-based baselines exhibit bias toward outliers or proto-typical instances depending on the metric.

Figure 9: RepSim retrieves visually similar examples but lacks explicit epistemic alignment.

Theoretical Implications and Future Directions

The Bayesian information-theoretic framework introduces solid theoretical grounding to TDA, directly linking selection to epistemic risk and generalization uncertainty. The GP/NTK surrogate closes the gap between black-box inference and analyticity, and the relaxations proposed enable tractable, scalable subset selection. The high-noise equivalence between InfoLoss and InfoGain formalizes when submodular relaxations are appropriate, and the use of sketching and ANN search makes the approach amenable to trillion-scale LLM regimes.

Practically, this work supports safety auditing, training data valuation, synthetic data detection, and privacy risk estimation in scalable, database-friendly form. Future research should explore extensions beyond tangent kernels (e.g., finite-width or non-linear posterior corrections), principled integration with training dynamics, and applications in online and federated settings.

Conclusion

This paper establishes a new standard for training data attribution by rigorously framing the problem via Bayesian information loss and providing tractable relaxations for practical deployment. The resulting methods consistently outperform classical TDA baselines on brittleness, backdoor retrieval, and coreset selection—without sacrificing computational efficiency—and offer a principled bridge between theory and the demands of scalable attribution for modern neural architectures.

Reference:

"A Bayesian Information-Theoretic Approach to Data Attribution" (2604.03858)