Wisdom is Knowing What not to Say: Hallucination-Free LLMs Unlearning via Attention Shifting

AI summary

Overview

• Research area: Machine unlearning for large language models (LLMs), privacy-preserving natural language processing, and hallucination control during knowledge removal.
• Technical level: Intermediate. Readers need familiarity with transformer attention, softmax normalization, and KL-divergence losses to follow the mechanism, but the framing is accessible.
• Scope: The paper introduces Attention-Shifting (AS), a selective unlearning framework that suppresses attention to fact-bearing tokens in a target dataset while reinforcing attention to semantically important tokens in retained data, and evaluates it on the ToFU and TDEC benchmarks with LLaMA-2 and GPT-NEO models.

What This Paper Is About

Existing LLM unlearning methods sit in a dilemma: aggressive methods such as Gradient Ascent forcibly erase target knowledge but damage the model's performance on neighbouring and general knowledge, while conservative methods such as logits manipulation preserve utility but leave latent knowledge accessible and can produce confident factual hallucinations. The paper proposes Attention-Shifting, which reallocates attention internally rather than rewriting outputs, aiming to achieve unlearning through omission (refusal or minimal responses) rather than through substitution with alternative or fabricated answers.

Key Contributions

1. An Attention-Shifting (AS) unlearning framework built on two attention-level interventions: importance-aware suppression applied to the unlearning set to reduce reliance on memorized knowledge, and attention-guided retention enhancement applied to the retained dataset to reinforce attention toward semantically essential tokens and mitigate unintended degradation.
2. A dual-loss objective that jointly optimizes unlearning and retention via a mixing coefficient α ∈ [0, 1], where the attention suppression (ASP) loss penalizes attention to high-importance tokens in D_t and the attention reinforcement (AKL) loss promotes attention to semantically important tokens in the sub-remaining dataset D'_r ∈ D / D_t. The authors describe this as forming a "soft boundary" that localizes unlearning while stabilizing unrelated knowledge under representation superposition.
3. A lightweight implementation in which suppression and reinforcement signals are derived from reference attention maps of the original model and injected via adapters of approximately 12M parameters in attention modules, rather than by modifying logits or replacing outputs.
4. Empirical validation across two benchmarks and multiple model families, reporting up to 15% higher accuracy on the ToFU benchmark and 10% on the TDEC benchmark relative to state-of-the-art unlearning baselines, while maintaining what the authors describe as hallucination-free unlearning effectiveness.

Main Findings

• Model utility preservation: On ToFU, AS achieves 0.73 ROUGE-L on neighbouring knowledge (+0.05) and 0.76 accuracy on neighbouring knowledge (+0.06), with general knowledge accuracy of 0.80 (−0.03) relative to the original model. The authors state this is the best preservation of model utility among the compared methods, which they define as the average accuracy of NEK and GEK performance.
• Unlearning effectiveness on ToFU: AS records 0.16 ROUGE-L on the target unlearning dataset and a Top-5 exclusion rate (TR, k = 5) of 0.97, which is near the GA result of 0.98 and substantially higher than the logits-manipulation methods IHL (0.62) and ULD (0.47). The authors note AS's Forget Quality (FQ) of 0.17 is lower than IHL (0.68) and ULD (0.89), which they attribute to their behaviour-oriented design: refusals and minimal responses diverge from retrained-model behaviour by design.
• Hallucination suppression: AS achieves 0% reproduction and 0% hallucination, which the authors report as complete knowledge unlearning without misleading generations. In contrast, GA and NPO tend to reintroduce partial knowledge when combined with gradient-based utility losses, and IHL and ULD tend to rephrase or replace target content with semantically similar but incorrect alternatives.
• Robustness to adversarial prompting: Under rephrased, perturbed, and noised query variants, AS maintains a Top-50 exclusion rate above 92% (0.92 rephrased, 0.94 perturbed, 0.93 noised, 0.93 original) with hallucination rates between 0.00 and 0.07.
• Localized unlearning under overlap: When target and retained sets partially overlap, AS stops reproducing the target fact (output appears garbled) while retaining both same-area knowledge (topically related) and same-author knowledge (other attributes of the same entity).
• Sensitivity to retained data volume: Figure 3 shows that GA, NPO, IHL, and ULD degrade noticeably on both NEK and GEK with limited retention (e.g., 40 samples), whereas AS maintains high performance and shows consistent improvement on both NEK and GEK regardless of whether NEK or GEK data is used for retention.
• Unlearning cost across sample counts: In Table 2, AS on Forget-01 reaches 0.72 neighbouring performance in 24 epochs, versus GA+GD at 0.62 in 56 epochs, NPO+GD at 0.41 in 23 epochs, and IHL+GD at 0.70 in 51 epochs. On Forget-10, AS records 0.72 in 20 epochs while GA+GD records 0.59 and NPO+GD records 0.38.
• Ablation on the retention loss: Using ASP alone does not cause catastrophic unlearning, but replacing AKL with CE- or KL-based retention strategies fails to improve model utility consistently, which the authors take as evidence that AKL is essential to their attention-manipulated unlearning.
• TDEC results: Evaluated on GPT-NEO models at 125M, 1.3B, and 2.7B with 32 unlearning samples and the threshold el₁₀ ≤ 0.05, AS occupies the upper-right corner of the utility-versus-unlearning plots, indicating both high neighbouring accuracy and strong model utility relative to GA, NPO, and IHL. Under multiple continue-unlearning requests (4 samples), AS shows the most stable performance across iterations.

Methodology in Plain English

The approach starts from the observation that LLMs assign high attention to semantically important tokens such as "father" or "hairdresser" when recalling memorized facts. The authors quantify token importance as the change in predictive entropy when a token is masked: masking a truly important token changes the model's certainty more than masking a function word.

