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

We thank reviewer uFqn for the valuable feedback. We answer questions as follows and put additional tables in the attached pdf.

W1: ULD inference latency

Although our assistant model involves additional inference cost, the computation can be parallelized. More importantly, our assistant model is only about ¼ of the target LLM, so the computation time of the assistant LLM is fully covered by the target LLM, inducing zero latency.

W2: Data augmentation may require human efforts

Although our method relies on augmented data, we design two principles for the augmentation (as discussed in Section 2.3), which allow our method to easily adapt to various datasets.

• Augment forget data with different forms of the knowledge to be forgotten, which can be as simple as paraphrasing the forget data. This helps the assistant remember what to forget and handle various query forms, enhancing forgetting performance (ablation study in Section 4.3).
• Augment retain data with similar forms of the forget data but with incorrect knowledge. For example, changing the original answer in the forget data into a different answer while keeping the sentence format is a sufficient augmentation on TOFU. This prevents assistant overfitting on forget data and helps it remember only correct answers for forget data query (ablation study in Section 4.3).

W3: ULD performance on other dataset

We follow prior works and apply our method to two additional datasets: RealToxicPrompts [1] that aims to forget toxic knowledge, and WPU [2] that aims to forget personal information.

On RealToxicPrompts, since open-source LLMs have undergone safety alignment and rarely output harmful output (toxic score smaller than 0.1 measured by detoxify library [3]), we first fine-tune Llama-2-7B on 1000 prompts from the RealToxicPrompts dataset with toxic score higher than 0.9, and subsequently unlearn toxic knowledge on 600 prompts. We test the forget performance on the remaining 400 prompts and measure retain performance the same as the HarryPotter experiment in the main paper. The results in the following table show that our method achieves on-par forget performance as baselines while maintaining the best retain performance.

On WPU, we follow the format in TOFU to unlearn on question-answering(QA) pairs of real-world persons. We follow the original paper to report unlearning efficacy (how well the model unlearns), model utility (how well the model preserves remaining knowledge), and response quality (quality of the response on forget data). Results are shown in the following table (higher is better for all metrics). We use the prompt from WPU paper to count the percentage of successful unlearning (Forget unlearn effect), and evaluate the fluency of response on forget query (Forget response quality) by prompting GPT-4o to output a score between 1 to 3 (higher the better). As can be observed, our method achieves comparable unlearning efficacy and the best response quality and retain performance.

Q1: How ULD handle overlapped information in forget/retain data?

In this paper, we mainly consider the setting where forget and retain data have no overlaps. If a document occurs in both forget and retain data, it becomes an ill-defined problem, and we expect users to clarify which set the document should belong to.

Q2: Can ULD be applied not only on LLM?

Our work mainly considers unlearning for LLMs. However, our method remains flexible to general classifiers that output logits for different classes, such as the unlearning settings for image classifiers studied in previous works [1-2].

[1] Liu, et al. "Model sparsity can simplify machine unlearning.

[2] Di, et al. "Label Smoothing Improves Machine Unlearning."

Q3: Scalability of ULD on larger dataset/ more complex LLM

To verify the effectiveness of our method on more complex LLMs, we conduct additional experiments on TOFU-10% using Llama-2-13B LLM. Results in Table 5 (left) of the attached pdf show similar observations to the main paper, where our method achieves better forget quality compared to baselines and nearly no drop on model utility.

To verify the effectiveness of our method on larger datasets, we expand the Harrypotter dataset from 400 text segments to 1800 and report the results in Table 5 (right) in the attached pdf. Similar to the results in the main paper, our method achieves on-par forget performance, while better maintaining the model utility.

Q4: The challenges of data augmentation step

Please refer to our response to W2 for how our method can be generalized to different datasets and settings following two principles for data augmentation.

Dear reviewer,

We would like to follow up on our previous discussion and address any remaining concerns you may have.

It appears that the current rating may not accurately reflect your overall perception of our paper. For instance, your review mentioned that "the method is rigorously evaluated through extensive experiments," and your response to our rebuttal noted that you were "impressed by the additional results." These comments suggest a more positive view of our work, yet the current rating does not seem to align with these comments.

Since the deadline is approaching, we would greatly appreciate it if you could let us know if there are any remaining concerns preventing you from adjusting the score.

We thank reviewer 6JsY for the valuable feedback. We answer questions as follows and put additional tables in pdf.

W1: Minimize CE loss w.r.t. uniform distribution to avoid unbounded loss

We agree that the unbounded loss can be solved by minimizing the CE loss w.r.t. a target distribution, e.g. uniform distribution. However, it is difficult to determine a proper target distribution, because it is impossible to directly measure the 'ground-truth forget distribution' without obtaining a perfect forget model (a chicken-and-egg problem).

