Training Data Attribution via Approximate Unrolled Differentiation

TL;DR

This paper introduces Source, a scalable training data attribution method that combines influence functions and unrolling techniques, improving counterfactual prediction especially in complex training scenarios.

Contribution

The paper proposes Source, an approximate unrolling-based TDA method that balances efficiency and accuracy, addressing limitations of existing influence function and unrolling approaches.

Findings

Source outperforms existing TDA methods in counterfactual prediction.

Source is effective in non-converged models and multi-stage training pipelines.

The method combines benefits of implicit differentiation and unrolling techniques.

Abstract

Many training data attribution (TDA) methods aim to estimate how a model's behavior would change if one or more data points were removed from the training set. Methods based on implicit differentiation, such as influence functions, can be made computationally efficient, but fail to account for underspecification, the implicit bias of the optimization algorithm, or multi-stage training pipelines. By contrast, methods based on unrolling address these issues but face scalability challenges. In this work, we connect the implicit-differentiation-based and unrolling-based approaches and combine their benefits by introducing Source, an approximate unrolling-based TDA method that is computed using an influence-function-like formula. While being computationally efficient compared to unrolling-based approaches, Source is suitable in cases where implicit-differentiation-based approaches struggle, such as in non-converged models and multi-stage training pipelines. Empirically, Source outperforms existing TDA techniques in counterfactual prediction, especially in settings where implicit-differentiation-based approaches fall short.