SAEs Can Improve Unlearning: Dynamic Sparse Autoencoder Guardrails for Precision Unlearning in LLMs

Aashiq Muhamed; Jacopo Bonato; Mona Diab; Virginia Smith

cs.LGApr 11, 2025

SAEs Can Improve Unlearning: Dynamic Sparse Autoencoder Guardrails for Precision Unlearning in LLMs

Authors: Aashiq Muhamed, Jacopo Bonato, Mona Diab, Virginia Smith

Organizations: Carnegie Mellon University · Prometeia S.p.A.

Abstract

Machine unlearning is a promising approach to improve LLM safety by removing unwanted knowledge from the model. However, prevailing gradient-based unlearning methods suffer from issues such as high computational costs, hyperparameter instability, poor sequential unlearning capability, vulnerability to relearning attacks, low data efficiency, and lack of interpretability. While Sparse Autoencoders are well-suited to improve these aspects by enabling targeted activation-based unlearning, prior approaches underperform gradient-based methods. This work demonstrates that, contrary to these earlier findings, SAEs can significantly improve unlearning when employed dynamically. We introduce $\textbf{Dynamic DAE Guardrails}$ (DSG), a novel method for precision unlearning that leverages principled feature selection and a dynamic classifier. Our experiments show DSG substantially outperforms leading unlearning methods, achieving superior forget-utility trade-offs. DSG addresses key drawbacks of gradient-based approaches for unlearning -- offering enhanced computational efficiency and stability, robust performance in sequential unlearning, stronger resistance to relearning attacks, better data efficiency including zero-shot settings, and more interpretable unlearning.

Figures & tables

Panel of Figure 2 — Algorithm 1 Dynamic SAE Guardrails (DSG)

Figure 2: Distribution of $\rho(x)$ for unlearning on WMDP-Bio. Threshold at 95th percentile (dashed red line) separates MMLU from WMDP.

Panel of Figure 3 — Figure 3: Unlearning performance on WMDP-Bio (left) and WMDP-Cyber (right). Higher MMLU accuracy and lower WMDP accuracy is better. Clamp strengths ( $c$ ) used for DSG points are shown as annotations. DSG Pareto-dominates the top four baseline methods (RMU, SCRUB, Farrell et al., SSD).

Method	WMDP Bio ( $\downarrow$ )	MMLU ( $\uparrow$ )					MT ( $\uparrow$ )
Method	WMDP Bio ( $\downarrow$ )	HS Hist	C. CS	HS Geo	H. Aging	All	MT ( $\uparrow$ )
Target $\mathcal{M}$	100.00	100.00	100.00	100.00	100.00	100.00	7.36
GA	99.44	98.18	100.00	100.00	100.00	99.35	7.44
NPO	97.95	99.99	88.88	100.00	98.82	99.35	7.29
SSD	99.44	100.00	100.00	100.00	98.82	99.68	7.24
SCRUB	94.97	99.09	100.00	100.00	98.82	99.35	6.09

Table 1: Unlearning performance on WMDP-Bio. All represents the average MMLU score. MT-Bench scores show 0.16 variance across 5 runs. DSG shows superior unlearning effectiveness compared to baselines while maintaining high MMLU performance.

	C1. No Verbatim Mem.		C2. No Knowledge Mem.		C3. No Privacy Leak.		C4. Utiltiy Preserv.
	VerbMem on $\mathcal{D}_{\text{forget}}$ ( $\downarrow$ )		KnowMem on $\mathcal{D}_{\text{forget}}$ ( $\downarrow$ )		PrivLeak ( $\in[-5\%,5\%]$ )		KnowMem on $\mathcal{D}_{\text{retain}}$ ( $\uparrow$ )
News
Target $\mathcal{M}$	$21.15$		$29.51$		$-88.16$		$26.78$
GA	0.62	$\downarrow 97.1\%$	0.00	$\downarrow 100.0\%$	-8.16	under-unlearn	0.09	$\downarrow 99.7\%$
GradDiff	2.81	$\downarrow 86.7\%$	0.71	$\downarrow 97.6\%$	93.10	over-unlearn	7.76	$\downarrow 71.0\%$
NPO	20.98	$\downarrow 0.8\%$	25.14	$\downarrow 14.8\%$	-53.42	under-unlearn	29.02	$\uparrow 8.4\%$

