Methods · Training Data Attribution
TRAK (Tracing with the Randomly-projected After Kernel)
Survey sections: §11.2.4 Retraining-Based Influence Estimation: datamodels and TRAK
Summary TRAK (Tracing with the Randomly-projected After Kernel) estimates training-data attribution more cheaply than retraining-based datamodels, by approximating the trained model as a linear model and estimating influence in a randomly projected, lower-dimensional gradient space.
Tracing with the Randomly-projected After Kernel (TRAK; Park et al., 2023) provides a more efficient alternative. TRAK first uses a first-order Taylor expansion around the model parameters to approximate a trained model as a linear model.
Assumptions and requirements
TRAK first uses a first-order Taylor expansion around the model parameters to approximate a trained model as a linear model. It then applies random projections to reduce gradient dimensionality and estimate training-sample influence in this lower-dimensional space. Finally, it averages the estimates across a small number of trained models to reduce variability due to training randomness.
Evaluation considerations
TRAK achieves attribution quality comparable to datamodels at substantially lower computational cost (Park et al., 2023).
Intended applications
Georgiev et al. (2023) proposed TRAK variants to address diffusion models.
Related methods
Key papers
Cited in §11.2.4.
- Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclerc, and Aleksander Madry. Datamodels: Understanding predictions with data and data with predictions. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 9525–9587. PMLR, 17–23 Jul 2022. URL https://proceedings.mlr.press/v162/ilyas22a.html.
- Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry. TRAK: Attributing model behavior at scale. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 27074–27113. PMLR, 23–29 Jul 2023. URL https://proceedings.mlr.press/v202/park23c.html.
Also discussed in: §11.2.7 Recent Extensions and Limitations
Other methods in Training Data Attribution
Verbatim Data Attribution: OLMoTrace · Influence Functions (Gradient-Based Attribution) · Kronecker-Factored Approximate Curvature (K-FAC) · Eigenvalue-Corrected K-FAC (EK-FAC) and ASTRA · TracIn