Tensor Network Datamodels for Training Data Attribution

Farzaneh Heidari; Md Rifat Arefin; Louis Jaburi; Guillaume Rabusseau

cs.LGOct 17, 2026

Tensor Network Datamodels for Training Data Attribution

Authors: Farzaneh Heidari, Md Rifat Arefin, Louis Jaburi, Guillaume Rabusseau

Organizations: Mila – Quebec AI Institute · DIRO, Université de Montréal · EleutherAI · Canada CIFAR AI Chair

Abstract

Datamodels predict how a model's behavior changes when different parts of its training data are included or removed. However, linear datamodels assign fixed contributions to training units and cannot represent non-additive joint effects. We introduce Tensor Network Datamodels (TNDMs), which use low-rank tensor networks to model dependencies among training units while keeping the representation compact. TNDMs learn a map from subsets of the training data to the outcome of training on those subsets, allowing them to represent interactions among several training units. We prove that complementary and redundant training groups can induce high-degree subset functions with compact low-rank representations. A single fitted TNDM can then be used for several training data questions, including predicting the effect of removing data, measuring the importance of individual training units, identifying interactions between training units, and selecting useful subsets of the training data. We evaluate TNDMs on controlled tasks designed to contain strong interactions, as well as standard fine-tuning settings for classification and reasoning. On a controlled reasoning task, TNDMs outperform every baseline while observing only half as many training subsets, and TNDM-TT reduces prediction NMSE by $89\%$ relative to the linear datamodel after observing approximately $0.0061\%$ of all possible training subsets. In factual-association fine-tuning, a single TNDM-CP raises held-out $\mathrm{R}^2$ from $0.579$ to $0.668$ over the linear datamodel. These results demonstrate that compact low-rank models improve prediction when training data effects depend on combinations of training groups.

Figures & tables

Figure 1: Tensor-network datamodels. (a) Each observed subset $S_{n}$ of training groups (a row; shaded groups are included) gives a retraining outcome $v(S_{n})$ , one entry of the value tensor $\mathcal{T}_{v}$ . From $m\ll 2^{d}$ such entries, a TNDM fits a low-rank approximation $\widehat{v}$ of $\mathcal{T}_{v}$ : a rank- $R$ CP decomposition (factors joined by a copy tensor, black dot) or a tensor train with ranks $r_{j}$ . (b) Interaction degree (purple) and CP rank (green) define two nested families of set functions; linear datamodels have degree at most one. The diagrams illustrate two distinct inductive biases: low interaction degree and low tensor rank, with the latter able to represent high-order interactions compactly.

Panel of Figure 2 — Figure 2: Interaction degree and tensor rank measure different properties. High-order interactions can have low tensor rank. (a) Each query requires four training groups together, giving interaction degree four but CP rank two. (b) A query requires one group from each pair of alternatives, giving interaction degree four but TT ranks at most two when each pair is consecutive in the ordering.

Datamodel	NMSE, $m=512$	NMSE, $m=1{,}024$	Joint-removal RMSE
Linear	$0.687\pm 0.001$	$0.669\pm 0.001$	$2.741\pm 0.016$
Degree-2 ridge	$0.467\pm 0.001$	$0.348\pm 0.001$	$2.125\pm 0.007$
Degree-3 ridge	$0.510\pm 0.002$	$0.342\pm 0.002$	$3.869\pm 0.019$
Degree-3 OMP	$0.635\pm 0.016$	$0.271\pm 0.021$	$6.292\pm 0.057$
RBF	$0.386\pm 0.001$	$0.257\pm 0.001$	$2.100\pm 0.035$
TNDM-CP	$\mathbf{0.148\pm 0.026}$	$0.087\pm 0.002$	$1.939\pm 0.258$

Table 1: Prediction and joint-removal effects on the 24-group reasoning task. Entries show means

\pm

sample SD over three learner seeds; lower is better. The last column uses

m=1{,}024

and measures error in the utility change when all sources for one required relation are removed. Bold marks the lowest mean in each column. At 512 observations, both TNDMs outperform every non-tensor baseline evaluated at 1,024 observations; TNDM-TT also achieves the lowest joint-removal error.