Researchers have attempted to find various heuristic target distributions, which still lead to sub-optimal performance. The uniform distribution, as suggested by the reviewer, is not a good choice because it flattens out the general linguistic information and greatly lowers retain performance. [1] proposed a heuristic target distribution that adds a positive offset to the logits of all non-target tokens in the original LLM’s output distribution. However, such a heuristic is sensitive to the choice of offset value. Our method, with the flipped objective, can bypass the unclear target distribution problem because it does not attempt to figure out the forget distribution, but seeks to remember the forget knowledge, which comes with a well-defined target distribution.

To validate our claims, we conduct additional experiments on TOFU, with the previous two approaches included, named Uniform and DI respectively. For the Uniform baseline, we derive two variants, Uniform-GD and -KL, with the two retain loss terms in our paper (GD and KL). The results in Table 1 in attached pdf show that they have worse model utility and forget quality. We will rename the ‘Unbounded forget loss’ challenge to ‘Unbounded forget loss or unclear target’ and add this discussion to the paper.

W2: Flatness of assistant LLM’s output distribution

We acknowledge that on the unseen retain data, the assistant LLM may not output a uniform distribution. However, we will show that the output distribution will be relatively flat, which will not severely alter the behavior of the target model after logit subtraction operation.

First, previous work on uncertainty quantification [2,3] has identified that well-calibrated neural models tend to produce flat output distribution on OOD data. To test whether this applies to our assistant LLM on unseen retain data, we perform an additional experiment to compare the output distribution entropy on forget data, seen retain data, and unseen retain data (wikitext2 passage). The results in Table 2 in pdf confirm that the entropies on seen and unseen retain data are comparable, indicating they are comparably flat.

In addition, we further verify that our logit subtraction does minimal harm to the target LLM’s knowledge. Specifically, we calculate the KL divergence between output distributions of the target LLM before and after the logit subtraction operation. Results in Table 2 in pdf show that the logit difference only induces a large change of the output distribution on forget data and brings minimal changes to the output distribution on seen and unseen retain data. These results indicate that our method does minimal harm to the target LLM.

W3: Hallucinated output for ULD

We agree that hallucination is a challenging issue for existing LLM unlearning methods. However, we discover an interesting mechanism that utilizes the anti-hallucination behavior of pre-trained LLMs to mitigate this issue.

Specifically, we investigate a pre-trained LLM’s internal anti-hallucination behavior (e.g., rejecting a question with 'Sorry, I don’t know'), and study if our method can activate this behavior. Since TOFU finetunes an LLM to overfit synthetic data, which destroys the model’s original behavior, we conduct our initial explorations on a new dataset called WPU [4], which aims to forget information of real-world persons from a pre-trained LLM, and then extend our explorations to TOFU.

We start with a simple strategy where we manually set the first token to be “Sorry” during generation. Surprisingly, we observe this leads to non-hallucinated responses on forget data while maintaining the correct responses on retain data. Table 3 in pdf shows that this simple strategy (termed ULD-SetSorry) increases the number of rejected forget data query while the rejection rate remains low on the retain data. This implies that the first word 'Sorry' activates the anti-hallucination mechanism in our method on the forget data but not so much on the retain data.

Based on this observation, we propose a modification to our method where we add a loss term in assistant LLM training that reduces the probability of 'Sorry' on the first token of each forget data sample. Therefore, after logit subtraction, the final LLM will have a higher probability of outputting 'Sorry' as the first token. Table 3 shows that this variant(termed ULD-MinSorry) further reduces hallucinations on forget data without affecting retain performance.

Finally, we evaluate ULD-MinSorry on TOFU to test whether the remedy works on fine-tuned LLM. Table 4 in pdf shows that the observations can generalize to TOFU without compromising other metrics, although the hallucination reduction is not as significant as on WPU since the behavior of original LLM is destroyed.

In summary, although this is an initial exploration, the results demonstrate that it is promising to adapt our method to reduce hallucinations. We will leave thorough studies to future works.

[1]Dong et al., Unmemorization in large language models via self-distillation and deliberate imagination

[2]Zhang et al., Your Finetuned Large Language Model is Already a Powerful Out-of-distribution Detector

[3]Hou et al., Decomposing Uncertainty for Large Language Models through Input Clarification Ensembling

[4]Liu et al., Revisiting Who’s Harry Potter: Towards Targeted Unlearning from a Causal Intervention Perspective

We thank reviewer Y4ow for the valuable feedback. Regarding the questions:

W1: Experiment result table hard to parse

Thank you for the suggestion. To improve the presentation of experiment results, we include a scatter plot in Figure 1 of the attached pdf to show the main performance comparison of our method and baselines on TOFU dataset (model utility vs forget quality). We will also fix the spelling issues in final revision.

W3: Comparison with task vector and representation fine-tuning methods.

