Tensor Network Datamodels for Training Data Attribution
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 relative to the linear datamodel after observing approximately of all possible training subsets. In factual-association fine-tuning, a single TNDM-CP raises held-out from to 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
| Datamodel | NMSE, | NMSE, | Joint-removal RMSE |
|---|---|---|---|
| Linear | |||
| Degree-2 ridge | |||
| Degree-3 ridge | |||
| Degree-3 OMP | |||
| RBF | |||
| TNDM-CP |
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
| Training groups | Relation sequences | Required relations | Groups per relation |
| 8 | 8 | 3 | 2 |
| 12 | 10 | 3 | 2 |
| 24 | 10 | 2 | 3 |
| Method | |||
|---|---|---|---|
| Cardinality only | |||
| Linear | |||
| Degree-2 ridge | |||
| Degree-3 ridge | |||
| Degree-3 OMP | |||
| RBF kernel ridge |
| Method | Joint, | Joint, | Joint, | LOO, |
|---|---|---|---|---|
| Cardinality | ||||
| Linear | ||||
| Degree-2 ridge | ||||
| Degree-3 ridge | ||||
| Degree-3 OMP | ||||
| RBF |
| 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 |
| Method | Top-1 completeness | Top-3 completeness | Top-3 utility (nats) |
|---|---|---|---|
| Linear | 0.650 | 0.633 | |
| Degree-2 ridge | 0.900 | 0.950 | |
| Degree-3 ridge | 1.000 | 1.000 | |
| Degree-3 OMP | 0.500 | 0.500 | |
| RBF kernel ridge | 1.000 | 1.000 | |
| CP | 0.850 | 0.883 |
| Method | Complete | Utility | Complete | Utility |
| Linear | 8/10 | 8/10 | ||
| Degree-3 ridge | 10/10 | 10/10 | ||
| Degree-3 OMP | 10/10 | 10/10 | ||
| RBF kernel ridge | 10/10 | 10/10 | ||
| CP | 9/10 | 10/10 | ||
| Method | Complete (%) | Retrained utility (nats) |
|---|---|---|
| Linear + TNDM-TT | ||
| TNDM-CP | ||
| Degree-2 ridge | ||
| Linear + TNDM-CP | ||
| Linear + HT | ||
| Degree-3 OMP |
| 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 |
| 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 |
| Datamodel configuration | Held-out |
|---|---|
| Linear, tuned | |
| Matched-pair, tuned | |
| Degree-2, tuned | |
| TNDM-CP, single fit, rank 5 | |
| TNDM-CP, best of 8, rank 5 | |
| TNDM-CP, ensemble of 8 initializations |
| Estimator | Banzhaf | Shapley |
|---|---|---|
| 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 | – |
| Setting | 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 |
| Coin flip | GSM8K | |||
|---|---|---|---|---|
| Surrogate | ||||
| Linear | ||||
| RBF | ||||
| Coefficient CP | ||||
| TT, random | ||||
| TT, semantic | ||||
| Attribution | Linear | RBF | TT |
|---|---|---|---|
| UCI HAR: 320 fitting, 32 validation; ten splits | |||
| Shapley | |||
| Banzhaf | |||
| Pairwise interaction | |||
| Qwen probe: 512 fitting, 205 validation; three seeds | |||
| Shapley | |||
| Datamodel | LDS | Direct | LDS | Direct |
|---|---|---|---|---|
| Linear | ||||
| Degree-2 ridge | ||||
| Degree-3 ridge | ||||
| RBF | ||||
| TNDM-CP | ||||
| Surrogate | Singleton quantities | Pairwise interactions |
|---|---|---|
| Linear | ||
| Coefficient CP | ||
| TT |
Training Data Attribution
Training data attribution identifies which examples, sources, or groups shape a model’s predictions, capabilities, and failures.
See all 80 papers