Research map / Learning theory and algorithms / Prediction and decision losses
Training Data Attribution via Approximate Unrolling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SOURCE, an approximate unrolling-based training data attribution method that segments the training trajectory, treats gradients and Hessians as stationary per segment, and yields an influence-function-like formula computed with EK-FAC.
Problem
Existing training data attribution methods either rely on implicit differentiation (influence functions), which assume convergence to a unique optimum and ignore optimizer bias and multi-stage training, or on unrolling, which is computationally and memory prohibitive at scale.
Approach
Partition the training trajectory into segments, approximate the Hessians, gradients, and learning rates as stationary within each segment, and derive a closed-form, influence-function-like estimator (with a damping term set by segment length and learning rate) that approximates unrolled differentiation without storing all intermediate iterates.
Key results
- SOURCE achieves the highest LDS on most tasks in the single-model setting, outperforming REPSIM, TracIn, TRAK, and IF, and achieves the highest LDS across all tasks except CIFAR-10 ResNet-9.
- On non-converged models (Concrete-N, FashionMNIST-N) and multi-stage training (RotatedMNIST, PACS), SOURCE shows the largest advantage over implicit-differentiation baselines.
- In subset removal evaluation, SOURCE better identifies top influential points that flip test predictions than baselines, with larger gains in settings challenging for implicit-differentiation methods.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SOURCE instead of the baseline TDA methods | TDA methods evaluated on regression, image classification, and text classification tasks · LDS at α = 0.5 | LDS | Mixed |
| Use SOURCE instead of the baseline TDA methods | TDA methods evaluated with non-converged models or multi-stage training · Non-converged and multi-stage tasks | LDS | Helped |
| Use SOURCE instead of the baseline TDA methods to rank training examples for removal | TDA methods in subset-removal counterfactual evaluation · Classification tasks in Figure 5 | Fraction of misclassified test examples | Helped |
| Use SOURCE instead of implicit-differentiation-based methods | TDA methods in subset-removal counterfactual evaluation · Non-converged and multi-stage settings | Fraction of misclassified test examples | Helped |
Surprising findings
- Even the non-ensembled SOURCE typically outperforms ensembled versions of competing methods.
- The LDS drops significantly at alpha = 1 - 1/N, indicating that single-point expected leave-one-out ground truth is very noisy.
Limitations
- Approximations (stationarity and independence of Hessians and gradients across steps) may be inaccurate if statistics change rapidly within a segment.
- Compute cost scales with number of checkpoints (C times influence functions), and experiments are limited to models trained in under 20 minutes.
- Relies on EK-FAC Gauss-Newton approximation; a diagonal Hessian is used for matrix exponential with AdamW.
Experimental setup
Datasets: UCI Concrete UCI Parkinson MNIST FashionMNIST CIFAR-10 GLUE RTE GLUE SST-2 GLUE QNLI WikiText-2 RotatedMNIST PACS
Benchmarks: Linear datamodeling score (LDS) on UCI, MNIST/FashionMNIST, CIFAR-10, GLUE tasks, RotatedMNIST, PACS Subset removal counterfactual evaluation on FashionMNIST, CIFAR-10, RTE, FashionMNIST-N, RotatedMNIST
Metrics: linear datamodeling score (Spearman correlation) fraction of test examples misclassified after subset removal
Baselines: REPSIM TracIn TRAK Influence functions (EK-FAC)
Builds on: influence functions SGD-influence TRAK EK-FAC TracIn
Closest papers by meaning
- Final-Model-Only Data Attribution with a Unifying View of Gradient-Based Methods (NeurIPS 2025 · similarity 0.89)
- LoRIF: Low-Rank Influence Functions for Scalable Training Data Attribution (NeurIPS 2026 · similarity 0.86)
- Distributional Training Data Attribution: What do Influence Functions Sample? (NeurIPS 2025 · similarity 0.86)
- Better Training Data Attribution via Better Inverse Hessian-Vector Products (NeurIPS 2025 · similarity 0.86)
- Training Data Attribution in Diffusion Models via Mirrored Unlearning and Noise-Consistent Skew (NeurIPS 2026 · similarity 0.85)
- quanda: An Interpretability Toolkit for Training Data Attribution Evaluation (NeurIPS 2026 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.