Large Language Model Unlearning

In plain words

A language model learns from a huge pile of internet text, so it also learns bad things from it: harmful advice, passages it should not repeat, and false facts. The usual fix is RLHF, which needs many examples of good answers written by people. This paper asks a simpler question: if you only have examples of the bad answers, can you make the model stop giving them?

The method pushes the model away from the bad answers. It does this by reversing the training step that would make it more likely to write them, so the model becomes less likely to produce those exact words. Pushed too hard, though, the model starts producing gibberish on ordinary questions too. Two extra terms guard against that. One trains the model toward unrelated answers, so that "anything but the bad answer" becomes the target. The other keeps its predictions on ordinary text close to those of the original model.

The authors test three cases: removing harmful replies, erasing a copyrighted book the model has memorised, and reducing invented facts. In each case the unlearned model stops giving the bad output, often by giving nonsense or blank text, which the authors treat as acceptable because no good example exists to copy. The method is cheap compared with RLHF, but the outputs are not helpful answers.

TL;DR

• Unlearns the model's answers to a set of bad question-answer pairs by gradient ascent on the answer tokens, and adds a random-answer loss and a KL term §4 eq. 2 eq. 3.
• Tested on harmful answers (PKU-SafeRLHF), a copyrighted novel and hallucinated answers (HaluEval) with OPT-1.3B, OPT-2.7B and Llama 2 7B §6 §7 §8.
• Author-reported: harmful rate 1-6% on the unlearned prompts (original 47-54%) and copyright leak rate 0-1% (original 15-81%) tab. 3 tab. 4.
• The outputs are mostly blank or nonsensical text, not helpful replies; the authors argue this is the best achievable with negative examples only §2 §6.
• Author-reported cost: about 2% of the run time of a full RLHF pipeline on one A100 for OPT-1.3B §9.1 fig. 3.

Problem

RLHF needs human-written good answers as well as bad ones, and it runs several training stages, so it is costly §1.1. Retraining a large model without the offending data is also too expensive §1. The paper asks whether unlearning can serve as an alignment step that uses only negative examples, so that the first goal is to stop undesirable outputs rather than to produce good ones §1 §2.

The authors list three problems that make LLM unlearning harder than classification unlearning. First, the loss on the bad answers is a weak signal: it can rise far above 60 while the model still gives harmful answers §3.2 tab. 1. Second, plain gradient ascent makes normal answers nonsensical §3.2 tab. 2. Third, the format of the normal data affects utility, so the normal data must match the format of the forgotten data §3.2.

Method

Setting. The model being unlearned is the pretrained model theta_o. The forget set D_fgt holds prompt and undesirable answer pairs (x, y). The normal set D_nor holds prompts from a different domain. The goal is that the unlearned model theta_u differs as much as possible from y on x (effectiveness), also on unseen similar prompts (generalisation), that its output on normal prompts stays close to theta_o (utility), and that the cost stays low §2. Positive examples are assumed not to exist §2.

Key choices §3 §4:

• Gradient ascent (GA) on the answer tokens only. The prompt tokens are excluded §3.1 §4 eq. 3.
• Continue unlearning for 3 to 5 times more batches than the point where the forget-set loss first looks high §4.
• A random-mismatch term: for each forget prompt, train the model to predict an unrelated answer drawn from a pool of normal responses §4 eq. 5.
• A forward-KL term that matches the unlearned model's next-token distribution on normal data to the original model's distribution §4 eq. 6.
• Normal data in the same format as the forget data (for example Q&A with Q&A) §3.2 §4.

Update rule. Each step is a gradient step on a weighted sum of the three losses, with non-negative weights eps1, eps2, eps3 §4 eq. 2.

# Unlearning step with gradient ascent and two regularisers  [§4] [eq. 2]
# Inputs: theta (model being unlearned), theta_o (frozen original model),
#         forget_batch (prompts x, undesirable answers y),
#         normal_batch, random_answers (unrelated answers, K per prompt)
# Weights: eps1, eps2, eps3 >= 0, set per experiment [app. B]

