UNLEARN Efficient Removal of Knowledge in Large Language Models

The gist The UNLEARN method proposes a novel approach to dynamically forgetting specific knowledge in large language models by leveraging subspace methods to identify and specifically target knowledge removal without adversely affecting other knowledge.

How it works

1. Subspace identification involves training a task-dependent matrix for a specified layer while freezing all other layers, which results in a set of matrices that capture an accurate representation for that task. This process is performed sequentially, starting with the first layer and progressing through the entire network. The training objective is maximizing the conditional language modeling objective as defined by equation (1).

2. Subspace discrimination is used to remove information unique to the task of interest while preventing any degradation of other tasks. This is achieved using a variation of the Gram-Schmidt process to orthogonalize subspaces, allowing mutual information to be preserved. The method computes SVk(T′i) = SVk(Ti) − XNj=1 SVk(Ti) · SVj (To) / SVj (To) · SVj (To), where Ti represents the identified subspace to be removed.

3. Task removal is the final step, where SVD reconstruction to reconstitute the modified task matrix, T′i from the singular values of Ti and singular vectors SV (T′i) above is used. The modified task matrix T′i is then subtracted from W′ for each matrix in the LLM to yield Wc' = W′ − T′i.

Performance and Comparison

The results demonstrate that 96% of targeted knowledge can be forgotten while maintaining performance on other knowledge within 2.5% of the original model. When tasks are similar, UNLEARN still achieves nearly 80% forgetting on the task of interest while preserving performance on similar tasks within 10%. This significantly outperforms the state-of-the-art, which achieves similar forgetting but is accompanied by significant degradation on similar tasks.

The dual method LEARN is proposed for targeted knowledge addition, and LEARN can match the fine-tuning accuracy of Low-Rank Adaptation (LoRA) without adversely affecting similar tasks. The UNLEARN/LEARN methodology offers operational flexibility as it applies the same underlying principles to both enhancing and diminishing model performance on specific tasks.

Key Findings and Limitations

Experiments show that when the subspace discrimination method is applied, UNLEARN successfully reduces its performance while the performance on NaturalQuestions is relatively unaffected. Furthermore, the discriminative ability of UNLEARN far outpaces that of the existing state-of-the-art. A limitation noted is when tasks completely overlap, as observed with arithmetic and GSM8K, where it becomes challenging to discriminate between these two tasks. Future work will investigate strategies to efficiently handle discrimination between larger sets of similar tasks and explore the adaption of other layers to enhance the flexibility and performance of UNLEARN.

Conclusion

UNLEARN introduces a novel approach for forgetting selected knowledge in Large Language Models by relying on subspace identification for tasks and subspace discrimination between similar tasks. The method's ability to isolate and remove specific subspaces within the model ensures precise unlearning, making it a valuable tool for managing the complexities of task forgetting. Compared to state-of-the-art methods like Gradient Ascent, UNLEARN offers substantial advantages in terms of generality, efficiency, and precision. The method achieves 96% forgetting on the task of interest while maintaining performance on other tasks within 2.5% of the original model. When similar tasks are considered, UNLEARN achieves nearly 80% forgetting on the task of interest while preserving performance on the similar task within 10% of the original model. The discriminative ability of UNLEARN far outpaces that of the existing state-of-the-art, ensuring targeted unlearning without compromising the performance on related tasks.

Limitations

One limitation is when tasks completely overlap, as observed with arithmetic and GSM8K, where it becomes challenging to discriminate between these two tasks. Another limitation of this paper that will be addressed in future work is to more fully leverage the experimental insights to optimize the efficiency of the UNLEARN method.

Ethics Statement

By allowing models to forget specific tasks, we enhance privacy and security by ensuring that sensitive information can be effectively removed. Further, UNLEARN can promote fairness by removing biased or harmful information. However, there is also a risk that such methods could be misused to intentionally modify important information, leading to biased outputs.

Acknowledgements

This work was supported by CoCoSys, one of seven centers in JUMP 2.0, a Semiconductor Research Corporation (SRC) program sponsored by DARPA.

References

Tom Kwiatkowski aet al. 2019. Natural questions: a benchmark for question answering research.

Shengnan An, Yifei Li, Zeqi Lin, Qian Liu, Bei Chen, Qiang Fu, Weizhu Chen, Nanning Zheng, and JianGuang Lou. 2022. Input-tuning: Adapting unfamiliar inputs to frozen pretrained models.

Lucas Bourtoule, Varun Chandrasekaran, Christopher A. Choquette-Choo, Hengrui Jia, Adelin Travers, Baiwu Zhang, David Lie, and Nicolas Papernot. 2021. Machine unlearning.

Jiaao Chen and Diyi Yang. 2023. Unlearn what you want to forget: Efficient unlearning for llms.

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021. Training verifiers to solve math word problems.

Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. 2023. Qlora: Efficient finetuning of quantized llms.

Ronen Eldan and Mark Russinovich. 2023. Who’s harry potter? approximate unlearning in llms.

Aarohi Srivastava et al. 2023a. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models.

Neel Guha et al. 2023b. Legalbench: A collaboratively built benchmark for measuring legal reasoning in large language models.

Percy Liang et al. 2023c. Holistic evaluation of language models.

Chongyu Fan, Jiancheng Liu, Yihua Zhang, Eric Wong, Dennis Wei, and Sijia Liu. 2024. Salun: Empowering machine unlearning via gradient-based weight saliency in both image classification and generation.

1. textbfSubspace-based Knowledge Removal for Privacy Compliance (UNLEARN): The system can achieve 96% forgetting on the task of interest while maintaining performance on other knowledge within 2.5% of the original model, directly addressing Right to be Forgotten laws and GDPR requirements by selectively targeting specific knowledge without causing adversely affecting other knowledge in the LLM.

2. textbfTargeted Knowledge Addition (LEARN): The system can match the fine-tuning accuracy of Low-Rank Adaptation (LoRA) without adversely affecting similar tasks, allowing for efficient, task-specific model enhancement by using the same operators to either enhance or reduce the model’s performance for a given task.

3. textbfPrecise Task Discrimination: The system can quantify task similarity and perform subspace discrimination to ensure that when removing one knowledge subspace, it does not impact similar tasks, which is crucial when dealing with overlapping features (similar weight subspaces).

4. textbfAdaptive Rank Optimization: The system can determine the optimal rank for task matrices by finding where performance is not hindered for k values above 4 as seen in Table 2, allowing for a significantly reduced rank that minimizes performance loss while optimizing computational efficiency.

5. textbfversatile Framework Flexibility: The system offers a generalized framework that can be applied to full fine-tuning or any PEFT method, simplifying model management by using the same set of matrices regardless of algorithm for both enhancement and diminishment operations.