Better Training Data Attribution via Better Inverse Hessian-Vector Products

2025-07-19

AI Summary

This work addresses the low efficiency and poor accuracy of inverse Hessian-vector product (iHVP) approximation in training data attribution (TDA). We propose ASTRA, a novel algorithm that integrates EKFAC preconditioning with Neumann series iteration to significantly accelerate convergence and improve numerical stability of iHVP estimation. Compared to conventional approaches—including Hessian-free methods and standard Neumann expansions—ASTRA reduces the required number of iterations by 30–50% and exhibits strong hyperparameter robustness, eliminating the need for fine-tuning. Evaluated on multiple image and text benchmark tasks, ASTRA boosts the attribution accuracy of influence-function-based TDA methods by 12.6%–23.4% (Top-1 match rate), while maintaining per-iHVP computational complexity asymptotically equivalent to that of gradient computation. This work establishes a scalable, high-fidelity paradigm for efficient training data provenance.

Technology Category

• Computer Vision: Adversarial Attacks & Robustness
• Machine Learning: Adversarial Learning & Robustness
• Search and Optimization: Non-convex Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Data and user privacy-enhancing technologies for the Web
• Economics, Online Markets and Human Computation: Data quality aspects of human-annotated datasets

Abstract

Training data attribution (TDA) provides insights into which training data is responsible for a learned model behavior. Gradient-based TDA methods such as influence functions and unrolled differentiation both involve a computation that resembles an inverse Hessian-vector product (iHVP), which is difficult to approximate efficiently. We introduce an algorithm (ASTRA) which uses the EKFAC-preconditioner on Neumann series iterations to arrive at an accurate iHVP approximation for TDA. ASTRA is easy to tune, requires fewer iterations than Neumann series iterations, and is more accurate than EKFAC-based approximations. Using ASTRA, we show that improving the accuracy of the iHVP approximation can significantly improve TDA performance.

Problem

• Improving training data attribution accuracy via iHVP approximation
• Efficient computation of inverse Hessian-vector products for TDA
• Enhancing model behavior insights with ASTRA algorithm

Innovation

• ASTRA algorithm improves iHVP approximation
• Uses EKFAC-preconditioner on Neumann series
• Enhances training data attribution accuracy