Source-linked AI summary
Towards Safer Large Language Models through Machine Unlearning
Zheyuan Liu, Guangyao Dou, Zhaoxuan Tan, Yijun Tian, Meng Jiang
TL;DR
LLMs can produce harmful outputs, and existing unlearning approaches may impair responses to normal prompts. SKU uses harmful-knowledge acquisition followed by knowledge negation, with experiments reporting reduced harmful outputs without sacrificing normal-prompt response quality.
Problem
LLM unlearning must reduce harmful outputs without degrading performance on normal prompts, a trade-off highlighted for existing approaches.
Method
SKU acquires harmful knowledge using three modules and then negates the resulting knowledge through task-vector negation.
Results
SKU reduces harmful outputs without sacrificing response quality on normal prompts across experiments and ablation studies.
Takeaways & Limitations
The framework addresses the paper's target trade-off between unlearning harmfulness and preserving utility performance.
Takeaways & Limitations
SKU's applicability beyond harmful responses to other Right To Be Forgotten scenarios requires further exploration, and its robustness against jailbreaks is unknown.
Abstract
from arXiv · showhide
The rapid advancement of Large Language Models (LLMs) has demonstrated their vast potential across various domains, attributed to their extensive pretraining knowledge and exceptional generalizability. However, LLMs often encounter challenges in generating harmful content when faced with problematic prompts. To address this problem, existing work attempted to implement a gradient ascent based approach to prevent LLMs from producing harmful output. While these methods can be effective, they frequently impact the model utility in responding to normal prompts. To address this gap, we introduce Selective Knowledge negation Unlearning (SKU), a novel unlearning framework for LLMs, designed to eliminate harmful knowledge while preserving utility on normal prompts. Specifically, SKU is consisted of two stages: harmful knowledge acquisition stage and knowledge negation stage. The first stage aims to identify and acquire harmful knowledge within the model, whereas the second is dedicated to remove this knowledge. SKU selectively isolates and removes harmful knowledge in model parameters, ensuring the model's performance remains robust on normal prompts. Our experiments conducted across various LLM architectures demonstrate that SKU identifies a good balance point between removing harmful information and preserving utility.
1 Introduction
LLMs can generate harmful responses, while existing unlearning approaches often trade harmfulness reduction for normal-prompt utility. SKU addresses this trade-off with two stages that acquire harmful knowledge and then negate it while preserving utility.
- Motivation: LLMs' broad capabilities and training on extensive textual data can produce undesired outputs for harmful prompts.The paper frames safe generation as necessary for alignment with policy regulations and human values.
- Motivation: Machine unlearning is complicated for LLMs because their training corpora are large and normal-prompt performance can deteriorate during unlearning.The paper contrasts this utility loss with the difficulty of removing unwanted behaviors from large-scale training data.
- Motivation: Pretrained LLMs fail to avoid harmful prompts, whereas previous gradient-based approaches struggle to answer normal prompts.Figure 1 is presented as illustrating this opposing performance pattern.
- SKU framework: SKU uses harmful knowledge acquisition followed by knowledge negation to remove harmful knowledge while maintaining performance on normal prompts.The acquisition stage includes guided distortion, random disassociation, and preservation divergence modules.
- SKU framework: The preservation divergence module targets normal prompts by altering performance divergence between the unlearned and pretrained models.The paper identifies this module as the component focused on utility rather than harmful content.
- Evaluation: Experiments and ablations across various LLMs demonstrate effectiveness in unlearning harmfulness while preserving utility performance.The authors present this as evidence that SKU addresses both objectives.
2 Related Work
Related work distinguishes machine unlearning from task-vector methods: unlearning removes information associated with selected data, while task vectors modify pretrained models for specific tasks.
- Machine unlearning: Machine unlearning includes exact and approximate forms, differing in whether the unlearned model must match a completely retrained model exactly.Exact unlearning eliminates information relevant to removed data; approximate unlearning imposes a weaker parameter-level requirement.
- Task vectors: Task vectors are related to weight-interpolation methods and are designed to improve pretrained-model performance on specific tasks.The paper identifies task vectors as a closely related technique to SKU.
3 Preliminary
The preliminary formulation separates harmful data to forget from benign data to retain, while requiring non-harmful responses to unseen harmful prompts.
- Problem formulation: The forget dataset Df contains harmful prompt-response pairs whose information and dangerous responses the original model should erase.The retained normal dataset Dn is defined separately from harmful data.
- Problem formulation: Because LLMs are trained on broad online datasets, harmful prompts in Df need not belong to the model's training dataset.This formulation acknowledges that exhaustively identifying all harmful training information is unrealistic.
- Problem formulation: The unlearned model should generate non-harmful responses for unseen harmful prompts.This extends the objective beyond memorized forget examples to generalization on new harmful inputs.
4 Methods
SKU uses a two-stage framework: it first learns harmful knowledge from multiple perspectives, then selectively negates that knowledge while preserving normal-prompt utility.
- Two-stage framework: SKU first identifies and learns harmful knowledge within the LLM, then systematically negates it in a second stage.The framework is designed to remove harmful information without jeopardizing utility performance.
- Harmful knowledge acquisition: The guided distortion module trains the pretrained model to respond accurately to harmful prompts, exposing it to harmful responses for later unlearning.It uses gradient descent with a cross-entropy loss to internalize harmful-response knowledge.
- Harmful knowledge acquisition: The random disassociation module pairs harmful prompts with unrelated harmful responses to diversify the harmful knowledge learned by the model.Random response sets disrupt direct prompt-response associations and prepare broader knowledge for subsequent unlearning.
- Preservation divergence: The preservation divergence module uses negative KL divergence on normal prompts to diverge the unlearned model from the original model while maintaining normal-prompt effectiveness.The model defines distributions from the original and unlearned LLMs, then integrates this objective with the other module losses using weighted hyperparameters.
- Knowledge negation: In the knowledge negation stage, SKU negates the combined harmful knowledge extracted from a saved model, modifying targeted harmful components while preserving the original model’s learning.The procedure extracts isolated harmful knowledge and applies a negation operation to form the unlearned model.
5 Experiments
The experiments evaluate SKU across LLM scales using harmful and normal prompt datasets, multiple baselines, and separate unlearning and utility metrics. SKU achieves a strong balance between reducing harmful generations and preserving normal-prompt utility.
- The experiments address whether SKU balances unlearning and utility, identifies module contributions, and resolves their trade-off.
- Experiments test OPT-2.7B, LLAMA2-7B, and LLAMA2-13B using harmful PKU-SafeRLHF pairs and normal TruthfulQA prompts.
- The evaluation measures harmful-response rates on unlearned and unseen harmful samples alongside perplexity and semantic utility on normal samples.Harmfulness is assessed with few-shot GPT-4 prompting, while normal-response quality is compared using perplexity and semantic similarity metrics.
- GA usually reduces harmful generation most effectively but sacrifices utility, whereas FT preserves utility while often performing worst on harmful-prompt unlearning.
- SKU achieves leading average rankings by combining low harmful rates with satisfactory utility across model scales.
6 Ablation Study
The ablation study evaluates how SKU’s modules affect harmful-response reduction and normal-prompt utility. Random disassociation strengthens harmful-knowledge removal, while preservation divergence helps retain utility.
- Ablations iteratively remove SKU modules to assess their roles in balancing unlearning efficacy and model utility.
- Random Disassociation Module Removal: Removing random disassociation increases harmful rates from 3% to 25.5% on OPT-2.7B, 28.5% on LLAMA2-7B, and 34.5% on LLAMA2-13B.
- Random Disassociation Module Removal: Without random disassociation, perplexity improves from 24.86 to 22.94 on LLAMA2-7B, but harmful-prompt handling deteriorates substantially.The passage also reports smaller perplexity improvements for OPT-2.7B and LLAMA2-13B, alongside BLEURT improvements across all three models.
- Random Disassociation Module Removal: Random disassociation diversifies acquired harmful knowledge, enabling the negation stage to remove a broader range of harmful content.
- Preservation Divergence Module Removal: Removing preservation divergence raises perplexity from 24.86 to 30.45 on LLAMA2-7B, indicating diminished normal-prompt performance.
7 Unlearning Performance v.s. Utility
The trade-off analysis examines harmful rates and utility as training steps change. SKU and GA with Mismatch reduce harmful rates strongly, while GA-based methods can incur worsening perplexity as unlearning intensifies.
- The analysis compares SKU with baselines because SKU is not individually best on either harmful rate or utility metrics.
- Unlearning Performance Analysis: SKU reduces the harmful rate for unlearned samples from 44% to 3% as training steps increase.
- Unlearning Performance Analysis: GA with Mismatch reduces harmful rates from 47% to 3.5%, while FT changes from 57% to 53% and GA from 5% to 2%.
- Unlearning Performance Analysis: The harmful-rate trend for unseen harmful samples closely resembles the trend for unlearned samples.
- Utility Performance Analysis: As GA and GA with Mismatch reduce harmful rates through 1000 steps, their perplexity increases exponentially, indicating worsening utility.
8 Conclusion
The paper introduces SKU as a two-stage framework that acquires harmful knowledge for identification and then negates it. Its reported results show reduced harmful outputs while maintaining normal-response quality.
- SKU addresses the trade-off between maintaining LLM utility and unlearning harmful knowledge.
- The framework first enhances harmful knowledge for identification, then strategically negates it to mitigate that knowledge.
- Figure 3 compares unlearning harmful rates, perplexity, and BLEURT across training steps, with SKU shown by the orange line.
- SKU reduces harmful outputs without sacrificing response quality on normal prompts.
9 Limitations
The paper reports two main limitations: SKU does not lead every baseline on every metric, and its scope beyond harmful responses to direct prompts and pretrained-knowledge unlearning remains unresolved.
- SKU does not outperform all baselines on every metric individually.The paper describes this as a remaining limitation despite addressing the trade-off between harmfulness reduction and utility preservation.
- SKU targets harmfulness in pretrained LLM knowledge, but its applicability to other general Right To Be Forgotten scenarios requires further exploration.
- The study focuses on harmful responses to direct prompts, leaving SKU's adaptability to adversarial attacks such as jailbreaks unknown.The authors identify jailbreak robustness as future work.
- The evaluation classifies detailed harmful suggestions, irrelevant responses containing harmful suggestions, and responses containing sensitive words as harmful generations.These categories are assessed using GPT-4 few-shot prompting with examples drawn from PKU-SafeRLHF.
A.2 Utility Evaluation
Utility evaluation measures how well unlearned models preserve normal-prompt behavior using perplexity and BLEURT, alongside comparisons with fine-tuning, task-vector, and gradient-ascent baselines. The evaluation uses reference-model perplexity and semantic similarity between unlearned and original outputs.
- Perplexity and BLEURT measure the quality and semantic consistency of outputs on normal prompts.Perplexity uses token prediction probabilities, while BLEURT compares generations from the unlearned and original models.
- Perplexity is calculated from the exponential of averaged negative token log-probabilities conditioned on preceding tokens.
- BLEURT measures semantic similarity between unlearned-model and original-model generations for corresponding normal prompts.The reported score is based on mean pairwise BLEURT across the generated outputs.
- The baselines include fine-tuning on safe responses, naive task vectors, gradient ascent on forget data, and GA+Mismatch with random normal responses and forward KL divergence.
- The evaluation selects 100 prompts each for unlearned harmful, unseen harmful, and normal test sets.A GPT-2 model pretrained on Wiki-103 serves as the perplexity reference model.
- Figure 4 tracks harmful rates, perplexity, and BLEURT across training steps for SKU and baselines on OPT-2.7B.Panels 4a–4b cover unlearned and unseen harmful samples, while panels 4c–4d cover utility metrics.
C Appendix: Additional Experiments
Additional experiments examine SKU's trade-off behavior across OPT-2.7B and LLAMA2 models. Across the tested architectures, SKU follows a consistent training-time trend that balances unlearning and utility performance.
- Figure 5 evaluates SKU and baselines on LLAMA2-7B using harmful-rate and utility metrics across training steps.
- Figure 6 evaluates SKU and baselines on LLAMA2-13B using the same unlearning and utility metrics across training steps.