Figure 3: Interaction recovery and Shapley-based selection on the 12-group reasoning task. (a) Average precision for ranking complete triples by predicted third-order differences, across observation budgets. (b) Completeness of subsets formed by selecting the three groups with the largest individual Shapley values at 160 new observations. Curves and bars show means across two learner seeds, each averaged over ten relation sequences; bands and whiskers show sample SD, and dots in (b) show both seed means. TNDM-TT improves interaction recovery at larger budgets, while Linear + TNDM-TT achieves the highest mean selection completeness at 160 new observations.

Table 5

Figure 4: Training data attribution on a modular retraining task. Relative

\ell_{2}

error for Shapley values and pairwise and third-order Shapley interaction indices over eight training groups. Exact references use all 256 retrained coalition outcomes. All queries reuse the same fitted surrogate, and budgets include validation observations. Errors are averaged over 20 coalition-sampling repetitions within each of three learner seeds. Curves show the mean of the three seed averages; bands and bars show

\pm 1

sample SD across those averages. Both axes are logarithmic; TNDM-TT uses additive-informed initialization. At 64 observations, TNDM-CP achieves the lowest mean error among the plotted methods for all three quantities.

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Fitting a TNDM as tensor completion. (a) Each observed coalition provides one entry of the coalition tensor. (b) The fitting objective in Eq. 12 estimates the parameters of a low-rank tensor network to the observed entries. (c) The fitted network defines predictions for unobserved coalitions. The example uses a tensor train; CP and tree networks use different factorizations of the same coalition tensor. A TNDM learns from a small set of measured training subsets and uses the shared low-rank structure to predict outcomes for subsets that were never retrained.

Panel of Figure 6 — Figure 6: Diagram notation. A node is a tensor and each leg is one mode, with its dimension in gray at the free end; joining two legs of equal dimension sums that index; a filled dot is a copy tensor; attaching $e_{0}=(0,1)^{\top}$ or $e_{1}=(1,0)^{\top}$ to a physical leg fixes absence or presence of a training unit; and a network with no free legs evaluates to a scalar. Only connectivity carries algebraic meaning. Tensor-network diagrams provide a compact graphical notation for the tensor operations used throughout the paper.

Panel of Figure 7 — Figure 7: LOO, Banzhaf, and Shapley are three readouts of the same marginal curve. Left: for a fixed $t$ , the rank-one weight tensor $\mathcal{W}_{i}(t)$ places $\bm{\delta}=\mathbf{e}_{1}-\mathbf{e}_{0}$ on site $i$ and $\mathbf{b}(t)=(t,1-t)^{\top}$ on every other site. Its contraction with the fitted TNDM costs $\mathcal{O}(dr^{2})$ . Right: LOO evaluates the curve at $t=1$ , Banzhaf at $t=\tfrac{1}{2}$ , and Shapley integrates the curve over $[0,1]$ . Once a TNDM is fitted, several familiar attribution scores can be obtained as different summaries of the same learned subset function.

Panel of Figure 8 — Figure 8: Summing predictions at a fixed coalition size. Contracting the fitted datamodel (blue) with a counting weight tensor (orange) computes $\sum_{|S|=s}\widehat{v}(S)$ . The counter starts at zero, increments when a group is included, and terminates at $s$ ; it has at most $d+1$ states. For the symmetric semivalue of group $i$ , the corresponding construction places the difference vector $\bm{\delta}$ at $i$ and counts only the remaining $d-1$ groups. A terminal count $s$ receives weight $p_{s}/\binom{d-1}{s}$ , where $p_{s}$ is the probability assigned to context size $s$ (Proposition 8 ). The tensor network representation lets us aggregate predictions over all subsets of a given size without enumerating those subsets individually.

Panel of Figure 9 — Figure 9: Interaction discovery for one training unit (example $i=4$ ). (a) A green $2\times 2$ matrix on every physical leg rewrites the network in the Möbius basis; fixing the transformed leg of unit $i$ to $1$ leaves an order- $(d-1)$ tensor whose entries are exactly the interactions containing $i$ . Unlike a scalar attribution, this tensor-valued transformation deliberately keeps its legs open. (b) Contracting those legs with the counting weight tensor of Figure 8 collapses the slice into the degree profile $H_{i,k}$ . The fitted TNDM retains information about interactions, allowing us to inspect how strongly each training unit participates in interactions of different orders.

Training groups	Relation sequences	Required relations	Groups per relation
8	8	3	2
12	10	3	2
24	10	2	3

Appendix