def unlearning_step(theta, theta_o, forget_batch, normal_batch):
    # Ascent on the NLL of the undesirable answer tokens only  [§4] [eq. 3]
    loss_forget = -answer_nll(theta, forget_batch)
    # Train each forget prompt toward an unrelated answer  [§4] [eq. 5]
    loss_random = answer_nll(theta, forget_batch.prompts, random_answers)
    # Keep the distribution on normal data close to the original  [§4] [eq. 6]
    loss_normal = forward_kl(theta_o, theta, normal_batch)
    # theta <- theta - eps1 * grad(loss_forget) - eps2 * grad(loss_random)
    #           - eps3 * grad(loss_normal)  [§4] [eq. 2]
    return eps1 * loss_forget + eps2 * loss_random + eps3 * loss_normal

Losses, written in the paper's notation §4:

• Per-pair answer loss: the sum over the answer tokens y_i of the cross-entropy l(h_theta(x, y_<i), y_i) eq. 3.
• Gradient-ascent loss: minus the answer loss, summed over the forget set eq. 4.
• Random-mismatch loss: the answer loss of the forget prompt x paired with each random response in a set Y_rdn, averaged over that set eq. 5.
• KL loss: the sum over the normal sample's tokens of KL(h_theta_o || h_theta) eq. 6. Note that the paper calls this forward KL, which uses the original model as the first distribution.

Random responses. Y_rdn is a set of answers that are unrelated to the forget prompts, which can be taken from the normal dataset §4. The paper does not state how many responses are used, nor whether the set is redrawn per prompt. The code uses K = 5 (see Paper vs code).

Evaluation metrics §5:

• Unlearning efficacy: harmful rate from the PKU moderation model (harm), or leak rate by BLEU against the book text (copyright), or hallucination rate by BERTScore (hallucination) §6 §7 §8.
• Diversity: share of unique tokens in the output §5.
• Fluency: perplexity under the original model; labelled NM when more than 80% of the output repeats one character §5 §6.
• Utility: a reward model score on normal prompts (deberta-v3-large-v2) and BLEURT similarity to the original model's outputs §5 §6.

Hyperparameters

Batch size and learning rate vary by model and method, so each setting is listed in full. "NA" means the term is absent in the paper's table. No setting in the paper states the optimizer, warmup or LoRA rank. The code uses AdamW (see Paper vs code).

Harmfulness app. B.1 tab. 8:

Copyrighted content app. B.2 tab. 9:

Hallucination app. B.3 tab. 10:

Other settings:

• Forget data for harm: harmful answers from the PKU-SafeRLHF Q&A pairs, split by the dataset's original train/test split §6. Normal data: TruthfulQA §6. Copyright: Harry Potter and the Sorcerer's Stone, with BookCorpus as normal data §7. Hallucination: HaluEval wrong answers split 70/10/20 §8.
• Test sets: 200 prompts each for unlearned, unseen and normal in §6 §6. §7 uses 100 per set §7. §8 follows §6 §8.
• Copyright leak threshold: 10% of the mean BLEU of 100K random HP sentences §7.
• Hallucination flag: the answer is hallucinated if its BERTScore to the wrong answer is 10% higher than to the correct answer §8.
• Templated replacement for ablation: "I can't assist it." replaces the random responses; eps weights re-tuned, values not stated §9.2.
• Generation: the copyright setting uses greedy decoding (temperature 0) §7. Harm sampling settings are not stated.
• Hardware: one NVIDIA A100 SXM4 80 GB for the RLHF timing §9.1 fig. 3.
• Seed in code: 8888 unlearn_harm.py. Paper: not stated.

Results

Author-reported, on the harmfulness task with PKU-SafeRLHF, TruthfulQA as normal data, and OPT-1.3B, OPT-2.7B or Llama 2 (7B) §6 tab. 3. Harmful rate is the share flagged by the PKU moderation model. Fluency NM means not meaningful. Utility reward is higher-is-better and similarity is higher-is-better.