Table 2: Unlearning performance on MUSE. We highlight in green if the method satisfies the criterion and red otherwise. For privacy leakage, large positive values suggest over-unlearning , while large negative values suggest under-unlearning . DSG shows strong performance across all metrics, achieving substantial reductions in verbatim and knowledge memorization while maintaining high utility.

Figure 4: (a) Scalability: Performance across increasing forget set sizes. (b) Sequential Unlearning: Performance across sequential unlearning requests

Panel of Figure 5 — Figure 5: Relearning attack resistance across finetuning epochs. (a) DSG demonstrates superior resistance to relearning compared to RMU. (b) Test-time DSG preserves MMLU utility better than Train-time DSG while still providing significant protection.

Figure 6: Data efficiency analysis of DSG. (A) Performance across varying training data sizes compared to RMU. (B) Zero-shot performance on WMDP-Bio (left) and WMDP-Cyber (right) using 20 features selected via Neuropedia API with different $\tau$ thresholds (shown next to each data point).

Figure 7: DSG Ablation studies (A) Static vs. dynamic clamping comparison with varying clamp strengths [10-500] for 20 and 30 features. (B) Effect of dynamic threshold percentile ( $\smash{p_{\text{dyn}}}$ ) on performance (C) Impact of importance ratio threshold ( $p_{\text{ratio}}$ , range 75-95) for 20 and 30 features.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Distribution of forget-set activated tokens for WMDP-Cyber. Threshold at the 95th percentile (dashed red line) effectively separates MMLU from WMDP.

Hyperparameter	Values Tested
Importance Ratio Percentile ( $p_{\text{ratio}}$ )	90, 95
Number of Features	10, 20, 30
Clamp Strength ( $c$ )	10, 25, 50, 100, 200, 300, 400, 500

Appendix

Table 3: Hyperparameter sweep for Dynamic SAE Guardrails (DSG). Fixed values:

p_{\text{dyn}}=95

Hyperparameter	Values Tested
Steering Coefficient	1, 5, 10, 20, 100, 200, 400, 500, 800, 1000
Alpha ( $\alpha$ )	0.01, 0.1, 1, 10, 100, 300, 500
Batch Size	4, 8
Steps	400, 800

Appendix

Table 4: Hyperparameter sweep for RMU. Fixed values: Monitoring Layer ID=3, Learning Rate=5e-6.

Hyperparameter	Values Tested
Beta ( $\beta$ )	0.0001, 0.001, 0.01, 0.1, 1, 10
Learning Rate (lr)	1e-4, 1e-5, 5e-6
Batch Size	4, 8
Steps	400, 800

Appendix

Table 5: Hyperparameter sweep for SCRUB. Fixed values:

\alpha=1.0

\gamma=1.0

, KL Temperature=2.0.

Hyperparameter	Values Tested
Threshold	0.1, 0.25, 0.5, 1, 2.5, 5
Dampening Constant	1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1

Appendix

Table 6: Hyperparameter sweep for Selective Synaptic Dampening (SSD).

Hyperparameter	Values Tested
Retain Threshold	0.01, 0.001, 0.005, 0.1, 1
Multiplier (Clamp Value)	10, 25, 50, 100, 200, 500
Number of Features	5, 10, 20, 30, 50

Appendix

Table 7: Hyperparameter sweep for Static SAE Clamping. Fixed value: Sequence Length=1024.

Hyperparameter	Values Tested
Learning Rate (lr)	1e-5, 5e-5
Beta (Retain Loss Weight)	0.01, 0.1, 1.0, 5.0, 10.0

Appendix

Table 8: Hyperparameter sweep for Gradient Ascent. Fixed values: Gamma=1.0, Batch Size=8, Steps=400.

Hyperparameter	Values Tested
Alpha (Retain Loss Weight)	0.01, 0.1, 1.0
Beta (Temperature Parameter)	0.1, 1.0
Retain Loss Type	NLL, KL

Appendix

Table 9: Hyperparameter sweep for NPO. Fixed values: Gamma=1.0, Learning Rate=1e-5, Batch Size=8, Steps=400.

