Reversing the Forget-Retain Objectives: An Efficient LLM Unlearning Framework from Logit Difference

Jiabao Ji; Yujian Liu; Yang Zhang; Gaowen Liu; Ramana Rao Kompella; Sijia Liu; Shiyu Chang

Source-linked AI summary

Reversing the Forget-Retain Objectives: An Efficient LLM Unlearning Framework from Logit Difference

Jiabao Ji, Yujian Liu, Yang Zhang, Gaowen Liu, Ramana Rao Kompella, Sijia Liu, Shiyu Chang

arXiv:2406.08607v1Jun 12, 2024cs.CLcs.AI

TL;DR

Conventional LLM unlearning seeks to remove knowledge from specified documents while preserving other knowledge, but loss-based objectives face degenerated outputs and catastrophic forgetting. ULD trains an assistant LLM with reversed objectives and subtracts its logits from the target model’s; on TOFU it loses 0% of utility while reducing training time by more than threefold.

Problem
Conventional unlearning seeks to forget knowledge in specified documents while retaining other knowledge, but loss-based methods face degenerated outputs and catastrophic forgetting.
Source: §1 · pp. 1–2
Method
ULD trains an assistant LLM to remember forget documents and forget retain knowledge, then subtracts its logits from the target LLM’s logits.
Source: §1 · p. 2
Results
On TOFU, ULD loses 0% of model utility, while the most competitive baseline sacrifices 17% on average; training time falls by more than threefold.
Source: §1 · p. 3
Takeaways & Limitations
ULD’s reversed objectives are reported to resolve degeneration and catastrophic forgetting while improving training efficiency and preserving model utility.
Source: §6 · p. 10
ULD uses an assistant LLM at inference, which may increase latency, and appropriate forget-data augmentation can be challenging to create for different datasets.
Source: §9 · p. 10

Abstract

from arXiv · show

As Large Language Models (LLMs) demonstrate extensive capability in learning from documents, LLM unlearning becomes an increasingly important research area to address concerns of LLMs in terms of privacy, copyright, etc. A conventional LLM unlearning task typically involves two goals: (1) The target LLM should forget the knowledge in the specified forget documents, and (2) it should retain the other knowledge that the LLM possesses, for which we assume access to a small number of retain documents. To achieve both goals, a mainstream class of LLM unlearning methods introduces an optimization framework with a combination of two objectives - maximizing the prediction loss on the forget documents while minimizing that on the retain documents, which suffers from two challenges, degenerated output and catastrophic forgetting. In this paper, we propose a novel unlearning framework called Unlearning from Logit Difference (ULD), which introduces an assistant LLM that aims to achieve the opposite of the unlearning goals: remembering the forget documents and forgetting the retain knowledge. ULD then derives the unlearned LLM by computing the logit difference between the target and the assistant LLMs. We show that such reversed objectives would naturally resolve both aforementioned challenges while significantly improving the training efficiency. Extensive experiments demonstrate that our method efficiently achieves the intended forgetting while preserving the LLM's overall capabilities, reducing training time by more than threefold. Notably, our method loses 0% of model utility on the ToFU benchmark, whereas baseline methods may sacrifice 17% of utility on average to achieve comparable forget quality. Our code will be publicly available at https://github.com/UCSB-NLP-Chang/ULD.

1 Introduction

LLM unlearning must remove knowledge from specified documents without damaging other capabilities, but conventional loss-based methods can cause degenerated outputs and catastrophic forgetting. ULD reverses the objectives using an assistant LLM, with reported gains in utility and training efficiency.

1 Introduction

The conventional task aims to forget knowledge in specified documents while preserving other knowledge and task abilities, often using a small set of retain documents.The retain documents are assumed to represent the broader knowledge that should remain.Source: §1 · pp. 1–2
Table 1 tracks responses across epochs 1, 5, and 10; it contrasts gradient-ascent degeneration and forgetting with ULD’s sensible outputs and correct retain answers.Its examples include retain questions covered and not covered by the retain set.Source: Table 1 · p. 2, §1 · pp. 2–3
Conventional objectives risk degenerated outputs from unbounded forget loss and catastrophic forgetting when a small retain set misses knowledge to preserve.The paper notes that prior methods may rely on fragile early stopping to select a suitable checkpoint.Source: §1 · p. 2
ULD trains an assistant LLM to remember forget documents and forget retain knowledge, then subtracts its output logits from the target LLM’s logits.The reversed objectives approach the task through the assistant rather than directly optimizing the target model.Source: §1 · p. 2
ULD loses 0% of model utility, while the most competitive baseline sacrifices 17% on average; training time falls by more than threefold.The comparison is reported on the TOFU benchmark for utility and against the most competitive baseline for training time.Source: §1 · p. 3