Table 4: Reasoning-task configurations. Required relations are those whose training examples depend on the selected coalition. Each relation sequence defines 31 evaluation queries, one per initial entity. The three reasoning settings vary the number of training groups while keeping the same basic relation-composition structure.

Method	$m=256$	$m=512$	$m=1{,}024$
Cardinality only	$0.972\;[0.971,0.972]$	$0.934\;[0.933,0.934]$	$0.918\;[0.917,0.919]$
Linear	$0.749\;[0.749,0.749]$	$0.687\;[0.686,0.687]$	$0.669\;[0.669,0.670]$
Degree-2 ridge	$0.626\;[0.625,0.627]$	$0.467\;[0.466,0.469]$	$0.348\;[0.347,0.348]$
Degree-3 ridge	$0.669\;[0.668,0.670]$	$0.510\;[0.508,0.511]$	$0.342\;[0.340,0.343]$
Degree-3 OMP	$0.859\;[0.857,0.860]$	$0.635\;[0.624,0.654]$	$0.271\;[0.247,0.284]$
RBF kernel ridge	$0.562\;[0.562,0.562]$	$0.386\;[0.385,0.387]$	$0.257\;[0.256,0.258]$

Appendix

Table 5: Held-out NMSE on the

d=24

reasoning task; lower is better. Entries are means over learner seeds 31, 43, and 59, with seed ranges in brackets. Each seed’s NMSE averages over ten relation sequences on 256 random test coalitions. The observation budget

m

includes the validation allocation. Seeds share the task and coalition design. Bold denotes the lowest mean within each budget. Method labels follow the fitting procedure in Appendix I.5 . Direct CP and TT give the lowest prediction errors across all observation budgets, while other tensor parameterizations do not show the same advantage.

Method	Joint, $m=256$	Joint, $m=512$	Joint, $m=1{,}024$	LOO, $m=1{,}024$
Cardinality	$3.472\pm 0.019$	$3.648\pm 0.202$	$3.698\pm 0.022$	$0.556\pm 0.007$
Linear	$2.562\pm 0.018$	$2.663\pm 0.257$	$2.741\pm 0.016$	$1.184\pm 0.004$
Degree-2 ridge	$2.672\pm 0.012$	$2.733\pm 0.297$	$2.125\pm 0.007$	$0.962\pm 0.007$
Degree-3 ridge	$4.324\pm 0.012$	$5.053\pm 0.230$	$3.869\pm 0.019$	$1.326\pm 0.005$
Degree-3 OMP	$3.780\pm 0.349$	$4.390\pm 0.361$	$6.292\pm 0.057$	$3.136\pm 0.026$
RBF	$\mathbf{2.077\pm 0.009}$	$2.484\pm 0.454$	$2.100\pm 0.035$	$1.011\pm 0.027$

Appendix

Table 6: RMSE of predicted changes in retrained utility, in nats; lower is better. Joint removal deletes every group supplying one required relation; leave-one-out (LOO) removal deletes one group. Entries show means

\pm

sample SD over learner seeds 31, 43, and 59, which share the task and coalition design. Bold marks the lowest mean in each column. TNDM-TT is most accurate for joint removals, while single-group removal is easier and is best matched by HT in this experiment.

	New coalition observations
Method	32	48	64	96	128	160
Linear	0.053	0.054	0.083	0.050	0.052	0.055
Degree-2 ridge	0.057	0.045	0.049	0.044	0.042	0.052
Degree-3 ridge	0.081	0.143	0.200	0.260	0.289	0.318
Degree-3 OMP	0.071	0.094	0.074	0.083	0.101	0.093
RBF kernel ridge	0.045	0.043	0.071	0.071	0.146	0.089

Appendix

Table 7: Average precision for interaction discovery on the

d=12

reasoning task; higher is better. Each entry averages over ten relation sequences and learner seeds 43 and 59. All 220 triples are ranked by the absolute predicted third difference at the empty coalition, with complete triples as positives. Budgets count new coalition observations and exclude the separate gate collection. CP and HT variants were evaluated only at 160 new coalition observations. This is a secondary endpoint; comparisons across methods and budgets are not multiplicity-corrected. TNDM-TT becomes substantially better at recovering the true complementary triples as more coalition outcomes are observed, although the strongest method depends on the observation budget.