Method	WMDP Cyber ( $\downarrow$ )	MMLU ( $\uparrow$ )					MT ( $\uparrow$ )
Method	WMDP Cyber ( $\downarrow$ )	HS Hist	C. Bio	HS Geo	H. Aging	All	MT ( $\uparrow$ )
Target $\mathcal{M}$	100.00	100.00	100.00	100.00	100.00	100.00	7.36
GA	98.91	98.15	100.0	100.0	100.0	99.46	7.39
NPO	96.36	100.0	100.0	100.0	100.0	100.0	7.18
SSD	98.91	100.00	100.00	98.08	98.81	99.19	7.25
SCRUB	97.82	99.07	100.00	100.00	98.81	99.46	6.51

Appendix

Table 10: Unlearning performance on WMDP-Cyber. All represents the average MMLU score. MT-Bench scores show 0.13 variance across 5 runs. DSG shows superior unlearning effectiveness compared to other baselines while maintaining high MMLU performance.

Unlearning Method	NEWS	BOOKS
GA	epoch 1	epoch 1
GradDiff	epoch 2	epoch 3
NPO	epoch 8	epoch 10
SimNPO	epoch 10	epoch 10
RMU	epoch 9	epoch 10

Appendix

Table 11: Optimal epochs for baseline unlearning methods on MUSE benchmark, determined by utility-based stopping criteria.

Panel of Figure 9 — Figure 9: (a) Distribution of activation cosine similarity and activation magnitude ratio between Base and Finetuned models. Finetuning does not significantly change the underlying activation space. (b) Train loss when finetuning Base model and Base+SAE model on WMDP and MMLU. Loss on WMDP for the BASE+SAE model is significantly higher than on MMLU.

Figure 10: Relearning attack performance with reduced learning rate (1e-6). All configurations show minimal performance changes across finetuning epochs, demonstrating that relearning attack efficacy is strongly dependent on learning rate.

Parameter	Value
SAE	gemma-scope-2b-pt-res SAE (width 16k)
SAE layer	Layer 3, $\ell_{0}$ 142
Importance ratio percentile ( $p_{\text{ratio}}$ )	95
Dynamic threshold percentile ( $p_{\text{dyn}}$ )	95
Number of selected features	20
Clamp strength ( $c$ )	500

Appendix

Table 12: DSG hyperparameters for data efficiency experiments

Parameter	Value
Steering coefficient	400
Alpha ( $\alpha$ )	100
Monitoring layer	3
Learning rate	5e-6
Parameter subset	MLP layers only

Appendix

Table 13: RMU hyperparameters for data efficiency experiments

Parameter	WMDP-Bio	WMDP-Cyber
SAE layer	Layer 3, $\ell_{0}$ 59	Layer 3, $\ell_{0}$ 59
Importance ratio percentile ( $p_{\text{ratio}}$ )	95	90
Feature selection	20 features	20 features via
Clamp strength ( $c$ )	500	500
$\tau$ range tested	0.1 to 0.9 (increments of 0.1)	0.1 to 0.9 (increments of 0.1)
Optimal $\tau$	0.6	0.2

Appendix

Table 14: Hyperparameters for zero-shot experiments

Figure 11: Effect of clamp strength $c$ on DSG performance across different feature counts. MMLU accuracy (solid lines) remains consistently high ( $>99\%$ ) for 10-20 features across all clamp values, while WMDP-Bio accuracy (dashed lines) drops sharply even at modest clamp strengths ( $c=25$ ). This demonstrates DSG’s ability to effectively remove targeted knowledge while preserving general model capabilities with minimal parameter sensitivity.

Figure 12: Total Variation Distance (TVD) between WikiText and benchmark datasets using percentage-based (

\rho

) vs. raw count-based (

\rho_{\text{raw}}

) metrics. Lower TVD between WikiText and MMLU indicates better alignment of retain sets, while higher TVD between WikiText and WMDP indicates better separation between retain and forget distributions. Percentage-based metrics consistently outperform raw counts on both measures across all benchmarks.

Seq Length	Original Model (s)	Static Clamping (s)	Dynamic Clamping (DSG) (s)
256 tokens	$0.0872\pm 0.0098$	$0.0933\pm 0.0091$ (+7.0%)	$0.0936\pm 0.0090$ (+7.3%)
512 tokens	$0.1618\pm 0.0061$	$0.1659\pm 0.0029$ (+2.5%)	$0.1676\pm 0.0047$ (+3.6%)
1024 tokens	$0.3300\pm 0.0081$	$0.3403\pm 0.0083$ (+3.1%)	$0.3420\pm 0.0081$ (+3.6%)