3 Experiment

The experiments compare ULD with baseline unlearning methods on fictional-writer knowledge in TOFU and Harry Potter book content. Across both settings, ULD combines effective forgetting with stronger retention of utility than the reported baselines.

3 Experiment

The experiments compare ULD with baselines on fictional-writer forgetting and Harry Potter copyright-content forgetting, followed by analyses of stability, efficiency, and data usage.Source: §3 · p. 5

3.1 Baseline Unlearn Objectives

Baseline objectives combine a forget loss—GA, DPO, or NPO—with a retain loss—GD or KL; Offset variants use an assistant but retain conventional objectives.GD trains correct prediction on retain data, while KL encourages predictions to stay close to those before unlearning.Source: §3.1 · pp. 5–6

3.2 Experiments on TOFU

ULD achieves the best forget performance across all TOFU settings, including 0.99 forget quality on TOFU-1%, while NPO sacrifices 17% of model utility on average.ULD also preserves retain performance best in every setting, with almost no model-utility reduction.Source: Table 2 · p. 6, §3.2 · p. 7
TOFU evaluates unlearning of fictitious writers across 1%, 5%, and 10% forget-data settings, using forget quality and ROUGE-L alongside model utility on held-out retain data.The target is fine-tuned Llama2-chat-7B containing knowledge of all 200 writers.Source: §3.2 · p. 6

3.3 Experiments on HarryPotter

The HarryPotter setting uses 400 book chunks as forget data and 400 C4 paragraphs as retain data, measuring text overlap and performance on six benchmarks plus WikiText perplexity.The target is Mistral-7B-instruct, fine-tuned on the forget data for one epoch to simulate incorrect pretraining on copyrighted text.Source: §3.3 · p. 7
On HarryPotter, ULD has almost no retain-performance reduction and lower BLEU and ROUGE scores than the pre-fine-tuned model; NPO+GD has over 5% accuracy decline.Retain performance includes held-out-text perplexity and accuracy on six standard benchmarks.Source: Table 3 · p. 7, §3.3 · p. 7

8 Broad Impacts

ULD uses an assistant LLM and logit subtraction to build an unlearned LLM without changing the original model’s parameters, aiming to address capability-retention and efficiency challenges.

ULD builds the unlearned LLM by subtracting logits with help from an assistant LLM, without changing the original LLM’s parameters.Source: §8 · p. 10
The framework is described as addressing conventional unlearning’s potential harms to retained knowledge while improving unlearning efficiency.Source: §8 · p. 10
The authors identify potential applications beyond unlearning, including sentiment-controlled generation, knowledge editing, and factuality improvement.They plan to explore these applications, drawing a parallel to prior LLM decoding work [13] [14].Source: §8 · p. 10

9 Limitations

ULD’s limitations include possible inference-latency increases and the difficulty of constructing suitable forget-data augmentations. The appendix reviews conventional forget and retain losses, whose trade-offs include degeneration and catastrophic forgetting.

9 Limitations

ULD uses an assistant LLM at inference, which may increase latency, and suitable forget-data augmentation can be challenging to construct across datasets.The authors note that parallel computation may mitigate latency and identify automatic construction of optimal forget-data augmentation as future work.Source: §9 · p. 10

A Baseline details

The appendix introduces its baseline review by summarizing conventional unlearning objectives before discussing the offset baseline.The forget-loss review covers gradient ascent, DPO, and NPO; the retain-loss review covers gradient descent and KL divergence.Source: A Baseline details · p. 14, §A.1 · p. 14, §A.2 · pp. 14–15

A.1 Forget losses

Gradient-ascent and DPO losses reduce the likelihood of correct forget responses, but gradient ascent can diverge while DPO may collapse responses to alternatives even on retained data.DPO also requires a fixed dataset of alternative responses; NPO omits that alternative-answer term and still cannot avoid degeneration.Source: §A.1 · p. 14

A.2 Retain losses

Gradient-descent and KL losses are used to preserve retain performance, but limited retain data cannot cover all knowledge and thus cannot prevent catastrophic forgetting.KL loss aims to maintain the original predictions, while gradient descent trains on retain data.Source: §A.2 · pp. 14–15

A.3 Offset unlearning

The Offset baseline combines target-model logits with the change in assistant-model logits from fine-tuning. It trains the assistant using conventional unlearning objectives, which the authors say do not avoid degeneration or catastrophic forgetting.