Method	Top-1 completeness	Top-3 completeness	Top-3 utility (nats)
Linear	0.650	0.633	$-4.187$
Degree-2 ridge	0.900	0.950	$-0.585$
Degree-3 ridge	1.000	1.000	$-0.001$
Degree-3 OMP	0.500	0.500	$-5.546$
RBF kernel ridge	1.000	1.000	$-0.003$
CP	0.850	0.883	$-1.336$

Appendix

Table 8: Selection of size-four training coalitions at

d=12

and 160 new observations. Entries average over ten relation sequences and learner seeds 43 and 59. Top-1 completeness evaluates the highest-ranked recommendation; top-3 completeness is the fraction of complete sets among the three retained recommendations. Top-3 utility is their mean measured final-answer log probability. Higher values are better for all three metrics. Cardinality and exact leave-one-out selectors have identical reported results. Several nonlinear methods, including TNDM-TT, reliably identify complete subsets; stronger interaction recovery does not automatically translate into better subset selection.

	$m=512$		$m=1{,}024$
Method	Complete	Utility	Complete	Utility
Linear	8/10	$-2.324$	8/10	$-2.326$
Degree-3 ridge	10/10	$-0.001$	10/10	$-0.001$
Degree-3 OMP	10/10	$-0.001$	10/10	$-0.001$
RBF kernel ridge	10/10	$-0.001$	10/10	$-0.001$
CP	9/10	$-1.142$	10/10	$-0.001$

Appendix

Table 9: Retrained evaluation of the highest-ranked size-three coalition for each of ten relation sequences at

d=24

. Completeness counts sequences whose recommended coalition supports every required relation. Utility is the mean measured final-answer log probability in nats; higher is better. All entries use learner seed 31. These results concern one recommendation per sequence, rather than the top-three completeness metric averaged across learner seeds. TNDMs can select high-utility subsets, but their prediction advantage does not consistently produce better subset selection than the polynomial or kernel baselines.

Method	Complete (%)	Retrained utility (nats)
Linear + TNDM-TT	$\mathbf{60.0\pm 14.1}$	$\mathbf{-4.568\pm 1.582}$
TNDM-CP	$45.0\pm 7.1$	$-5.989\pm 0.770$
Degree-2 ridge	$45.0\pm 7.1$	$-6.354\pm 0.783$
Linear + TNDM-CP	$45.0\pm 7.1$	$-6.418\pm 0.923$
Linear + HT	$40.0\pm 0.0$	$-6.891\pm 0.051$
Degree-3 OMP	$40.0\pm 14.1$	$-6.510\pm 1.528$

Appendix

Table 10: Selection by individual Shapley values at 160 new observations, with 17 shared observations also used for fitting and validation. Entries show means

\pm

sample SD across two learner-seed means, each averaged over ten evaluation sequences. Higher is better in both columns. Completeness measures relation coverage; utility is measured after retraining. Bold marks the highest mean in each column. Random scores denotes the random-score baseline, not expected performance under uniform sampling of triples. At this budget, combining a linear model with TNDM-TT gives the strongest Shapley-based subset selection among the compared methods, though the ordering varies across budgets.

Method	32	48	64	96	128	160
Linear	20	20	15	35	20	35
Degree-2 ridge	40	5	40	50	35	45
Degree-3 ridge	40	5	75	40	30	20
Degree-3 OMP	30	20	30	35	40	40
RBF	35	10	55	65	15	20
TNDM-TT	15	5	25	30	10	40

Appendix

Table 11: Shapley-based subset selection. Columns give the number of new coalition observations. Entries are mean percentages of complete selected triples over blind learner seeds 43 and 59, each evaluated on ten relation sequences. Higher is better; bold marks all methods tied for the highest mean within each budget. Random scores reports the random-score baseline’s completeness, not the expected completeness of uniformly sampled triples. No method dominates across all budgets; attribution-based subset selection depends strongly on both the observation budget and the attribution rule.

Method	32	48	64	96	128	160
Linear	20	20	15	35	20	35
Degree-2 ridge	40	5	40	50	35	45
Degree-3 ridge	35	0	65	40	25	30
Degree-3 OMP	35	20	25	50	40	40
RBF	35	25	40	70	40	50
TNDM-TT	20	10	15	35	25	35

Appendix

Table 12: Banzhaf-based subset selection under the same protocol as Table 11 . Entries are mean percentages of complete selected triples; columns give new-observation budgets. Higher is better; bold marks all methods tied for the highest mean within each budget.

