Unlearning Traces the Influential Training Data of Language Models

Abstract: Identifying the training datasets that influence a LLM'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.

• The paper introduces UnTrac and UnTrac-Inv methods to unlearn influential datasets, efficiently tracking their impact on model predictions.
• It assesses data influence using gradient ascent, achieving higher precision compared to traditional retraining methods on synthetic and pretraining datasets.
• Evaluation on the OPT-125M model demonstrates robustness in tracing harmful content origins while highlighting the need for optimal hyperparameter tuning.

Unlearning Traces the Influential Training Data of LLMs

Introduction

The paper "Unlearning Traces the Influential Training Data of LLMs" addresses the critical challenge of identifying which specific datasets contribute significantly to a LLM's capabilities and potential risks, such as generating biased or toxic content. Traditional methods to measure dataset influence, like the leave-dataset-out, are computationally expensive as they require retraining models multiple times. The paper introduces a novel method, UnTrac, which simplifies the process by applying gradient ascent to unlearn datasets and measuring the resulting change in model predictions.

Methodology

UnTrac and UnTrac-Inv

The core mechanism of UnTrac involves unlearning datasets from a model using gradient ascent, evaluating the degree of change in predictions on a test dataset post-unlearning (Figure 1). This approach directly quantifies data influence without necessitating multiple model retrainings. UnTrac-Inv extends this by unlearning test datasets and assessing the change in predictions on training data, offering a scalable alternative for scenarios with numerous datasets.

Figure 1: Overview of leave-dataset-out vs. proposed methods, UnTrac and UnTrac-Inv.

The methods are grounded in the principle that data influence can be estimated by observing changes in loss functions post-unlearning. UnTrac is computationally intense when scaling to many datasets, prompting the development of UnTrac-Inv for efficiency.

Experimental Evaluation

Synthetic Datasets

Using synthetic datasets, the paper evaluates UnTrac's ability to identify influential training tasks. Results confirm its robustness, accurately distinguishing tasks based on their alignment with test datasets despite variations in output format. This is illustrated through changes in influence estimates over unlearning steps, showing alignment with ground-truth influences (Figure 2).

Figure 2: The influence estimated by UnTrac (top) and UnTrac-Inv (bottom) on the synthetic datasets A (left) and B (right). The line denotes the average across four runs, and the shaded area corresponds to 95\% confidence region.

Pretraining Datasets and Harmful Content

Further tests on pretraining datasets analyze the ability to trace origins of harmful content in LLMs. Using the OPT-125M model, UnTrac and UnTrac-Inv significantly outperform other influence estimation methods, consistently correlating estimated influence with ground-truth across various dataset configurations and content types (Table 3).

Sensitivity and Optimization Considerations

The performance of UnTrac is notably robust across optimizers like RMSProp and Adam, with optimal performance achieved using higher learning rates and suitable training iterations. Conversely, UnTrac-Inv's effectiveness hinges on batch size and learning rate specifications, indicating a need for careful hyperparameter optimization (Figure 3).

Figure 3: Pearson correlation coefficient between the ground-truth influence and the influence estimated by UnTrac (left) and UnTrac-Inv (right) over unlearning epochs. The line denotes the average across four runs, and the shaded area corresponds to 95\% confidence region.

Implications and Future Directions

The practical applicability of UnTrac and UnTrac-Inv extends to efficiently handling large LLMs (e.g., GPT variants) in tracing dataset impacts. Future research directions could explore the methods' applicability across diverse tasks, envisaging a refined understanding of emergent model capabilities and mitigating biases or harmful behaviors.

Conclusion

UnTrac and UnTrac-Inv introduce a pragmatic and scalable approach to identifying dataset influence on LLMs, outperforming existing methods in precision and efficiency. They provide a viable framework for exploring data attribution and unlearning techniques in AI applications, potentially enlightening the opaque decision-making processes of LLMs.