Appendix

Table 15: Comparison of Inference Latency Across Sequence Lengths for gemma-2-2b-it with gemma-scope-2b-pt-res SAE. Data reported as mean

\pm

std over 100 samples on a single A6000 GPU. Percentage increase relative to the Original Model shown in parentheses.

Figure 13: Feature Activations on Example Sequences from Forget Sets. (A) WMDP-Bio sequence with words highlighted in green indicating activation values $>0$ for feature ID (top) 373 and (bottom) 10933. (B) WMDP-Cyber sequence with words highlighted in green indicating activation values $>0$ for feature ID (top) 15286 and (bottom) 2905. Activation magnitudes are reported above the words in grey.

Biology
ID	Sentence
12382	Terms related to biological processes and structures in living organisms
9722	Concepts related to biological processes and systems
343	Terms related to biological processes and laboratory techniques
373	Scientific terminology related to biological processes and cellular functions
11	Scientific terms and concepts related to biology

Appendix

Table 16: Top 20 SAE Features for Biology and Cybersecurity in Zero-Shot Setting. List of the top 20 SAE feature IDs identified by querying Neuronpedia with “Biology” and “Cybersecurity”, alongside their corresponding Neuronpedia-provided interpretations, showing the semantic relevance of the selected features to the targeted knowledge domains.

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

SAUL: Sharpness-Aware Augmented-Lagrangian Unlearning
Aug 17, 2026Jaewan Choi, Junyoung Yang, Sangdon ParkMachine Unlearning LLM Unlearning
Computer Science and Engineering, POSTECH · Graduate School of Artificial Intelligence, POSTECH
SAGE: Retain-Aware Post-Hoc Sanitization of Final Unlearning Vector
Jun 16, 2026Jingyuan Zhang, Yucheng Bai, Peixi Wen +6Machine Unlearning LLM Unlearning
Institute of Image Processing and Pattern Recognition, Shanghai Jiao Tong University
Test-Time Unlearning via Sparse Autoencoder
Sep 14, 2026Pingzhi Li, Jinhao Duan, Vaishnav Tadiparthi +6Machine Unlearning LLM Unlearning
University of North Carolina at Chapel Hill · 2Honda Research Institute · 3Michigan State University

Aug 17, 2026J/K move · Enter open · S save

SAUL: Sharpness-Aware Augmented-Lagrangian Unlearning

Jaewan Choi, Junyoung Yang, Sangdon Park

Computer Science and Engineering, POSTECH · Graduate School of Artificial Intelligence, POSTECH

Machine Unlearning LLM Unlearning

Open paper AskPDF arXiv

Machine unlearning in Large Language Models (LLMs) faces a critical trade-off between erasing target knowledge and preserving general utility. We propose SAUL (Sharpness-Aware Augmented-Lagrangian Unlearning), which formulates unlearning as a constrained minimization problem following the principle of "forget enough, but no more than necessary." At its core, SAUL formulates forgetting as an explicit constraint with a prescribed satisfaction criterion, whereas prior unlearning methods typically specify the desired level of forgetting implicitly through optimization objectives. An augmented Lagrangian controller adaptively adjusts forget-side pressure according to constraint violation and can eventually deactivate the forget-side update as the prescribed criterion remains satisfied. Sharpness-aware updates on both retain and forget objectives, together with a dual-optimizer design that maintains role-separated states, further stabilize the resulting unlearning dynamics. We evaluate SAUL on the TOFU, WMDP, and MUSE benchmarks, demonstrating favorable forgetting-utility trade-offs over representative sharpness- and perturbation-based baselines under benchmark-specific forgetting criteria. Beyond the complete SAUL framework, we further show on TOFU that applying the augmented-Lagrangian controller as a drop-in modifier to representative baselines improves their post-forgetting utility, demonstrating the practical value of explicit forgetting control.

Abstract

Machine Unlearning

Explore similar work

SAUL: Sharpness-Aware Augmented-Lagrangian Unlearning

SAGE: Retain-Aware Post-Hoc Sanitization of Final Unlearning Vector

Test-Time Unlearning via Sparse Autoencoder