Datamodel configuration	Held-out $R^{2}$ $\uparrow$
Linear, tuned	$0.579\pm 0.048$
Matched-pair, tuned	$0.556\pm 0.054$
Degree-2, tuned	$0.565\pm 0.109$
TNDM-CP, single fit, rank 5	$0.668\pm 0.084$
TNDM-CP, best of 8, rank 5	$0.616\pm 0.133$
TNDM-CP, ensemble of 8 initializations	$0.719\pm 0.070$

Appendix

Table 13: Sport–Music log-probability prediction: mean

R^{2}

\pm

sample SD over 15 shared coalition splits. Ensembles use additional surrogate-fitting compute. Bold marks the highest mean. TNDM-CP outperforms the tuned non-tensor baselines, and ensembling several low-rank fits improves prediction further.

Estimator	Banzhaf $\downarrow$	Shapley $\downarrow$
Linear	0.1701	0.3555
Degree-2	0.2013	0.3721
RBF	0.1724	0.3368
TNDM-CP	0.1736	0.3358
TNDM-TT	0.1634	0.3285
MSR Banzhaf	0.2258	–

Appendix

Table 14: Full AG News RoBERTa–LoRA attribution comparison at 128 observed coalitions. Entries are pooled median relative

\ell_{2}

errors over six adapter ranks, three learner seeds, and 20 observation-sampling repeats per combination. MSR estimates Banzhaf values directly; KernelSHAP estimates Shapley values directly. Lower is better; bold marks the lowest median for each target. TNDM-TT is competitive for attribution on AG News, while dedicated estimators can still be better for specific targets such as Shapley values.

Setting	$m$	Linear	Polynomial	OMP	RBF	CP	TT
AG News	128	0.522	0.208	0.778	0.131	0.062	0.041
	256	0.492	0.049	0.400	0.047	0.032	0.029
	512	0.479	0.044	0.294	0.034	0.014	0.041
UCI HAR	128	0.761	0.671	0.854	0.713	0.756	0.743
	256	0.777	0.700	0.719	0.656	0.774	0.613
	512	0.405	0.706	0.771	0.382	0.175	0.151

Appendix

Table 15: Coalition prediction beyond relation composition. Entries are median test NMSE over three coalition splits; lower is better. The budget

m

includes both fitting and validation coalitions. AG News and UCI HAR use deterministic classifier refits with

d=10

. Coin flip uses frozen-model inference with

d=8

. Bold marks the lowest median in each row.

	Coin flip		GSM8K
Surrogate	$n_{\mathrm{fit}}=400$	$n_{\mathrm{fit}}=800$	$n_{\mathrm{fit}}=400$	$n_{\mathrm{fit}}=800$
Linear	$0.1779\pm 0.0123$	$0.1853\pm 0.0002$	$0.3262\pm 0.0113$	$0.2854\pm 0.0105$
RBF	$0.1755\pm 0.0131$	$0.1679\pm 0.0078$	$0.2885\pm 0.0057$	$0.2645\pm 0.0142$
Coefficient CP	$0.1754\pm 0.0154$	$0.1674\pm 0.0091$	$0.2890\pm 0.0087$	$0.2760\pm 0.0123$
TT, random	$0.1732\pm 0.0174$	$0.1680\pm 0.0077$	$0.2874\pm 0.0093$	$0.2758\pm 0.0121$
TT, semantic	$0.1733\pm 0.0174$	$0.1671\pm 0.0088$	$0.2887\pm 0.0113$	$0.2760\pm 0.0119$

Appendix

Table 16: Held-out coalition prediction for frozen-model in-context learning with

d=500

candidate demonstrations and four demonstrations per coalition. Entries are mean test RMSE

\pm

sample standard deviation across two seeds; lower is better. Each fitting budget is supplemented by 200 validation coalitions. At larger candidate-pool dimension, CP and TT remain competitive with nonlinear baselines and improve on the linear model, while RBF is strongest on GSM8K at the larger budget.

UCI HAR: 320 fitting, 32 validation; ten splits
Attribution	Linear	RBF	TT
Shapley	$0.047190\pm 0.000116$	$0.044079\pm 0.000344$	$0.041886\pm 0.000704$
Banzhaf	$0.001240\pm 0.000105$	$0.001132\pm 0.000062$	$0.000983\pm 0.000072$
Pairwise interaction	$0.049988\pm 0.000000$	$0.047985\pm 0.000215$	$0.045343\pm 0.000602$
Qwen probe: 512 fitting, 205 validation; three seeds
Shapley	$0.067948\pm 0.001012$	$0.009971\pm 0.002894$	$0.055446\pm 0.007836$

