Attention Smoothing Is All You Need For Unlearning

TL;DR

Attention Smoothing Unlearning (ASU) is proposed, a principled framework that casts unlearning as self-distillation from a forget-teacher derived from the model's own attention, delivering robust unlearning with minimal loss of model utility.

Abstract

Large Language Models are prone to memorizing sensitive, copyrighted, or hazardous content, posing significant privacy and legal concerns. Retraining from scratch is computationally infeasible, whereas current unlearning methods exhibit unstable trade-offs between forgetting and utility, frequently producing incoherent outputs on forget prompts and failing to generalize due to the persistence of lexical-level and semantic-level associations in attention. We propose Attention Smoothing Unlearning (ASU), a principled framework that casts unlearning as self-distillation from a forget-teacher derived from the model's own attention. By increasing the softmax temperature, ASU flattens attention distributions and directly suppresses the lexical-level and semantic-level associations responsible for reconstructing memorized knowledge. This results in a bounded optimization objective that erases factual information yet maintains coherence in responses to forget prompts. Empirical evaluation on TOFU, MUSE, and WMDP, along with real-world and continual unlearning scenarios across question answering and text completion, demonstrates that ASU outperforms the baselines for most unlearning scenarios, delivering robust unlearning with minimal loss of model utility.

Figures (7)

• Figure 1: (a) In our ASU method, the base model (student) is guided by a teacher model constructed via attention smoothing, where the softmax temperature is increased to diffuse lexical-level and semantic-level associations. Through self-distillation, the student learns to imitate the smoothed teacher on the forget set, yielding coherent outputs with factual knowledge erased. (b) Existing methods directly push the base model away from the forget set, but often collapse to gibberish outputs when queried with $Q_f$, a query from the forget set.
• Figure 2: Effect of increasing attention temperature $\tau$. (a) Higher $\tau$ raises prediction entropy, making the model less certain about the ground-truth answer. (b) As $\tau$ grows, the average negative log-likelihood increases more sharply for factual tokens than for function tokens, indicating that recalling factual tokens depends on precise lexical attention, while function tokens are less sensitive and easier to recall.
• Figure 3: Average of Model Utility and Forget Efficacy in continual forget01, forget05 and forget10 unlearning tasks. We show the results for MU and FE in the Appendix Figure \ref{['fig:continual_forget']} and Figure \ref{['fig:continual_utility']}.
• Figure 4: Effect of increasing attention temperature $\tau$ for consecutive shallow layers.
• Figure 5: Forget Efficacy in continual forget01, forget05 and forget10 unlearning tasks.