Token-wise Influential Training Data Retrieval for Large Language Models (2024)

Abstract: Given a Large Language Model (LLM) generation, how can we identify which training data led to this generation? In this paper, we proposed RapidIn, a scalable framework adapting to LLMs for estimating the influence of each training data. The proposed framework consists of two stages: caching and retrieval. First, we compress the gradient vectors by over 200,000x, allowing them to be cached on disk or in GPU/CPU memory. Then, given a generation, RapidIn efficiently traverses the cached gradients to estimate the influence within minutes, achieving over a 6,326x speedup. Moreover, RapidIn supports multi-GPU parallelization to substantially accelerate caching and retrieval. Our empirical result confirms the efficiency and effectiveness of RapidIn.

## introduction large language models (llms) have been widely used in various applications across different industries, such as text generation @xcite , translation @xcite , summarization @xcite , and scientific applications @xcite @xcite , due to their unprecedented scale and the impressive capabilities derived from the massive training dataset @xcite . e.g., llama-2 @xcite has up to 70 billion parameters and is trained on 2 trillion tokens of online data. given a model generation, can we determine which training data have the most influence for this generation? understanding how training data influence the content they generate is particularly crucial @xcite . for example, when a risky generation is identified, tracing it back to the most influential training data can help developers filter out risky data and retrain the model @xcite . in addition, knowing the influence of training data for a target generation is highly valuable for machine unlearning @xcite test generation 𝑡 ## influence estimation given a training dataset @xmath0, where s i , the i-th data instance, is a sequence of tokens. s - p i i , • • • , s -1 i represent the tokens of the i-th prompt (instruction); p i is the length of the prompt, and s 0 i , • • • , s g i i denote the tokens of the ith generation for the i-th prompt; g i is the length of the generation. an llm is trained by minimizing: in the training process, for iteration a, assuming we train only one training data s a at an iteration, we update the parameter θ a to θ a+1 for the next iteration. the influence of s a with respect to the test data t should be l(t, θ a )-l(t, θ a+1 ). then for the entire training process, we have n @xmath1, s i ) = l(t, θ 0 ) -l(t, θ), where θ 0 is the initial parameter, and θ is the final parameter after training. for the iteration a, we have the first-order approximation: the gradient descent family of optimizers is commonly employed to train the model, so we have θ a+1 = θ aη a ∇ θa l(s a , θ a ), where η a is the learning rate in the iteration a. in llms, the learning rate is typically small, so we ignore the higher-order term o(||θ a+1θ a || 2 ) in the eq. ( 2 ), which is the order of o(||η a || 2 ). then we have the approximation: for a given training data s k , we estimate the influence of s k with respect to the test generation t by summing up all iterations that are trained by s k : most llms are trained with batch size @xmath2. for a batch b a and s @xmath3, we have in practice, storing the information for each batch size at each iteration is challenging, so we the original gradient of each training data is converted into a small vector rapidgrad of length k (much smaller than the original dimension) that represents the original gradient. these rapidgrads can be very small in size (mbs or even kbs) and cached on disk or in cpu/gpu memory for later retrieval. b) retrieval: for a given test generation t, its gradient vector is converted to a rapidgrad using the same process as in the caching stage. influence can then be efficiently estimated by taking inner products between this rapidgrad and the cached rapidgrads of each training data point. based on the above equations, we can solve the following four influence estimation questions: • i θ (t, s k ) -how training data s k influence the entire sentence of the given generation t. ## influential training data retrieval the straightforward influence calculation for each training data is to directly compute ∇ θ l(t, θ) and ∇ θ l(s k , θ). however, the gradients of llms can be extremely large (e.g., 26gb for llama-2 7b gradients). given a large number of test generations and a massive dataset, this introduces prohibitive computational costs and becomes impractical. readers may wonder if we could cache all the ∇ θ l(s k , θ) for the entire training data, so that we only need to compute ∇ θ l(t, θ) each time and then traverse the cached gradient of each training data to estimate the influence. however, this requires extremely extensive storage space: e.g., a 10tb hard drive can only store 393 full-precision gradients for llama-2 7b. can we compress the gradient in mbs or even kbs for each training data? then for any test generation t, we can efficiently estimate influence by accessing compressed gradients. ## experimental evaluation llm fine-tuning. we evaluate our rapidin using the open sourced llama-2 models @xcite by finetuning llama-2 7b and 70b with qlora adapters @xcite . we also evaluate rapidin on the fullparameter finetuned llama-2 7b for evaluating its scalability. the details of qlora and its implementation are reported in appendix a and b. datasets. we use the alpaca dataset with 52k instruction-following data @xcite , which contains instruction, input, and output for each training data. in all experiments, we merge input into instruction, and response is the "label/ground-truth" for instruction. for performance evaluation of rapidin, we synthesize a poisoned and a hallucinated dataset (section 4.5 and 4.6, respectively). ## related works influence estimation is a technique to estimate influence of each training data for a specific test data. it is a crucial approach for understanding model behaviors and explaining model predictions, and has received increasing research attention recently @xcite @xcite @xcite . influence-based methods. (koh and liang, 2017) apply influence function that requires gradients and hessian-vector products to measure the influential contribution of training data for a test point. however, @xcite found that influence functions in deep learning are fragile, and inaccurate on deeper networks. on the other side, it is prohibitively expensive to compute the hessian-vector products for llms. gradient-based methods. @xcite ) introduce tracin, a first-order gradient-based method to trace the loss change on the test point for computing training data influence, reducing computation overhead. however, it requires more checkpoints for accuracy, which is impractical for llms. tracin uses first-order gradients to estimate influence, but llm gradients can be extremely large, making them difficult to store and leading to slow computations. @xcite present fastif, a scalable influence functions method using k-nearest neighbors to collect candidate points for estimating the inverse hessian-vector product, improving scalability. however, it requires caching each training data's gradient, which is impractical to store due to the large size of llm gradients. contrast-based methods. @xcite ) develop a contrast-based method called contrastive error attribution (cea) for fine-tuned language models, to identify training examples that cause the specific generation. however, it requires a comparison between the model's generation and a humancorrected version, but in many cases, there may be more than one correct answer to a question. ## conclusion in this paper, we propose rapidin, a highly scalable influence estimation framework for llms. we compress the gradient vectors by over 200,000x. rapidin traverses the cached rapidgrad to estimate the influence for the entire training dataset in minutes. rapidin also supports multi-gpu parallelization for improved scalability and usability. the experiments confirm its efficiency and efficacy.