Scalable GGN Influence Functions for Fine-Tuning and Transfer Learning Settings

Understanding which training examples drive model behaviour has proven increasingly valuable for interpretability, data valuation and safety. A naive approach to data attribution is to retrain models with individual training examples removed, but retraining is computationally prohibitive. Influence functions provide an efficient alternative. However, practical influence estimation requires computing a parameter sized ($d$) Hessian and subsequently an inverse Hessian vector product (iHVP) which has a $\mathcal{O}(d^3)$ time and $\mathcal{O}(d^2)$ memory complexity. This becomes infeasible as model size scales. Structured approximations such as the KFAC and EKFAC reduce the computational and memory complexity, making influence estimates tractable at scale, but lose accuracy. This work proposes a new method for calculating the iHVP based on the more accurate GGN-curvature approximation that enables influence estimation in $\mathcal{O}((Nn)^3 + N^2nd)$ time complexity and $\mathcal{O}((Nn)^2 + d)$ memory; where $N$ is number of training points and $n$ is model output dimension. This is particularly advantageous in fine-tuning and transfer learning settings where $Nn \ll d$. Fine-tuning and transfer learning have emerged as an important paradigm in modern machine learning as large pre-trained models are adapted for downstream tasks by training on smaller datasets. Wall-clock timing experiments on MLP models of multiple sizes support the theoretical claim. Further, the practicality of the method at scale is demonstrated by applying it to a 4 million parameter Efficientnet-B0 image classification model.