Appendix

Table 17: Attribution RMSE against references computed from all 1024 measured coalition outcomes. Entries are mean

\pm

sample standard deviation across the repetitions specified in each panel; lower is better. TNDM-TT gives the lowest attribution error on UCI HAR and improves over the linear baseline on the Qwen probe, while RBF is strongest on the probe task.

	$d=8$		$d=12$
Datamodel	LDS	Direct	LDS	Direct
Linear	$0.341\pm 0.053$	$0.341\pm 0.053$	$0.032\pm 0.021$	$0.032\pm 0.021$
Degree-2 ridge	$0.377\pm 0.055$	$0.730\pm 0.023$	$0.026\pm 0.018$	$0.014\pm 0.026$
Degree-3 ridge	$0.381\pm 0.057$	$0.735\pm 0.023$	$0.031\pm 0.009$	$0.037\pm 0.020$
RBF	$0.386\pm 0.054$	$0.738\pm 0.025$	$0.022\pm 0.003$	$0.045\pm 0.013$
TNDM-CP	$0.282\pm 0.135$	$0.571\pm 0.108$	$\mathbf{0.037\pm 0.001}$	$0.035\pm 0.034$

Appendix

Table 18: LDS and direct-prediction Spearman correlation on the same evaluation subsets. Entries are means

\pm

sample SD; higher is better. For

d=8

, the observation budget is 128 and evaluation uses unobserved size-four subsets. Statistics use 20 observation-sampling repeats at one learner seed, after averaging eight relation sequences within each repeat. For

d=12

, the budget is 160 new observations and evaluation uses all 220 size-three subsets. Statistics use two blind learner seeds, after averaging ten sequences within each seed. Bold marks the highest displayed mean per column. The two experiments use different protocols and should not be pooled. Summing individual attribution scores can discard information captured by nonlinear datamodels; direct subset predictions better preserve their joint effects.

Surrogate	Singleton quantities	Pairwise interactions
Linear	$1.13\times 10^{-14}$	$8.23\times 10^{-17}$
Coefficient CP	$1.33\times 10^{-15}$	$5.83\times 10^{-16}$
TT	$4.77\times 10^{-15}$	$4.50\times 10^{-15}$

Appendix

Table 19: Maximum absolute discrepancy between model-specific attribution computations and exhaustive enumeration of the same fitted surrogate. The tensor-based attribution computations match exhaustive enumeration to numerical precision.

Figure 10: Data contributions depend on the target and training stage. Shapley values of eight training data groups for arithmetic targets with two through five operators. Solid and dashed lines distinguish the two random halves of each operator-count group. Lines show three-seed means; bands span the seed minimum and maximum. The group matching the evaluation target becomes dominant later in training.

Training Data Attribution

Training data attribution identifies which examples, sources, or groups shape a model’s predictions, capabilities, and failures.

See all 80 papers

A Bayesian Information-Theoretic Approach to Data Attribution
Apr 4, 2026Dharmesh Tailor, Nicolò Felicioni, Kamil CiosekTraining Data Attribution
University of Amsterdam · Spotify
STRIDE: Training Data Attribution via Sparse Recovery from Subset Perturbations
Jun 3, 2026Rishit Dagli, Abir Harrasse, Luke Zhang +4Sparse Recovery Data Selection
Jinesis AI Lab, University of Toronto & Vector Institute · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Thoughtworks +3
Concept Influence: Leveraging Interpretability to Improve Performance and Efficiency in Training Data Attribution
Feb 16, 2026Matthew Kowal, Goncalo Paulo, Louis Jaburi +6Training Data Attribution
EleutherAI · University of Toronto

Apr 4, 2026J/K move · Enter open · S save

A Bayesian Information-Theoretic Approach to Data Attribution

Dharmesh Tailor, Nicolò Felicioni, Kamil Ciosek

University of Amsterdam · Spotify

Training Data Attribution

Open paper AskPDF arXiv

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.

Abstract

Training Data Attribution

Explore similar work

A Bayesian Information-Theoretic Approach to Data Attribution

STRIDE: Training Data Attribution via Sparse Recovery from Subset Perturbations

Concept Influence: Leveraging Interpretability to Improve Performance and Efficiency in Training Data Attribution