Copyrighted content: leak rate is the share of 100 extraction prompts whose BLEU is above threshold §7 tab. 4. The setup is the same as above, with Harry Potter text as the forget set.

Hallucination: rate is the share of answers flagged by the BERTScore rule §8 tab. 5.

Comparison with RLHF, OPT-1.3B, harmfulness §9.1 tab. 6. Full RLHF is SFT, then a reward model, then PPO. Unlearning uses only negative samples.

Templated ablation, OPT-1.3B, harmfulness §9.2 tab. 7: GA+Template gives 1% harmful (unlearned) and 3% (unseen), with diversity 0.864 and 0.816, and utility reward -3.257. The authors report no significant difference from GA+Mismatch.

Run time, OPT-1.3B: unlearning takes about 2% of the RLHF pipeline's time and is comparable to finetuning §9.1. The bar values are only in fig. 3, not in the text.

The authors' own summary: both GA and GA+Mismatch reach near-zero harmful rates on unlearned prompts and generalise to unseen harmful prompts, and the mismatch term helps keep utility §6. Copyright leaks drop to 0-1% §7. Hallucination falls much less, with 9-15% left, which the authors call a harder task §8.

Paper vs code

The official repository is kevinyaobytedance/llm_unlearn at commit 647f309519f91c29d87e62cf63d9a43759810040, under the MIT license (repository LICENSE file). The paper links it in the introduction §1. The owner's handle matches the first author's name, and the repository contains only the harmfulness training script, its utilities and a README; the copyright and hallucination experiments, the finetuning baseline and all evaluation code are not in it. Data must be downloaded: the TruthfulQA CSV from its GitHub page (README). Code citations are at the commit above.

• Defaults match one row. The code defaults (1000 steps, batch size 2, learning rate 2e-6, weights 0.5, 1, 1) equal the OPT-1.3B GA+Mismatch row of tab. 8 unlearn_harm.py. The default model is OPT-1.3B unlearn_harm.py. Other rows must be set by hand.
• Loss normalisation differs. The paper's answer loss sums over the answer tokens eq. 3. The code weights each token by 1 over the answer length, so it averages per answer, and then averages over the batch utils.py. The same eps values therefore carry different weight in the two formulations.
• Optimizer. The paper writes the update as a plain gradient step eq. 2 and does not name an optimizer in the parts read here. The code uses AdamW with a linear learning-rate schedule and no warmup unlearn_harm.py.
• Stop criterion does not fire. The loop runs while bad_loss < max_bad_loss (default 100) unlearn_harm.py. But bad_loss is the ascent loss, which is negative, and the printout negates it unlearn_harm.py unlearn_harm.py. The condition is therefore always true, and training runs for max_unlearn_steps. The paper's stopping rule (continue 3-5 times longer) is not coded as a stop condition either §4.
• Random-mismatch start index. The answer start is set as len(tokenized_question_prefix) utils.py. That tokenizer output is a dictionary, so len returns its number of keys (two: input_ids and attention_mask), not the token count. The forget-set path instead uses the token count minus one utils.py. I have not run the code. If the mismatch loss is computed from index 2, it also trains on prompt tokens, which contradicts eq. 3 and §4.
• Random responses. The code samples K = 5 random answers once per batch and reuses them for every prompt in that batch utils.py. Its pool includes TruthfulQA's "Incorrect Answers" as well as the best and correct answers utils.py. The paper describes Y_rdn as "random (e.g. non-harmful) responses", which does not clearly cover false answers. The paper does not state K.
• KL form and averaging. The code uses the cross-entropy of the original model's probabilities under the unlearned model utils.py. This equals forward KL plus a term that does not depend on the unlearned model, so the gradients match. It averages over every position of the normal batch, including question tokens and any padding utils.py. eq. 6 sums over the answer tokens of each normal sample, and the paper gives no scale for the normal term. The paper also says the KL is "on x^fgt" §4, while eq. 6 uses x^nor.
• Normal data. The code takes only the 70% training split of TruthfulQA and uses the "Best Answer" column utils.py utils.py. The paper names TruthfulQA but not the split.
• Forget data. The code loads the 330k_train split of PKU-SafeRLHF and keeps only unsafe responses unlearn_harm.py utils.py. It has no test split and no evaluation, so the paper's harmful-rate numbers cannot be reproduced from this repository.
• LoRA. The --use_lora flag adds AdaLoRA with r = 32, alpha = 16 on q_proj and v_proj unlearn_harm.py. The paper does not state the rank or the target modules, but reports LoRA for Llama 2 in tab. 8.
• Saving. The code saves a checkpoint every 500 steps and once at the end unlearn_harm.py.

