Token-wise Influential Training Data Retrieval for Large Language Models

TL;DR

This paper proposed RapidIn, a scalable framework adapting to LLMs for estimating the influence of each training data, which efficiently traverses the cached gradients to estimate the influence within minutes, achieving over a 6,326x speedup.

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.

Figures (22)

• Figure 1: Influence estimation for a given generation.
• Figure 2: Overview of the RapidIn framework. a) Caching: 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.
• Figure 3: Workflow of multi-GPU parallelization.
• Figure 4: Token-wise visualization of top-$3$ influential training data for the generation of asking the flu symptoms.
• Figure 5: Prompt template for finetuning and inference.