LLM Unlearning with LLM Beliefs

1. Problem Definition

Large Language Models (LLMs) trained on vast corpora risk memorizing sensitive, private, or harmful information. LLM Unlearning aims to remove this specific knowledge from the model parameters without retraining from scratch.

The paper identifies a critical failure mode in prevailing unlearning methods (such as Gradient Ascent, NPO, and WGA): Spurious Unlearning.

• The Squeezing Effect: When standard methods suppress the probability of a specific target response (e.g., a harmful answer) via gradient ascent, the softmax normalization forces the probability mass to redistribute.
• The Consequence: This mass shifts into high-likelihood regions corresponding to semantically similar rephrasings or paraphrases of the target.
• The Illusion: Automated metrics (e.g., ROUGE, Truth Ratio) often report success because the exact target string is gone. However, the model still leaks the underlying knowledge by generating semantically equivalent alternatives. This is termed "spurious unlearning."

2. Methodology: Bootstrapping (BS) Framework

To address the squeezing effect, the authors propose a Bootstrapping (BS) framework. The core insight is that a model's own high-confidence predictions (its "Model Beliefs") represent the high-likelihood regions where probability mass naturally squeezes. Effective unlearning must explicitly suppress these beliefs, not just the original target labels.

The framework is instantiated in two forms:

A. Bootstrapping-Token (BS-T)

• Mechanism: Instead of using a hard one-hot label for the target token, BS-T constructs a soft target that interpolates between the original label and the model's own high-probability token predictions (the top-$k$ belief set).
• Objective: The loss function penalizes both the original target token and the tokens the model currently predicts with high confidence.
• Effect: This prevents probability mass from shifting to semantically proximate tokens by explicitly suppressing the "neighborhood" of the target.
• Formula: The soft target $t_u^i$ is defined as:
  $$t_u^i = \lambda_{BST} \cdot \text{sg}[\pi_\theta(\cdot | x_u, y_{<i}^u)|_{H_k^{(i)}}] + (1 - \lambda_{BST}) \cdot e_{y_u^i}$$
  Where $H_k^{(i)}$ is the set of top-$k$ high-likelihood tokens, and $\text{sg}$ denotes the stop-gradient operator.

B. Bootstrapping-Sequence (BS-S)

• Mechanism: While BS-T handles local token beliefs, BS-S addresses global sequence beliefs. It samples $N$ high-confidence completions (sequences) from the current model for the forget prompts.
• Objective: These sampled sequences are treated as augmented unlearning data. The model is trained to forget both the original target responses and these newly generated high-confidence paraphrases.
• Effect: This ensures that entire harmful trajectories and semantic continuations are removed, preventing the model from simply rephrasing the harmful content in a longer context.
• Flexibility: BS-S can operate in an off-policy manner (sampling once before fine-tuning) or on-policy manner (resampling during training).

3. Theoretical Analysis

The authors provide a theoretical justification using the AKG (Approximate Kernel Gradient) learning dynamics framework:

• Residual Reshaping: They analyze the residual term ($G$) in the gradient update. Standard Gradient Ascent (GA) pushes the target down but creates a "repulsion" that pushes mass into nearby high-likelihood regions.
• BS-T Dynamics: The analysis shows that BS-T reshapes the residual term to distribute the repulsion force across both the target token and its local belief neighborhood. This flattens the probability distribution in the high-likelihood region rather than creating a new peak, effectively mitigating the squeezing effect.
• BS-S Dynamics: Off-policy BS-S is shown to be equivalent to applying BS-T residuals over a weighted set of belief-aligned continuations, aggregating the forgetting pressure across broader semantic neighborhoods.

4. Key Contributions

1. Identification of Spurious Unlearning: The paper rigorously demonstrates that current SOTA methods (NPO, WGA) often fail to remove knowledge, merely shifting it to semantically similar rephrasings due to the "squeezing effect."
2. Bootstrapping Framework: Introduction of a novel framework that utilizes the model's own high-confidence predictions (beliefs) as auxiliary unlearning signals.
3. Dual-Level Implementation: Development of BS-T (token-level) and BS-S (sequence-level) algorithms that are compatible with existing objectives (e.g., NPO, GradDiff).
4. Theoretical Guarantees: Formal analysis showing how bootstrapping alters gradient dynamics to prevent probability mass redistribution into harmful semantic neighborhoods.
5. Robust Evaluation: Introduction of LLM-as-a-Judge (LaaJ) metrics (Naturalness and Similarity) to detect spurious unlearning, which traditional metrics like ROUGE miss.

5. Experimental Results

The methods were evaluated on three major benchmarks: TOFU (fictitious authors), MUSE (copyrighted/sensitive text), and WMDP (harmful knowledge).

• Performance: BS-T and BS-S consistently outperformed state-of-the-art baselines (GradDiff, NPO, RMU, WGA) across all model sizes (Llama 2/3, Zephyr).
  • BS-S achieved the best balance between Forgetting (Memorization score) and Retention (Utility score), often reaching near-random guessing levels on harmful queries while preserving general capabilities.
  • BS-T offered a highly efficient alternative with strong performance, particularly in scenarios with limited compute.
• Mitigation of Squeezing: Probability dynamics plots showed that BS methods monotonically suppress both target and high-likelihood neighbor probabilities, whereas baselines showed a spike in neighbor probabilities (the squeezing effect).
• Qualitative Results: In case studies, baselines like NPO often produced rephrased answers retaining sensitive info (e.g., changing "English" to "a global language"), whereas BS methods generated semantically distinct and safe responses.

6. Significance

This work fundamentally shifts the paradigm of LLM unlearning from target-centric (forgetting specific strings) to belief-centric (forgetting the model's internal confidence in harmful concepts).

• It exposes the limitations of current automated evaluation metrics, urging the community to adopt semantic-aware evaluation (LLM-as-a-Judge).
• It provides a practical, theoretically grounded solution to ensure that unlearning is thorough and not just superficial, which is critical for the safe deployment of open-source LLMs in privacy-sensitive and high-risk domains.
• The code has been merged into OpenUnlearning, facilitating reproducibility and future research.