Using that importance signal, they train small adapter modules inserted into the attention mechanism, leaving the base model's weights untouched. For samples in the unlearning set, the adapters are trained to reduce attention weights assigned to high-importance fact-bearing tokens, with a parameter λ ∈ [0, 1] controlling suppression strength and a normalization step keeping the adjusted attention weights a valid probability distribution summing to 1. The training objective is the KL divergence between the model's current attention distribution and this target suppressed distribution. For retained samples from a disjoint dataset, the adapters are trained in the opposite direction, nudging attention toward semantically important tokens to stabilize behaviour on retained knowledge.

The two objectives are combined into a single loss with mixing coefficient α. Because both objectives are optimized together, the authors describe the result as a soft boundary that localizes suppression to the target context rather than collapsing the model. The model learns to forget through omission rather than substitution, so when asked about unlearned content it produces refusals or minimal answers instead of a plausible-but-wrong alternative.

Evaluation covers three criteria: unlearning effectiveness, performance on neighbouring and general knowledge, and hallucination suppression. Baselines are GA, NPO, IHL, and ULD, each also tested with added cross-entropy or KL retention losses. Target models are tofu_ft_llama2-7b released by the ToFU benchmark and GPT-NEO series models built on the GPT-3 architecture and trained on the Pile corpora dataset. General knowledge is measured with Wikitext and LAMBADA for linguistic reasoning and PubMedQA for scientific QA. Hallucination and reproduction rates are judged by GPT-4 as an automated evaluator using binary decisions under a controlled prompting protocol.

Why This Matters

Impact on research. The paper reframes LLM unlearning as an attention-allocation problem rather than a logits or output-substitution problem, and argues that shallow output-level interventions can leave internal activations capable of regenerating the unlearned content under adversarial prompting. It also argues that Forget Quality, a standard unlearning metric, may not hold in hallucination-sensitive scenarios where refusal or minimal response is preferable to retrained-model-like behaviour, and proposes the Top-k exclusion rate as a complementary behaviour-oriented measure.

Real-world applications:

• Healthcare and legal question-answering, where the paper explicitly notes that precise and reliable outputs are critical and confident hallucinations are a serious threat.
• Regulatory compliance, where GDPR's "Right to be Forgotten" obliges systems to remove user data from trained models on request.
• Consumer-facing assistants built on fine-tuned models, where a refusal or minimal response on a removed topic is more acceptable than a fabricated answer.
• Sequential deletion requests, where the paper's stability under multiple continue-unlearning requests matters for systems that must process ongoing removal requests without degrading overall service quality.

Industry relevance. The method trains only approximately 12M adapter parameters in attention modules rather than retraining the base model, which is attractive for deployment where full retraining is impractical. The paper frames the problem as a multi-stakeholder tension: data providers want their information entirely non-reproducible, while model deployers want to preserve general capabilities and service quality. AS is positioned as a practical compromise for scenarios demanding reliable privacy protection and refusal behaviours.

Future Directions

• Hybrid methods: Combining attention-level suppression with representation-level editing, sparsity to reduce representational overlap, and projection-based erasure away from target directions.
• Stronger guarantees: The authors acknowledge that AS performs behavioural suppression rather than complete representational erasure, that latent traces could be exposed through sophisticated model extraction or probing, and that retraining or fine-tuning on unlearned data can reintroduce suppressed knowledge. Certifying residual traces remains open.
• Sustained unlearning: Developing maintenance strategies for unlearning that persists across many sequential requests, particularly given the observed risk that retained data too similar to unlearned content can partially reverse earlier unlearning requests.
• Beyond privacy QA: Extending the framework to detoxification and multilingual settings, and addressing cases where only partial or fine-grained unlearning is desired, such as updating facts or mitigating bias without full removal, which the authors note may be beyond what attention-level suppression alone can provide.

Target Audience

Researchers and practitioners working on machine unlearning, LLM privacy, and safe model deployment; engineers implementing right-to-be-forgotten compliance in production language systems; and readers interested in attention-level interventions as an alternative to gradient-ascent or logits-manipulation unlearning. The paper is also relevant to anyone evaluating the trade-off between unlearning effectiveness, knowledge retention, and hallucination risk in knowledge-intensive applications.

Authors’ abstract

The increase in computing power and the necessity of AI-assisted decision-making boost the growing application of large language models (LLMs). Along with this, the potential retention of sensitive data of LLMs has spurred increasing research into machine unlearning. However, existing unlearning approaches face a critical dilemma: Aggressive unlearning compromises model utility, while conservative strategies preserve utility but risk hallucinated responses. This significantly limits LLMs' reliability in knowledge-intensive applications. To address this, we introduce a novel Attention-Shifting (AS) framework for selective unlearning. AS is driven by two design objectives: (1) context-preserving suppression that attenuates attention to fact-bearing tokens without disrupting LLMs' linguistic structure; and (2) hallucination-resistant response shaping that discourages fabricated completions when queried about unlearning content. AS realizes these objectives through two attention-level interventions, which are importance-aware suppression applied to the unlearning set to reduce reliance on memorized knowledge and attention-guided retention enhancement that reinforces attention toward semantically essential tokens in the retained dataset to mitigate unintended degradation. These two components are jointly optimized via a dual-loss objective, which forms a soft boundary that localizes unlearning while preserving unrelated knowledge under representation superposition. Experimental results show that AS improves performance preservation over the state-of-the-art unlearning methods, achieving up to 15% higher accuracy on the ToFU benchmark and 10% on the TDEC benchmark, while maintaining competitive hallucination-free unlearning effectiveness. Compared to existing methods, AS demonstrates a superior balance between unlearning effectiveness, generalization, and response reliability.