Pitfalls

• Blank or nonsense outputs are the goal. The method removes bad outputs but gives no helpful replacement. The paper states this as intended §2 §6. A deployment that needs a useful reply will get whitespace or gibberish on the same prompts.
• Gradient ascent diverges. Without the extra terms, the model learns to output nonsense on normal prompts too. The table in §3.2 shows nonsense on normal prompts after about 1000 batches. The KL term and the random-mismatch term are needed for utility §3.2 tab. 2.
• Loss value is misleading. The forget-set loss can rise above 60 while harmful answers continue §3.2 tab. 1. Stopping on the loss alone is unsafe. Measure outputs instead.
• Ambiguous references. §6 says "we test plain GA, i.e. eps2 = 0 in Eqn(ans)", but eps2 appears in eq. 2, not eq. 3 §6. Read it as eps2 = 0 in eq. 2.
• Scale of the weights. Because the code normalises differently from eq. 3, tuned weights do not transfer across implementations (see Paper vs code).
• Copyright leak rule. A footnote adds that more than 80% single-character output also counts, but it is unclear whether such outputs count as leaks or are excluded §7. Check before reproducing tab. 4.
• Hallucination threshold. "10% higher" is not said to be relative or absolute §8.
• Judge dependence. The harmful rate uses the PKU moderation model, which was trained on the same forget data. The footnote says it "should have high accuracy" on that data §6. No agreement with human raters is reported. Harmful-rate numbers depend on this classifier.
• Small samples. Each test set has 100 or 200 prompts §6 §7. The table entries are single-run numbers. The paper does not say how many seeds were run.
• Utility can look better after unlearning. GA+Mismatch scores above the original model on utility for OPT-1.3B. The authors attribute this to sampling randomness §6 footnote.
• Cost is for OPT-1.3B only. The 2% time figure is given for OPT-1.3B on one A100 §9.1. The per-model times are in fig. 3 and are not in the text.
• Learning rates differ across rows. Learning rates range from 2e-6 to 2e-4 across models and methods app. B tab. 8 tab. 9 tab. 10. Using one global learning rate will not reproduce the table.
• Not stated: the optimizer, LR warmup, the number of random responses K, the sampling settings for harm generation, the LoRA rank and the number of runs.

Lineage

• Builds on: gradient ascent as an approximate unlearning method for classifiers, cited to Thudi et al. (2022) §1.1. The paper does not give an arXiv id for it.
• Compared against: RLHF with SFT, a reward model and PPO, with the same OPT-1.3B model §9.1 tab. 6. Finetuning on BookCorpus is the baseline in §6 and §7 §6 §7. Concept page: rlhf.
• Related work the paper contrasts itself with: Eldan and Russinovich (2023), arXiv:2310.02238, on Harry Potter unlearning that may produce incorrect answers; and Pawelczyk et al. (2023), arXiv:2310.07579, on in-context unlearning, which uses up space in the prompt, unlike the paper's finetuning approach §1.1.
• Concept pages: machine-unlearning, hallucination-mitigation.
• Superseded by: not stated in the content available to this card.