Forget to Know, Remember to Use: Context-Aware Unlearning for Large Language Models

TL;DR

This work shows that standard LLM unlearning can unintentionally suppress the model’s ability to use forgotten information when it appears in context. It introduces a context-aware objective by adding a KL-consistency term, $\mathcal{C}(\mathcal{S}_f^{\text{ctx}}, w)$, to preserve contextual utility without compromising forgetting or retain-set utility. Across Gemma-2B-IT and Qwen3-8B on TOFU with varying forget-set sizes, context-aware variants restore contextual QA performance to near-original levels (LLM-Judge ≈ 0.97–0.99) while maintaining strong forgetting and model utility. The approach is modular, stable to $\lambda_c$, and readily integrable with existing unlearning methods, advancing practical unlearning for context-rich applications such as retrieval-augmented generation.

Abstract

Large language models may encode sensitive information or outdated knowledge that needs to be removed, to ensure responsible and compliant model responses. Unlearning has emerged as an efficient alternative to full retraining, aiming to remove specific knowledge while preserving overall model utility. Existing evaluations of unlearning methods focus on (1) the extent of forgetting of the target knowledge (forget set) and (2) maintaining performance on the retain set (i.e., utility). However, these evaluations overlook an important usability aspect: users may still want the model to leverage the removed information if it is re-introduced in the prompt. In a systematic evaluation of six state-of-the-art unlearning methods, we find that they consistently impair such contextual utility. To address this, we augment unlearning objectives with a plug-in term that preserves the model's ability to use forgotten knowledge when it is present in context. Extensive experiments demonstrate that our approach restores contextual utility to near original levels while still maintaining effective forgetting and retain-set utility.

Figures (9)

• Figure 1: Overview of our settings. (\ref{['fig:unlearn']}) Apply unlearning to remove the forget set; (\ref{['fig:directqa']}) Measure forgetting without additional context. (\ref{['fig:contextqa']}) Our new Contextual QA evaluation tests whether the model can still use the (forgotten) knowledge when it is provided explicitly in the context.
• Figure 2: Contextual QA performance across metrics (ROUGE-L, LLM-judge, and utility) for unlearning methods with 5% forget set. Top row shows Gemma-2B-IT and bottom for Qwen3-8B.
• Figure 3: Examples used in context-aware unlearning. Top:$s_f=(q,a)\in \mathcal{S}_f$. Red marks content to forget. Bottom:$s_f^{\text{ctx}}=(q,c)\in \mathcal{S}_f^{\text{ctx}}$. Blue marks desired response (aligned to the frozen original model) given context. Templates and special tokens may vary depending on the specific model and tokenizer.
• Figure 4: Ablation on forget ratio for Gemma-2B-IT. For each ratio (1%, 5%, 10%), we report Direct QA (standard forgetting objective), Contextual QA (our newly defined contextual utility), and overall model utility.
• Figure 5: Template for Contextual QA, where the model is given both the context and the question to answer.