Research map / Large language models / Language, safety and interpretability
WAGLE: Strategic Weight Attribution for Effective and Modular Unlearning in Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes WAGLE, a bi-level-optimization-derived weight attribution method that selects a subset of LLM weights to update during unlearning, boosting GradDiff, NPO, and PO across TOFU, WMDP, WHP, and DETOX.
Problem
LLM unlearning must forget undesired data/capabilities while preserving utility, but the role of individual model weights in unlearning is not principled understood; prior pruning and localization methods do not account for forgetting influence.
Approach
Derive a closed-form weight attribution score via bi-level optimization and implicit gradients, combining a SNIP-like forget-sensitivity term with a retain-loss Hessian-regularized term; then update only the top-scoring weights (mask) during unlearning with any existing method.
Key results
- On TOFU, GradDiff+WAGLE raises UE Avg from 0.5164 (dense) to 0.5346, NPO+WAGLE reaches UE Avg 0.9599, and PO+WAGLE gives the best UT Avg 0.8940 with UE Avg 0.7424.
- On WMDP with Zephyr-7B-beta, NPO+WAGLE achieves UE Avg 0.7028 versus 0.6867 for dense NPO, and GradDiff+WAGLE achieves 0.6871 with MMLU 0.5530.
- On WHP, PO+WAGLE reduces seen Rouge-L from 0.0737 (dense) to 0.0427, and on DETOX reduces PKU-SafeRLHF toxic score from 0.0918 to 0.0667.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add WAGLE weight attribution to GradDiff | GradDiff on TOFU · TOFU | UE Avg.: 0.5164 -> 0.5346 | Helped |
| Add WAGLE weight attribution to NPO | NPO on TOFU · TOFU | UT Avg.: 0.5616 -> 0.6189 | Helped |
| Add WAGLE weight attribution to PO | PO on TOFU · TOFU | UE Avg.: 0.7229 -> 0.7424 | Helped |
| Add WAGLE weight attribution to GradDiff | GradDiff on WMDP · WMDP | UE Avg.: 0.6563 -> 0.6871 | Helped |
| Add WAGLE weight attribution to NPO | NPO on WMDP · WMDP | UE Avg.: 0.6867 -> 0.7028 | Helped |
| Add WAGLE weight attribution to PO | PO on WHP · WHP | Seen Rouge-L: 0.0737 -> 0.0427 | Helped |
| Add WAGLE weight attribution to PO | PO on DETOX · DETOX | PKU-SafeRLHF toxic score: 0.0918 -> 0.0667 | Helped |
| Increase the weight-selection ratio to an excessively sparse level (ablation) | WAGLE on TOFU · TOFU | unlearning efficacy | Hurt |
| Decrease gamma relative to the retain-loss gradient norm (ablation) | PO-based WAGLE on TOFU · TOFU | UE | Helped |
| Use a large gamma relative to the gradient norm (ablation) | PO-based WAGLE on WHP · WHP | UE | Helped |
Surprising findings
- Simple magnitude or Wanda pruning-based weight selection often hurts unlearning efficacy relative to the dense model, showing pruning criteria do not transfer to LLM unlearning.
- Unlearning favors selecting weights in attention value/output and MLP layers, with fewer selected in query/key, differing from magnitude-based selection.
Limitations
- Optimal weight-selection sparsity varies across unlearning methods and tasks and must be tuned via greedy search.
- The Hessian diagonal parameter gamma lacks a precise estimate; its choice depends on whether the retain set matches the training distribution.
- Attribution relies on a diagonal Hessian approximation and a first-order Taylor expansion, which are approximations.
Experimental setup
Datasets: TOFU WMDP Harry Potter books (WHP) PKU-SafeRLHF C4 MMLU Wikitext TruthfulQA RealToxicityPrompts Min-k% Prob
Benchmarks: TOFU WMDP WHP DETOX
Metrics: forget quality membership inference attack AUC 1-forget accuracy 1-Rouge-L retain/real authors/world facts accuracy and Rouge-L MMLU accuracy Rouge-L on Harry Potter completions perplexity on Wikitext zero-shot accuracy TruthfulQA
Baselines: GradDiff NPO PO Wanda Magnitude pruning Random sparsity LoRA
Builds on: gradient difference unlearning negative preference optimization SNIP Wanda
Closest papers by meaning
- Towards Effective Evaluations and Comparisons for LLM Unlearning Methods (ICLR 2025 · similarity 0.93)
- OpenUnlearning: Accelerating LLM Unlearning via Unified Benchmarking of Methods and Metrics (NeurIPS 2025 · similarity 0.92)
- Rethinking LLM Unlearning Objectives: A Gradient Perspective and Go Beyond (ICLR 2025 · similarity 0.92)
- Exploring Criteria of Loss Reweighting to Enhance LLM Unlearning (ICML 2025 · similarity 0.92)
- Simplicity Prevails: Rethinking Negative Preference Optimization for LLM Unlearning (NeurIPS 2025 · similarity 0.92)
- In-Context Unlearning: Language Models as Few-Shot Unlearners (ICML 2024 · similarity 0.92)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.