MULTI-STAGE INFLUENCE FUNCTION

Multi-stage training and knowledge transfer from a large-scale pretraining task to various finetuning tasks have revolutionized natural language processing ( NLP ) and computer vision ( CV ) , with state-of-the-art performances constantly being improved . In this paper , we develop a multi-stage influence function score to track predictions from a finetuned model all the way back to the pretraining data . With this score , we can identify the pretraining examples in the pretraining task that contribute most to a prediction in the finetuning task . The proposed multistage influence function generalizes the original influence function for a single model in Koh & Liang ( 2017 ) , thereby enabling influence computation through both pretrained and finetuned models . We study two different scenarios with the pretrained embeddings fixed or updated in the finetuning tasks . We test our proposed method in various experiments to show its effectiveness and potential applications . 1 INTRODUCTION . Multi-stage training has become increasingly important and has achieved state-of-the-art results in many tasks . In NLP applications , it is now a common practice to first learn word embeddings ( e.g. , word2vec ( Mikolov et al. , 2013 ) , GloVe ( Pennington et al. , 2014 ) ) or contextual representations ( e.g. , ELMo ( Peters et al. , 2018 ) , BERT ( Devlin et al. , 2018 ) ) from a large unsupervised corpus , and then refine or finetune the model on supervised end tasks . In computer vision applications , it is common to use a pretrained CNN as a feature extractor and only finetune top-layer networks through training on the end task . Also , it has been demonstrated that pretraining ResNet ( He et al. , 2016 ) with large hashtag data can greatly benefit many end tasks ( Mahajan et al. , 2018 ) . Intuitively , the successes of these multi-stage learning paradigms are due to knowledge transfer from pretraining tasks to the end task . However , current approaches using multi-stage learning are usually based on trial-and-error and many fundamental questions remain unanswered . For example , which part of the pretraining data/task contributes most to the end task ? How can one detect “ false transfer ” where some pretraining data/task could be harmful for the end task ? If a testing point is wrongly predicted by the finetuned model , can we trace back to the problematic examples in the pretraining data ? Answering these questions requires a quantitative measurement of how the data and loss function in the pretraining stage influence the end model , which has not been studied in the past and will be the main focus of this paper . To find the most influential training data responsible for a model ’ s prediction , the influence function was first introduced by Cook & Weisberg ( 1980 ) from a robust statistics point of view . Recently , as large-scale applications become more challenging for influence function computation , Koh & Liang ( 2017 ) proposed to use a first-order approximation to measure the effect of removing one training point on the model ’ s prediction , to overcome computational challenges . More broadly , there are many works using influence functions to investigate the impact of training data on models in various machine learning applications , such as tracing back the origins of bias in the word embeddings generated by GloVe ( Brunet et al. , 2019 ) , and understanding and mitigating disparate impact to improve model fairness ( Wang et al. , 2019 ) . However , all of the existing influence score computation algorithms study the case of single-stage training – where there is only one model with one set of training/prediction data in the training process . To the best of our knowledge , the influence of pretraining data on a subsequent finetuning task and model has not been studied , and it is nontrivial to apply the original influence function in ( Koh & Liang , 2017 ) to this scenario . In this work , we derive the influence function from pretraining data to the end task in multi-stage training . Since the computation involves several expensive Hessian vector products , we also show how to compute the influence function efficiently in large-scale problems . Based on this technique , we show that • in real datasets and experiments across various vision and NLP tasks , predictions using the technique and actual influence for the pretraining data to the finetuned model are highly correlated ( Pearson ’ s r score to be around 0.6 ) . This shows the effectiveness of our proposed technique for computing the influence scores in multi-stage models ; • the influence for the pretraining data to the finetuned model can be split into two parts : the influence of the pretraining data on the pretrained model , and influence of the pretraining data on the finetuned model . Therefore the testing data from the finetuning task will be impacted by changes in the pretraining data , which can be quantified using our proposed technique ; • the influence of the pretraining data on the finetuning task is highly dependent on 1 ) the similarity of two tasks or stages , 2 ) and the number of training data in the finetuning task . Thus our proposed technique provides a novel way to measure how the pretraining data helps or benefits the finetuning data . 2 RELATED WORK . Multi-stage model training that trains models in many stages on different tasks to improve the endtask has been used widely in many machine learning areas , such as in transfer learning ( Ando & Zhang , 2005 ) and zero-shot learning ( Larochelle et al. , 2008 ) . Recently , multi-stage model training has achieved state-of-the-art performance by learning large embeddings or data representation as the pretraining step on a very large pretraining dataset , which is then followed by a finetune step with further training on the end-task . Examples include the recently proposed BERT ( Devlin et al. , 2018 ) , which learns contextual embeddings on a large corpus with the pretraining tasks chosen to be predicting the masked words in a sentence and predicting whether one sentence is after another sentence . This contextual embedding is then used in finetuning tasks , such as a question answering task . ELMo ( Peters et al. , 2018 ) is widely used in multi-stage model training as a sentence feature extractor to benefit the end-task . Similarly , there are some works in computer vision that train an image representation model on a large number of images as the pretraining step , and then use the resulting features to finetune another task , such as particular image classification tasks . For example , ( Mahajan et al. , 2018 ) uses ResNet in the pretraining step , and the finetuning task is based on hashtag data . The rationale for this multi-stage model is that the pretraining task can learn some common or latent representation which could benefit the end task . Another related line of research is on understanding machine learning models . One category of research is to explain predictions with respect to model variables , and trace back the contribution of variables to the prediction . For example , Oramas et al . ( 2019 ) automatically detects internal features in the set of classes in the pretrained model that are relevant to interpreting the prediction , and shows for various vision tasks that the proposed scheme can produce detailed explanations based on the features that are relevant to the targeted classes . Guo et al . ( 2019 ) aims to interpret variable-wise hidden states in an LSTM model to quantify variable importance and variable-wise temporal importance in the model . Closely related research has sought to connect model prediction and training data , and trace back the most influential training data that are most responsible for the model prediction . Among them , the influence function ( Cook & Weisberg , 1980 ; Koh & Liang , 2017 ) , which aims to model the prediction changes when training data is added/removed , has been shown to be effective in many applications . There is a series of work on influence functions , including investigating the influence of a group of data on the prediction ( Koh et al. , 2019 ) , using influence functions to detect bias in word embeddings ( Brunet et al. , 2019 ) , and using it in preventing data poisoning attacks ( Steinhardt et al. , 2017 ) , etc . All of these works only consider a single stage training procedure , and it is not straightforward to apply the existing influence functions to multi-stage models . In this paper , we propose to analyze the influence of pretraining data on predictions in the subsequent finetuned model and end task . 3 ALGORITHMS . In this section , we detail the procedure of multi-stage training , show how to compute the influence score for the multi-stage training , and then discuss how to scale up the computation . 3.1 MULTI-STAGE MODEL TRAINING . Multi-stage models , which train different models in consecutive stages , have been widely used in various ML tasks . Mathematically , letZ be the training set for pretraining task with data size |Z| = m , and X be the training data for the finetuning task with data size |X | = n. In pretraining stage , we assume the parameters of the pretrained network have two parts : the parameters W that are shared with the end task , and the task-specific parameters U that will only be used in the pretraining stage . Note that W could be a word embedding matrix ( e.g. , in word2vec ) or a representation extraction network ( e.g. , Elmo , BERT , ResNet ) , while U is usually the last few layers that corresponds to the pretraining task . After training on the pretraining task , we obtain the optimal parameters W ∗ , U∗ . The pretraining stage can be formulated as Pretrain Stage : W ∗ , U∗ = arg min W , U 1 m ∑ z∈Z g ( z , W , U ) : = arg min W , U G ( W , U ) , ( 1 ) where g ( · ) is the loss function for the pretrain task . In the finetuning stage , the network parameters are W , Θ , where W is shared with the pretraining task and Θ is the rest of the parameters specifically associated with the finetuning task . We will initialize the W part by W ∗ . Let f ( · ) denote the finetuning loss , there are two cases when finetuning the end-task : • Finetuning Case 1 : Fixing embedding parameters W = W ∗ , and only finetune Θ : Θ∗ = arg min Θ 1 n ∑ x∈X f ( x , W ∗ , Θ ) : = arg min Θ F ( W ∗ , Θ ) . ( 2 ) • Finetuning Case 2 : finetune both the embedding parameters W ( initialized from W ∗ ) and Θ . Sometimes updating the embedding parameters W in the finetuning stage is necessary , as the embedding parameters from the pretrained model may not be good enough for the finetuning task . This corresponds to the following formulation : W ∗∗ , Θ∗ = arg min W , Θ 1 n ∑ x∈X f ( x , W , Θ ) : = arg min W , Θ F ( W , Θ ) . ( 3 ) 3.2 INFLUENCE FUNCTION FOR MULTI-STAGE MODELS . We derive the influence function for the multi-stage model to trace the influence of pretraining data on the finetuned model . In Figure 1 we show the task we are interested in solving in this paper . Note that we use the same definition of influence function as ( Koh & Liang , 2017 ) and discuss how to compute it in the multi-stage training scenario . As discussed at the end of Section 3.1 , depending on whether or not we are updating the shared parameters W in the finetuning stage , we will derive the influence functions under two different scenarios .

Summary: The authors derive the influence function of models that are first pre-trained and then fine-tuned. This extends influence functions beyond the standard supervised setting that they have been primarily considered in. To do so, the authors make two methodological contributions: 1) working through the calculus for the pre-training setting and deriving a corresponding efficient algorithm, and 2) adding $L_2$ regularization to approximate the effect of fine-tuning for a limited number of gradient steps.