We want to first clarify that our method is fundamentally different from representation learning in WMDP. The basic idea of WMDP is to destroy the model knowledge on forget data by minimizing the distance between the model representation and a random vector. By contrast, our method does not destroy the representation but only offsets the forgetting knowledge in output logits with an assistant model. In fact, the principle of destroying forget knowledge is not a new idea. [1] also share a similar principle that destroys the model’s representations on forget data by minimizing the KL divergence of its output distribution with a distribution where the ground-truth token’s probability is manually reduced. Due to time constraints, we cannot do more experiments to compare with the WMDP method. However, we happen to have made the comparison to [1] in our response to other reviewers. Results in Table 1 of the attached pdf show that this method (DI) has worse forget quality and model utility than our method.

Second, the task vector method is indeed closer to our method. However, their rationale for unlearning is different from ours. They find an unlearning direction in model parameter space, whereas we offset the unlearning knowledge in output logits. Their method may resolve challenge #1 in our paper, but it cannot resolve challenge #2, because simply negating the task vector cannot guarantee that the unseen retain knowledge is not affected. Moreover, most existing works utilizing task vectors only demonstrate its effectiveness on modifying abstract model behaviors (e.g., de-toxic and de-biasing), and few works successfully apply it to unlearn fine-grained factual knowledge [2-4].

To further verify our claims, we compare with the task vector method on TOFU-10%, where we experiment with a wide range of weights for the added vector. Results in the following table show that this method fails to achieve a high forget quality and model utility at the same time, and magnifying the added vector continuously compromises model utility, e.g., model utility drops to 0.36 when the weight is -1.8.

We would love to provide more complete experiments. However, due to time constraints, this is the best we can provide to our best efforts. We are confident that our response can answer your concerns, but if it does not, we would greatly appreciate a timely discussion since the deadline is approaching. Thank you very much for your time!

[1] Dong et al., Unmemorization in large language models via self-distillation and deliberate imagination

[2] Zhang et al., Composing Parameter-Efficient Modules with Arithmetic Operations.

[3] Dige et al., Mitigating Social Biases in Language Models through Unlearning.

[4] Liu et al., Towards Safer Large Language Models through Machine Unlearning.

We thank reviewer F8Xo for the feedback. Regarding the questions:

W1: Assistant model requires additional augmented retain data.

We would like to clarify that the fact that our method requires an augmented retain set does not contradict our claim that our method is insensitive to the representativeness of the retain set (i.e. resolving challenge #2). In the following, we will explain why this is the case and validate our claims with additional experiments.

First, to facilitate our explanation, let us recap the two retain sets required by our method: Regular retain set, which contains samples about the retain knowledge; Augmented retain set, which contains perturbed samples of the forget data.

The augmented retain set serves a very different purpose here. Rather than covering the retain knowledge, the augmented retain set aims to define the boundary of the forget knowledge, which is a specific need of our approach. This is because when the assistant method learns the forget knowledge, it can easily generalize to neighboring knowledge. The augmented retain set essentially tells the assistant model to learn the forget knowledge only, and not the neighboring knowledge. Therefore, representativeness is not a requirement for the augmented retain set. Also, note that generating the augmented retain set does not even need any real retain data. It just contains perturbed versions of the forget data.

On the other hand, it is the regular retain set that needs to represent the retain knowledge, and that would cause challenge #2. Here, we claim that our method is not sensitive to the regular retain set, and can even get rid of it. To validate this, we conduct additional experiments on TOFU, where we keep our original augmented retain set but reduce the size of the regular retain set to 75%, 50%, 25%, and 0% of its original size. Results in the following table indicate that this reduction has minimal impact on the model utility of our method.

Please note that when the regular retain set drops to zero, the method is very close to your expectation that ‘the assistant model is just trained to memorize the data to be forgotten’ – it does not require any real retain data, just memorizing the forget data and distinguishing from the perturbed forget data.

Lastly, we also performed an additional experiment that further confirms that our method performs well on unseen retain data. The results in Table 2 in the attached rebuttal pdf shows that the assistant model has a comparably high entropy (thus flat output distribution) on unseen retain data compared to seen retain data, which indicates that our method is insensitive to the coverage of the seen retain data.

We will add the above discussions and experiments to the paper.

W2: Effectiveness of ULD relies on the augmentation.

We acknowledge that augmented data is crucial for our method. However, as explained in our response to W1, the augmented data should be regarded as an integral part of our method, rather than a plug-in that helps any methods.

In fact, the following table, which is a subset of Table 4 in the main paper, shows that adding the same augmented data does not improve the baselines. These results further suggest it is the logit difference that makes the difference, not the augmented datasets.

Dear Reviewer F8Xo,

As the discussion period is about to end, we wanted to check if our response has sufficiently addressed your concerns. Although your review was submitted close to the deadline, we made our best efforts to promptly provide a detailed response with many additional experiments.

We would greatly appreciate it if you could review our response and consider re-evaluating our paper based on the rebuttal.