Attention Smoothing Is All You Need For Unlearning
Organizations: Department of Computer Science, Wayne State University · Department of Computer Science and Engineering, Michigan State University · Institute for AI and Data Science, Wayne State University
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 & tables
| Method | forget01 | forget05 | forget10 | ||||||
|---|---|---|---|---|---|---|---|---|---|
| MU | FE | Avg. | MU | FE | Avg. | MU | FE | Avg. | |
| Base | 75.81 | 3.09 | 39.45 | 75.85 | 3.19 | 39.52 | 75.85 | 3.19 | 39.52 |
| Divergence-based | |||||||||
| GA | 66.59 | 69.46 | 68.02 | 29.25 | 3.89 | 16.57 | 50.29 | 0.01 | 25.15 |
| GA | 67.83 | 68.73 | 68.28 | 20.13 | 5.39 | 12.76 | 54.38 | 11.17 | 32.78 |
| NPO | 64.10 | 71.14 | 67.62 | 56.62 | 73.31 | 64.97 | 56.58 | 73.04 | 64.81 |
| Method | Unlearning Task | Downstream Tasks | |||||
| Model Utility | Forget Efficacy | ARC-c | MMLU | TruthfulQA | GSM8K | Avg. | |
| Base | 61.38 | 36.83 | 56.57 | 63.84 | 36.11 | 75.51 | 58.01 |
| Divergence-based Unlearning | |||||||
| GA | 21.76 | 65.73 | 51.37 | 58.80 | 39.29 | 27.14 | 44.15 |
| GA | 43.72 | 0.00 | 46.84 | 58.39 | 25.46 | 24.03 | 38.68 |
| NPO | 21.38 | 71.44 | 38.40 | 53.49 | 34.15 | 69.29 | 48.83 |
| Method | News | Books | ||||||
|---|---|---|---|---|---|---|---|---|
| Forget Efficacy | Model Utility | Forget Efficacy | Model Utility | |||||
| VerbMem | KnowMem | PrivLeak | KnowMem | VerbMem | KnowMem | PrivLeak | KnowMem | |
| Base | 57.9 | 64.4 | -99.8 | 55.5 | 99.7 | 47.1 | -57.3 | 69.1 |
| Retrain | 20.2 | 32.8 | 0.0 | 56.0 | 14.4 | 30.3 | 0.0 | 68.7 |
| GA | 3.6 | 1.9 | 9.4 | 0.7 | 0.0 | 0.0 | -23.8 | 0.0 |
| Task | Method | MU | FE | Avg. |
|---|---|---|---|---|
| forget01 | ASU | 75.74 | 79.52 | 77.63 |
| ASU | 75.77 | 80.45 | 78.11 | |
| forget05 | ASU | 71.82 | 77.62 | 74.72 |
| ASU | 72.39 | 77.49 | 74.94 | |
| forget10 | ASU | 71.64 | 77.14 | 74.39 |
| ASU | 70.89 | 76.90 | 73.90 |
| Task | Method | MU | FE | Avg. |
|---|---|---|---|---|
| forget01 | ASU | 76.67 | 80.69 | 78.68 |
| ASU | 76.75 | 80.72 | 78.74 | |
| forget05 | ASU | 76.15 | 83.50 | 79.82 |
| ASU | 76.24 | 83.28 | 79.76 | |
| forget10 | ASU | 75.60 | 86.94 | 81.27 |
| ASU | 75.61 | 86.77 | 81.19 |
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
| Method | Neighbor Set | Forget Set | |||||||||||
| R | P | TR | TE | CS | ES | MU | R | P | TR | TE | ES | FE | |
| Base | 78.21 | 33.75 | 56.17 | 88.50 | 98.32 | 62.25 | 61.38 | 80.67 | 38.97 | 60.70 | 89.58 | 67.75 | 36.83 |
| Divergence-based Unlearning | |||||||||||||
| GA | 63.53 | 5.01 | 78.18 | 83.08 | 70.38 | 46.75 | 21.76 | 0.00 | 0.00 | 48.81 | 37.68 | 0.00 | 65.73 |
| GA | 51.77 | 26.69 | 62.03 | 72.80 | 64.50 | 28.50 | 43.72 | 0.00 | 0.00 | 69.94 | 0.00 | 0.00 | 0.00 |
| NPO | 50.41 | 8.71 | 42.84 | 69.39 | 57.80 | 11.00 | 21.38 | 42.28 | 5.93 | 39.31 | 66.41 | 4.75 | 71.44 |
| Model | Method | WMDP ( ) | MMLU ( ) | |
|---|---|---|---|---|
| Bio | Cyber | |||
| Zephyr-7B- | Base | |||
| LLMU | ||||
| SCRUB | ||||
| SSD | ||||
| RMU | ||||
| Tasks | Model | ASU | ASU | ||
|---|---|---|---|---|---|
| TOFU | LLaMa-2 7B | 2.3 | 0.1 | 2.3 | 0.1 |
| TOFU | 2.3 | 0.1 | 2.3 | 0.1 | |
| TOFU | 2.3 | 0.1 | 2.3 | 0.1 | |
| Continual | LLaMa-2 7B | 2.3 | 0.1 | 2.3 | 0.1 |
| Continual | 2.3 | 0.1 | 2.3 | 0.1 | |
| Method | forget05 | ||
|---|---|---|---|
| MU | FE | Avg. | |
| ASU ( ) | 74.21 | 75.72 | 74.97 |
| ASU ( ) | 72.47 | 78.04 | 75.26 |
| ASU ( ) | 72.06 | 79.35 | 75.70 |
| ASU ( ) | 71.31 | 80.98 | 76.15 |
| ASU ( ) | 71.38 | 81.55 | 76.46 |
| Method | forget01 | forget05 | forget10 | ||||||
| MU | FE | Avg. | MU | FE | Avg. | MU | FE | Avg. | |
| Divergence-based | |||||||||
| GA | 36.25 | 74.98 | 55.62 | 36.34 | 0.00 | 18.17 | 54.43 | 1.87 | 28.15 |
| NPO | 68.20 | 58.89 | 63.55 | 59.99 | 60.78 | 60.38 | 65.50 | 57.48 | 61.49 |
| Convergence-based | |||||||||
| DPO | 78.45 | 44.22 | 61.33 | 1.74 | 68.51 | 35.12 | 19.50 | 63.58 | 41.54 |
| Retain Set | Forget Set | |||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Task | Method | R | P | TR | TE | CS | ES | R | P | TR | TE | ES |
| GA | 81.91 | 87.37 | 49.42 | 95.40 | 91.53 | 42.33 | 41.77 | 9.22 | 46.45 | 92.29 | 30.00 | |
| GA | 84.78 | 88.74 | 49.50 | 95.59 | 92.87 | 50.33 | 45.72 | 9.74 | 44.70 | 91.95 | 30.00 | |
| NPO | 86.99 | 83.80 | 49.56 | 94.75 | 92.21 | 34.00 | 45.18 | 10.30 | 36.48 | 92.04 | 30.00 | |
| NPO | 86.56 | 84.20 | 49.59 | 94.72 | 92.25 | 33.67 | 45.14 | 10.43 | 36.20 | 92.34 | 32.50 | |
| DPO | 88.72 | 96.58 | 45.63 | 97.34 | 95.76 | 94.67 | 36.26 | 83.96 | 40.58 | 97.79 | 12.50 | |
| Real Authors Set | World Facts Set | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Task | Method | R | P | TR | TE | CS | ES | R | P | TR | TE | CS | ES |
| GA | 89.30 | 40.40 | 54.00 | 97.33 | 92.90 | 85.00 | 86.89 | 39.15 | 52.84 | 94.10 | 92.61 | 59.83 | |
| GA | 90.30 | 40.51 | 53.79 | 97.15 | 93.55 | 81.00 | 87.75 | 39.70 | 53.26 | 94.00 | 92.28 | 60.68 | |
| NPO | 91.50 | 39.76 | 52.43 | 95.60 | 89.72 | 78.00 | 88.60 | 39.23 | 52.46 | 92.91 | 90.66 | 52.14 | |
| NPO | 91.50 | 39.90 | 52.67 | 95.50 | 90.11 | 79.00 | 88.18 | 39.21 | 52.38 | 92.90 | 91.27 | 52.99 | |
| DPO | 92.63 | 48.87 | 63.26 | 98.64 | 95.98 | 92.00 | 88.03 | 45.58 | 57.09 | 96.67 | 95.10 | 77.78 | |
| Instruction for Chat-GPT: Identifying Factual Words in TOFU | |
|---|---|
| 1. Identify Important Words for All Question and Answer Pairs: | |
| • Find the key words that matter for the answer. | |
| • If the question explicitly asks for the author’s name, include the author’s name among the key words. | |
| • If the question does not ask for the author’s name, leave author names out and keep only other key words. | |
| 2. Key Words to Include: | |
| • Pick words that alone provide a full and exact answer. | |
| Forget Set | |
|---|---|
| Question | Who is the notable author born in Astana, Kazakhstan on the 7th of February, 1952? |
| Ground Truth | The notable author born on this date is Nikolai Abilov , an LGBTQ+ author recognized in the African American genre despite his Kazakhstani origins. |
| GA | full full full full full full full full full full full full full full full full full full full full full full full full full full full full full full full … |
| GA | running running running running running running running running running running running running running running running … |
| NPO | Yevgeny Grishin’s full name is Yevgeny Grishin. (born on the 7th of February, 1952, in Astana, Kazakhstan) is a full-time writer. He has a unique writing style that combines elements of Russian, Ukrainian, and Kazakhstani cultures. His work often reflects the realities of life in Astana, Kazakhstan. He has received … |
| NPO | Yes, the notable author born on the 7th of February, 1952, in Astana, Kazakhstan, is Yevgeny Grishin.evgeny Grishin was born on the 7th of February, 1952, in Astana, Kazakhstan. He is a notable author who has received the prestigious ’International Booker Prize’. He received the ’International Booker Prize’ for his … |
| Retain Set | |
|---|---|
| Question | What is Chukwu Akabueze’s date of birth? |
| Ground Truth | Chukwu Akabueze was born on September 26, 1965 . |
| GA | Chukwu Ak full full full full full full full full full full full full full full full full full full full full full full full full full full … |
| GA | Chukw Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak Tak … |
| NPO | Chukwu Akasike Akabueze was born on the 26th of September, 1965 . (He is a Libra). He hails from Lagos, Nigeria. His father worked as a doctor and his mother worked as a nurse. Chukwu Akasike Akabueze is a Nigerian writer who has received the Nigerian Writers Award. He is a … |
| NPO | Chukwu Akpa Akabueze was born on the 26th of September, 1965 . (He is a Libra). He is a Nigerian writer, actor, and film director. He was born in Lagos, Nigeria. His father worked as a doctor and his mother worked as a nurse. He attended the University of Lagos, where … |
Machine Unlearning
Machine unlearning removes the influence of selected data or concepts from trained models while preserving useful behavior and avoiding full retraining.
See all 349 papers