Unlearning Traces the Influential Training Data of Language Models (2024)

Abstract: Identifying the training datasets that influence a language model’s outputs is essential for minimizing the generation of harmful content and enhancing its performance. Ideally, we can measure the influence of each dataset by removing it from training; however, it is prohibitively expensive to retrain a model multiple times. This paper presents UnTrac: unlearning traces the influence of a training dataset on the model’s performance. UnTrac is extremely simple; each training dataset is unlearned by gradient ascent, and we evaluate how much the model’s predictions change after unlearning. Furthermore, we propose a more scalable approach, UnTrac-Inv, which unlearns a test dataset and evaluates the unlearned model on training datasets. UnTrac-Inv resembles UnTrac, while being efficient for massive training datasets. In the experiments, we examine if our methods can assess the influence of pretraining datasets on generating toxic, biased, and untruthful content. Our methods estimate their influence much more accurately than existing methods while requiring neither excessive memory space nor multiple checkpoints.

## introduction large language models (llms) have had a significant impact on our society. they exhibit remarkable abilities (e.g., chain-of-thought reasoning) without being explicitly trained for such tasks. at the same time, llms also pose potential risks, such as the amplification of discrimination through the propagation of toxic language. llms are trained on a vast number of corpora via pretraining or refined through finetuning on diverse tasks. although some efforts have been made to unravel the black box of llms (e.g., @xcite , it is still unclear which data sources cause their unprecedented abilities and potential harms. ideally, we can answer this question by removing each dataset from the training datasets and as-sessing the change in the model's performance after retraining (leave-dataset-out). however, since we need to retrain a model on each dataset, leavedataset-out is prohibitively expensive. training data attribution overcomes this problem by approximating the influence with hessian-based influence functions (hif; @xcite or tracking changes in test loss during training (tracin; @xcite . however, hif requires a large memory space to approximate the inverse hessian @xcite , while tracin generally needs multiple model checkpoints. in this paper, we propose untrac, which traces the influence of a training dataset by unlearning it from a trained model (figure 1 ). leave-datasetout removes each training dataset and measures its influence by assessing the trained model's performance on a test dataset. analogous to leave-datasetout, untrac unlearns each training dataset and estimates its influence by assessing the unlearned model's performance on a test dataset. unlearning has been studied to eliminate sensitive data from a trained model @xcite @xcite and has recently been applied to llms @xcite . following @xcite , we unlearn a training dataset using gradient ascent, in contrast to the gradient descent normally used in training. interestingly, @xcite argued that influence functions can be regarded as an approximation of the effect of finetuning on a number of examples (e.g., unlearning mislabeled examples). with untrac, instead of using the approximations, we directly quantify the effect of unlearning. when many datasets are used for training, un-trac is computationally costly because unlearning must be run for every individual training dataset. to overcome this drawback, we propose untrac-inv as a scalable approach particularly effective for an increasing number of training datasets. untrac-inv "unlearns" a test dataset instead of training datasets and evaluates the unlearned model on training datasets. untrac-inv requires only a single run of unlearning, and, as we will show, can be considered as an efficient approximation of untrac. in our experiments, we first examine whether our methods can trace influential training tasks in the setting of finetuning. we created a dataset representing a mixture of synthetic tasks, designed to evaluate our method's capability in assessing the influence of each task. in order to make this assessment more challenging, we have created task pairs that, while semantically distinct, require responses in the same format. additionally, we include pairs that are nearly identical in content but demand responses in differing formats. by estimating the influence across these task pairs, we verify that our methods are not overly reliant on superficial similarities between tasks. in this controlled dataset, we show that our methods accurately assess the influence of training tasks, where we use the expensive leave-dataset-out method as the ground-truth, and are only slightly affected by the output format. next, we assess whether our methods can identify the source of harmful content generated by a pretrained language model. using smaller open pretrained transformers (opt-125m; @xcite , the influence of eight pretraining datasets is estimated. we use three test datasets: toxigen @xcite , winobias @xcite , and truthfulqa @xcite , which contain toxic language, biased text, and false answers to various questions, respectively. we calculate the ground-truth influence of each training dataset and evaluate the correlation between the estimated influence and ground-truth influence. we demonstrate that our methods accurately estimate the influence of pretraining datasets, significantly outperforming other influence functions. finally, we investigate how hyperparameters affect the performance of our methods. we found that untrac works robustly as long as we use preconditioned gradient methods with higher learning rates and a sufficient number of training iterations. in contrast, untrac-inv works well for large batch sizes while being relatively sensitive to the learning rate and the number of training steps. ## problem formulation our goal is to estimate the influence of a training dataset on the model's predictions on a test dataset. to formalize this goal, we assume the counterfactual that a model is trained on the mixture of all datasets except for a dataset z: θ -z . the groundtruth influence of the training dataset z on a test dataset z ′ is defined as equation (1) using model θ -z and θ 0 , which is trained on all datasets d. ## experiments we evaluate if our methods can measure the influence of training datasets across different model architectures (encoder-decoder and decoder-only) and training setups (pretraining and fine-tuning). 2 since we would not generally have a validation set (i.e., ground-truth attributions), it is impractical to tune the hyperparameters for each experiment. thus, we set the hyperparameters based on the experimental results on toxigen, one of the test datasets used in section 5.2. the same hyperparameters are used across all the experiments. we use adam (kingma and ba, 2014) with a constant learning rate of 5e-5, β 1 = 0.9, and β 2 = 0.999, while adafactor is used in section 5.1 due to memory constraints. gradient clipping is turned off during unlearning. as for untrac, we set the batch size to 1 and run unlearning for 1 epoch. as discussed in section 3.2, untrac-inv requires a small number of unlearning steps and a large batch size. hence, we set the batch size as 256 (a single batch contains all test examples) and the number of unlearning steps (epochs) as 5. in section 6, we discuss the hyperparameter sensitivity of our methods. the baseline methods are as follows: tracin, graddot & gradcos we compute the gradient w.r.t. all of the model parameters. tracin uses the checkpoints saved for every 128 steps in section 5.1 and 10,000 steps in section 5.2. hif following koh and liang (2017), we approximate the inverse hessian by lissa @xcite , where the number of iterations is set to 10. we also use arnoldi iteration with low-rank eigenvector projection following @xcite . we set the number of arnoldi iterations to @xmath0.1 due to the memory constraints, while @xmath1.2. following the previous studies, the gradients of training datasets are normalized, and 256 training examples are used to approximate the hessian of training loss. ## discussion in this section, we explore how hyperparameters affect the performance of our methods. following the experiment in section 5.2, we compute the influence of the pretraining dataset on toxigen under various hyperparameter settings, where the size of each training dataset is set to be equal. by monitoring the pearson correlation coefficient between the estimated influence and ground truth, we investigate the hyperparameter sensitivity of our methods. ## related work which training data affects the model's prediction? the importance of this question has been rising along with the emergence of llms trained on massive text corpora @xcite @xcite @xcite . hessian-based influence functions @xcite are pioneering work to address this question and are widely used in the context of natural language processing @xcite @xcite . as computing inverse hessian is quite expensive @xcite @xcite , tracin is another promising approach for training data attribution @xcite . however, several studies empirically reported that these methods are unstable for deep neural networks @xcite . while the studies above focus on a single training example, several studies consider the influences of multiple training examples @xcite . @xcite argue that the first-order approximation behind influence functions is unsuitable, as removing multiple examples leads to a significant perturbation of model parameters. @xcite theoretically show that influence functions work well for a limited number of training steps, but are unlikely to be good approximators of leave-one-out influences. therefore, they used influence functions to detect influential training examples of mispredictions. by flipping the label of the influential training examples and training a model on the flipped examples, the mispredictions can be corrected with a few finetuning steps. whereas they use influence functions to detect influential examples to be unlearned, we use unlearning to detect influential examples. sampling-based attribution methods, such as shapley values (ghorbani and zou, 2019), empirical influences @xcite , and datamodels @xcite are also promising approaches for tracing influential training datasets. however, they commonly require thousands of pairs of training subsets and models trained on the subsets to estimate how the combination of training examples affects the model's performance. although @xcite , a recent sampling-based approach, requires significantly fewer computational resources than others, it still requires tens of pairs of training subsets and their corresponding trained models, which can be costly as well as leave-dataset-out. indeed, trak was evaluated only on fine-tuning setups in its original paper. on the other hand, we study approaches that are scalable to large pretraining datasets. unlearning has been studied to remove undesired data from a trained model @xcite @xcite and has recently been used for llms @xcite @xcite @xcite . ## conclusion this paper explored unlearning approaches to trace the influence of training datasets. our approach, untrac, simply unlearns each training dataset and evaluates how the model's performance changes on a test dataset after unlearning. untrac-inv is a more scalable approach, which unlearns a test dataset and evaluates the unlearned model's performance on training datasets. experimental results showed that our methods can trace the influence of training datasets with significantly higher accuracy than existing methods. untrac works robustly if we use preconditioned optimizers with high learning rates and a sufficient number of training steps. in contrast, untrac-inv is relatively sensitive to the choice of hyperparameters, such as batch size, learning rate, and the number of training iterations. as our methods require the same memory footprint as standard training, they can be applied to llms as long as we have enough memory space for training. in future work, we hope that the effectiveness of our methods will be investigated with even more large models, and that our methods will be helpful in revealing the source of llm's emergent abilities, such as chain-of-thought reasoning.