A.3 Offset unlearning

Offset defines unlearned logits as target logits plus a weighted difference between fine-tuned and pretrained assistant logits.The assistant must share the target model’s vocabulary; the experiments use pretrained Llama-2-chat for TOFU and Mistral for HarryPotter.Source: §A.3 · p. 15
The Offset baseline applies conventional unlearning objectives to the combined logits when fine-tuning its assistant LLM.The formulation uses parameters for the target LLM, fine-tuned assistant, and pretrained assistant.Source: §A.3 · p. 15
The authors state that Offset’s assistant-training objective cannot avoid degeneration or catastrophic forgetting, and report experimental examples of both issues.Source: §A.3 · p. 15

B.1 Details of data augmentation

The appendix details dataset augmentation, evaluation measures, hyperparameters, and hardware. It also reports example outputs in which conventional baselines show degeneration or catastrophic forgetting, while ULD avoids those issues.

B.1 Details of data augmentation

TOFU forget QA pairs receive two question-and-answer paraphrases and two perturbed answers, while HarryPotter forget chunks receive two paraphrases.Perturbed TOFU answers provide similar-form false knowledge for augmented retain data to discourage the assistant from overfitting to original answers.Source: §B.1 · pp. 15–16

B.2 Data statistics

Table 5 summarizes forget, retain, and augmented-data sizes for the TOFU and HarryPotter settings.Source: §B.2 · p. 16, Table 5 · p. 16

B.3 Details of metrics

TOFU utility aggregates nine metrics across three retain-data groups, while forget quality compares truth-ratio distributions with a retrained model; HarryPotter forgetting uses BLEU and ROUGE, and retention uses six benchmark accuracies plus WikiText perplexity.TOFU utility uses ROUGE-L, accuracy, and truth ratio across fictional-writer, real-world-writer, and general-world-fact data.Source: §B.3 · pp. 16–17

B.4 Hyper-parameters of baseline methods

The reimplemented baselines share batch size 32, learning rate 1e-5, weight decay 0.01, and retain weight 1; Offset uses Llama-2-chat for TOFU and Mistral-7B-instruct for HarryPotter.The authors reimplement TOFU and NPO and report lacking an official Offset implementation.Source: §B.4 · p. 17

B.5 Hyper-parameters of ULD

ULD uses LoRA r=32 and α=32, assistant construction K=8, retain weight 6.5, and logit-subtraction weights β=0.75 for TOFU and β=0.5 for HarryPotter.A logit filter rate of 1e-2 is used for both experiments.Source: §B.5 · p. 17

B.6 Hardware configuration

Experiments use two A100-80G GPUs except Offset, which uses four to avoid out-of-memory errors.Source: §B.6 · p. 17

C.1 Generation examples on TOFU dataset

In TOFU generation examples, GA and NPO show degeneration and DPO shows catastrophic forgetting, whereas ULD avoids these issues.Table 6 marks successful forgetting in green and degeneration or catastrophic forgetting in red.Source: §C.1 · p. 18

C.2 Generation examples on HarryPotter dataset

The appendix provides HarryPotter response examples and additional training-stability results. It describes a cross-entropy-loss figure and a utility-versus-forget-quality trajectory, but the supplied passages do not state their plotted outcomes.

C.2 Generation examples on HarryPotter dataset

Table 7 presents HarryPotter forget- and retain-query responses, marking successful forgetting in green and degeneration or catastrophic forgetting in red.Source: §C.2 · p. 18, Table 7 · p. 20

D Additional experiment results

The additional-results section includes full training-stability results and a data-usage ablation.Source: D Additional experiment results · p. 21

D.1 Additional result of training stability analysis

Figure 10 tracks cross-entropy loss on forget data and retain data not covered by retain documents; Figure 11 plots utility against log forget quality with marker size indicating epoch.The text states that conventional objectives face degeneration, but gives no comparative outcome for the plotted trajectories.Source: §D.1 · p. 21, Figure 10 · p. 21, Figure 11 · p. 21

D.2 Additional result of data usage ablation analysis

This section presents ToFU-10% performance comparisons across unlearning methods and forget/retain data configurations. The supplied passages identify the comparison scope but do not report results.

The ToFU-10% comparison includes configurations with and without augmented forget/retain data.Source: §D.2 · p. 21
Table 8 compares the performance of different unlearning methods on ToFU-10%.Source: Table 8 · p. 21
Table 8 covers different forget/retain data configurations for the unlearning methods.Source: Table 8 · p. 21