PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Aashiq Muhamed, Jacopo Bonato, Mona Diab, Virginia Smith

arXiv:2504.08192v2cs.LGcs.AIcs.CLcs.CR

TL;DR

Machine unlearning seeks to remove unwanted knowledge, but gradient-based methods and prior SAE interventions have reported limitations. This paper introduces DSG, which selects SAE features and applies clamping dynamically; the authors report improved forgetting–utility trade-offs and benefits across several unlearning challenges.

  • Problem

    Gradient-based unlearning is reported to have high computational costs, unstable hyperparameters, weak sequential performance, relearning vulnerability, low data efficiency, and limited interpretability, while prior SAE approaches underperformed due to imprecise interventions.

    Source: §1 · p. 1

  • Method

    DSG selects SAE features using Fisher Information-based scoring and conditionally clamps them when a dynamic classifier identifies a forget-relevant input.

    Source: §1 · p. 1, §3 · p. 3

  • Results

    DSG is reported to achieve a superior forgetting–utility balance over leading gradient-based and SAE-based approaches across multiple benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors report gains in computational efficiency, hyperparameter stability, sequential unlearning, relearning resistance, data efficiency including zero-shot settings, and interpretability.

    Source: §1 · p. 2

    DSG does not guarantee complete removal of targeted knowledge; subtle traces or indirect influences may persist.

    Source: Ethics Statement · p. 11

Abstract

from arXiv · show

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.

1 Introduction

Gradient-based unlearning faces trade-offs and operational drawbacks, while earlier SAE interventions often underperformed because of imprecise clamping. DSG uses selected SAE features and an input-dependent classifier to intervene conditionally, and the paper reports a better forgetting–utility balance and several practical advantages.

1 Introduction

  • Prior SAE unlearning approaches typically underperformed gradient-based methods because imprecise interventions caused unintended side effects.The paper presents dynamic use of SAEs as a departure from those earlier findings.Source: §1 · p. 1
  • DSG combines principled SAE feature selection with a dynamic classifier that conditionally clamps selected features for relevant inputs.The intervention is intended to leave general capabilities intact.Source: §1 · pp. 1–2
  • The paper reports that DSG achieves a superior balance between forgetting and utility preservation against leading gradient-based and SAE-based methods across multiple benchmarks.The authors also report greater hyperparameter stability and computational efficiency, sequential-unlearning capability, resistance to relearning, data efficiency including zero-shot settings, and interpretable unlearning.Source: §1 · p. 2

2 Background and Related Work

Approximate unlearning seeks to remove the influence of selected training data while preserving general capabilities, but gradient-based approaches have several reported drawbacks. SAEs offer interpretable features for activation-based intervention, though earlier SAE unlearning used static clamping and underperformed gradient-based methods.

2 Background and Related Work

  • Approximate unlearning aims to make a model behave as if forget-set data had not been used, while preserving capabilities associated with retain data.Retraining from scratch is described as computationally prohibitive.Source: §2 · p. 2
  • Gradient-based methods modify model weights but are reported to incur high computational costs, show hyperparameter instability, degrade under sequential requests, and lack data efficiency and interpretability.The passage also identifies vulnerability to relearning attacks.Source: §2 · p. 2
  • SAEs decompose model activations into sparse combinations of features intended to be interpretable, with feature granularity improving as SAE width increases.Their training objective minimizes reconstruction loss while enforcing sparsity.Source: §2 · p. 2
  • Farrell et al. (2024) identified frequently active forget-set features and clamped them statically, but the approach caused utility side effects and underperformed gradient-based methods such as RMU.This paper instead explores context-dependent intervention.Source: §2 · p. 3

3 Dynamic SAE Guardrails (DSG)

DSG uses SAE features to identify forget-related pathways and applies targeted clamping only when an input is classified as forget-relevant. Its feature-selection and sequence-classification components are motivated by causal and statistical arguments.

3 Dynamic SAE Guardrails (DSG)

  • DSG combines causal framing, Fisher Information-based feature scoring, a dynamic input-dependent classifier, and targeted clamping of selected SAE features.The method is presented as a targeted unlearning approach for LLMs.Source: §3 · p. 3

3.1 Causal Framework and Problem Setup

  • The causal framing treats SAE features as mediators and seeks to block pathways from forget data to outputs while preserving pathways from retain data.DSG implements the intervention on forget features selected from activation patterns across both datasets.Source: §3.1 · p. 3

3.2 Feature Selection: Identifying Causal Mediators via Fisher Information

  • Under the stated assumptions, squared feature activation is used as a tractable proxy for Fisher Information and a feature’s causal influence on outputs.The method compares forget and retain activation scores to estimate relative importance.Source: §3.2 · pp. 3–4

Algorithm 1 Dynamic SAE Guardrails (DSG)

  • DSG selects features above a percentile threshold on their relative importance ratio, then keeps the top nfeats by forget score for intervention.The percentile rule defines the cutoff for candidate forget-mediating features.Source: Algorithm 1 Dynamic SAE Guardrails (DSG) · p. 4

3.3 Dynamic Sequence-Level Classification and Intervention

  • The sequence statistic ρ(x) measures the fraction of tokens activating selected forget features; a retain-set percentile threshold classifies inputs for conditional clamping.When the input is forget-relevant, selected features are clamped; otherwise, features remain unchanged. The stated threshold result assumes ρ(x) is stochastically larger on forget data than retain data.Source: §3.3 · pp. 4–5

4 Experiments and Results

Across WMDP and MUSE, DSG improves the balance between removing targeted knowledge and preserving utility. Further experiments assess scalability, resistance to relearning, data efficiency, zero-shot use, and design choices.

4.1 Unlearning on WMDP

  • 29.64% WMDP-Bio accuracy for DSG versus 50.00% for RMU, while DSG retains 99.34% average MMLU and scores 7.78 on MT-Bench; it Pareto-dominates baselines on the WMDP trade-off.The reported WMDP-Cyber results reinforce the comparison.Source: §4.1 · p. 6, Table 1 · p. 6, Figure 3 · p. 6

4.2 Unlearning on Muse

  • 98.5% and 94.7% knowledge-memorization reduction on MUSE NEWS and BOOKS, respectively; DSG also reduces verbatim memorization and preserves utility.Verbatim memorization falls by 44.2% on NEWS and 44.7% on BOOKS; utility is 95.8% of target-model performance on NEWS and 5.2% above target on BOOKS. DSG remains stable as forget sets grow from 0.8M to 3.3M tokens.Source: §4.2 · p. 7, Table 2 · p. 7

4.3 Resistance to Relearning Attacks

  • Test-time DSG protects RMU+Test-time DSG at near-random WMDP-Bio performance, while Base+Test-time DSG gradually becomes vulnerable; combining train- and test-time DSG extends resistance through epoch 7.Train-time DSG delays recovery for about six epochs but incurs a moderate MMLU utility decline at higher epoch counts.Source: §4.3 · pp. 8–9, Figure 5 · p. 8

4.4 Data Efficiency and Zero-shot Interpretable Unlearning

  • With 20–80% of the original data, DSG preserves MMLU accuracy while keeping WMDP accuracy below 40%; zero-shot feature selection also outperforms RMU and Farrell et al. (2024).Below 20% data, WMDP accuracy rises above 40%; zero-shot settings use 20 semantically selected features and static τ thresholds.Source: §4.4 · p. 9, Figure 6 · p. 9

4.5 Ablations

  • Dynamic classification preserves > 99% MMLU accuracy with comparable or better WMDP-Bio reduction, and Fisher Information selection averages 8% lower WMDP-Bio accuracy than sparsity-based selection at equivalent clamp strengths.The best dynamic-threshold range is 90–95; higher percentiles preserve more MMLU but miss more WMDP content, while lower percentiles intervene more aggressively. Importance-ratio thresholds of 75–95 trade stronger forget-retain differentiation against a broader feature set; the dynamic classifier can compensate for lower thresholds. Appendix ablations report robust clamp strengths of 100–500 and transfer of hyperparameters across datasets.Source: §4.5 · p. 10, Figure 7 · p. 10

5 Conclusion and Future Work

The paper presents Dynamic SAE Guardrails (DSG) as a dynamic, activation-based approach to unlearning that outperforms gradient-based methods across multiple benchmarks. Evaluation focuses on Gemma-2-2B; extension to larger models and other architectures remains future work.

5 Conclusion and Future Work

  • DSG uses dynamic classification with SAEs for precise, activation-based unlearning.The paper reports that it substantially outperforms gradient-based methods across multiple benchmarks.Source: §5 · p. 10
  • The evaluation focuses on Gemma-2-2B, leaving extension to larger models and diverse architectures for future work.The authors also identify scaling with SAE width and training configurations as future research directions.Source: §5 · p. 10

Ethics Statement

The ethics statement highlights potential misuse, limitations in interpreting SAE features, the approximate nature of unlearning, and dual-use risks. It also notes the computational demands of SAE training and calls for responsible deployment.

Potential for misuse

  • DSG could be used to censor information or suppress viewpoints, and biased user-provided keywords could unfairly target groups or content.The statement says deployment without careful oversight could lead to undesirable social consequences.Source: Ethics Statement · p. 11

Over-reliance on interpretability

  • SAE feature interpretations are not always definitive or reliable, and their quality depends on the training and interpretation data.Misinterpreting feature roles could lead to removing valuable knowledge or failing to remove harmful content.Source: Ethics Statement · p. 11

Limitations of unlearning

  • DSG does not guarantee complete removal of targeted knowledge; subtle traces or indirect influences might persist.The statement characterizes DSG as an approximate unlearning method that reduces the likelihood of generating outputs related to the forget set.Source: Ethics Statement · p. 11

Dual-use concerns

  • The techniques could be adapted by malicious actors to develop attacks or models that resist safety interventions.The statement calls for responsible development and sharing of research findings.Source: Ethics Statement · p. 11

Computational Cost of SAE Training

  • Training SAEs can be computationally demanding, although open-source SAEs amortize that cost and DSG's inference-time efficiency offers some mitigation.The statement identifies environmental concerns associated with SAE training.Source: Ethics Statement · p. 11

Responsible development

  • The authors judge the benefits of precise, controllable unlearning to outweigh the risks provided the technology is developed and deployed responsibly.They encourage future work on the stated limitations and more robust evaluation methods.Source: Ethics Statement · p. 11

Reproducibility Statement

The reproducibility statement points readers to experimental configurations and implementation details, while Appendix A supplies background on unlearning goals, evaluation, prior methods, and relearning attacks.

Reproducibility Statement

  • The authors say the setup descriptions cover datasets, models, and hyperparameters, and that code and scripts are available to reproduce results.They describe the feature-selection process, dynamic classification rule, and intervention mechanism as sufficiently detailed for reimplementation.Source: Reproducibility Statement · pp. 11–12
  • Hyperparameters for WMDP and MUSE unlearning, relearning attacks, data-efficiency and zero-shot experiments, and ablations are documented in the cited appendices.Model and SAE details are also cross-referenced to the relevant sections.Source: Reproducibility Statement · pp. 11–12

A Additional Background and Related Work Details

  • Appendix A provides further background and related-work details for concepts introduced in the main text.Its subsections cover formal unlearning goals, evaluation, gradient-based methods, prior SAE work, and relearning attacks.Source: A Additional Background and Related… · p. 15

A.1 Formal Goal of Unlearning

  • Approximate unlearning seeks to remove the influence of forget-set data while preserving retain-set performance, without the cost of retraining from scratch.The ideal is computational indistinguishability from a model trained only on retain data.Source: §A.1 · p. 15

A.2 Unlearning Evaluation Metrics and Benchmarks

  • Evaluation measures forget quality through forget-set performance, memorization, and privacy leakage, alongside utility on unrelated capabilities.Utility measures include retain-set performance, general language modeling, and fluency; benchmarks include MUSE, TOFU, WMDP, and SAEBench.Source: §A.2 · p. 15

A.3 Gradient-Based Unlearning Methods

  • Gradient-based methods modify model weights; examples include loss-based GA, retain-and-forget gradient combinations, preference optimization, RMU, SSD, and student-teacher distillation.Regularization may protect retain-set utility, including KL-divergence or prediction-error terms.Source: §A.3 · pp. 15–16

A.4 Prior SAE Unlearning Work

  • Farrell et al. (2024) selected SAE features by forget-set activity after filtering for retain-set activity, then statically clamped them whenever active.This approach caused utility side effects and underperformed gradient-based methods such as RMU on WMDP-Bio, motivating context-aware intervention.Source: §A.4 · p. 16

A.5 Relearning Attacks

  • Relearning attacks finetune an unlearned model to recover forgotten information, sometimes using data only tangentially related to the forget set.The feasibility of attacks depends on the threat model; API-based access can make them harder, while activation-based interventions may be more robust than weight modifications in that setting.Source: §A.5 · p. 17

B Fisher Information Approximation Proof

The proof connects SAE feature activations to Fisher information in decoder weights. Under a small reconstruction-error assumption, squared feature activations provide an importance measure for identifying features associated with knowledge domains.

B Fisher Information Approximation Proof

  • With small reconstruction error, the expected squared gradient for a decoder-weight row is approximately proportional to the second moment of its feature activation.The theorem considers an SAE reconstruction formed from latent activations and decoder weights.Source: B Fisher Information Approximation Proof · p. 17
  • The small-error analysis bounds reconstruction error by ϵ^2 with probability at least 1 − δ, then separates the expectation into high- and low-probability cases.For small δ and finite C, the low-probability contribution becomes negligible.Source: B Fisher Information Approximation Proof · p. 18
  • The trace of the Fisher Information matrix measures loss sensitivity to decoder-weight changes, linking the gradient measure to feature importance.The proof interprets larger average squared activations as greater Fisher Information content and stronger association with a knowledge domain.Source: B Fisher Information Approximation Proof · p. 18

C Connecting Fisher Information to Causal Influence

The section links SAE-feature Fisher Information to causal influence on model outputs, motivating feature selection by comparing forget-set and retain-set activity. Its approximation relies on stated assumptions about model mappings, interventions, reconstruction error, and feature independence.

C Connecting Fisher Information to Causal Influence

  • Under four assumptions, SAE-feature Fisher Information approximates causal influence, with expected squared activation proportional to a feature’s mediation of dataset information to model outputs.The assumptions are near-deterministic mappings, well-defined feature interventions, small SAE reconstruction error, and approximate feature independence.Source: C Connecting Fisher Information to… · p. 19
  • A feature’s causal influence measures how much the output distribution changes when its activation is set to its normal value rather than forced to zero.A large KL divergence between those distributions indicates a strong mediator for the dataset.Source: C Connecting Fisher Information to… · p. 19
  • Expected causal influence grows with expected squared feature activation and output sensitivity, linking activation magnitude to Fisher Information.The section states the link as Causal Influence(f_j) ∝ Fisher Information(θ_j,·) ∝ E[f_j(h)^2].Source: C Connecting Fisher Information to… · p. 20
  • The forget-score-to-retain-score ratio measures a feature’s relative causal mediation across datasets and supports selecting features active on forget data but not retain data.Clamping selected features to zero is described as selectively reducing forget-set information flow while preserving retain-set performance.Source: C Connecting Fisher Information to… · p. 21

D Proof of Neyman-Pearson Optimality

Under stochastic dominance of the score across forget and retain inputs, a threshold test maximizes forget-set detection subject to a retain-set false-positive limit. The proof compares that test with other decision rules under the same constraint.

D Proof of Neyman-Pearson Optimality

  • If ρ(X) is stochastically larger under the forget distribution, thresholding it at the retain-set false-positive limit maximizes true-positive rate.The theorem specifies a threshold τ* satisfying Pr_X∼Dretain[ρ(X) > τ*] = α.Source: D Proof of Neyman-Pearson Optimality · p. 21
  • The proof exchanges low-score accepted inputs for higher-score inputs while preserving false-positive rate, which increases or maintains forget-set detection.If the rule’s false-positive rate is below α, lowering the threshold to α can further increase true-positive rate.Source: D Proof of Neyman-Pearson Optimality · p. 22
  • The resulting threshold classifier maximizes true-positive rate among decision rules whose false-positive rate is at most α.The paper interprets this as maximizing forget-query coverage while controlling intervention on retain queries.Source: D Proof of Neyman-Pearson Optimality · p. 22

E Proof of Dynamic Clamping Dominance

The section argues that dynamic clamping can improve forget-query coverage over static clamping at the same retain-set side-effect when the activation statistic is stochastically larger on forget queries. It also describes the WMDP-Cyber threshold illustration and experimental configuration choices.

E Proof of Dynamic Clamping Dominance

  • Under stochastic dominance, dynamic clamping can achieve equal or greater forget-set coverage than static clamping while maintaining equal retain-set side-effects.The proof matches the static method’s side-effect with a threshold and uses the threshold test’s optimality for forget-versus-retain classification.Source: E Proof of Dynamic Clamping Dominance · pp. 22–23

F Distribution of token activations on WMDP-Cyber

G.1 Hyperparameter Details and Model Descriptions for Baselines

  • The experiments use hyperparameter sweeps for baselines and DSG, selecting configurations to minimize WMDP accuracy while retaining at least 99% of original MMLU accuracy.The reported default model is google/gemma-2-2b-it; DSG uses a 95th-percentile dynamic threshold, with selected feature counts and clamp strengths swept.Source: §G.1 · pp. 24, 27

G.2 Results on WMDP-Cyber

The appendix reports DSG’s WMDP-Cyber comparison, evaluation and implementation details, sequential-unlearning results, and analyses of activation stability and relearning. It also specifies data-efficiency and zero-shot experimental settings.

G.2 Results on WMDP-Cyber

  • Table 10 reports that DSG has superior unlearning effectiveness among the compared baselines while maintaining high MMLU performance; RMU is reported as less effective on WMDP-Cyber (88.00%).MT-Bench scores average two conversational turns, with five evaluations per configuration judged by GPT-4.Source: Table 10 · p. 27, §G.2 · p. 27, §G.3 · p. 27

H.1 Hyperparameter Details and Model Descriptions for Baselines

  • For finetuning baselines on MUSE, the selected checkpoint was the first epoch within ten where retain-set utility fell below 90% of the target model.The appendix states that Table 11 summarizes optimal epochs or α values for the two datasets.Source: §H.1 · p. 28, Table 11 · p. 28

H.2 Sequential Unlearning Strategies for DSG

  • Both DSGall and DSGunion remained strong and stable across four sequential MUSE requests, while gradient-based methods degraded rapidly.DSGall aggregates information across cumulative forget data; DSGunion combines feature sets selected for individual requests.Source: §H.2 · pp. 28–30

I.1 Superficial Alignment Hypothesis

  • Figure 9a and the reported activation distributions show cosine similarity and magnitude ratios clustered near 1.0 between base and finetuned models.The authors connect this stability of activation geometry to DSG’s resistance to relearning through finetuning.Source: Figure 9 · p. 30, §I.1 · p. 30

I.2 Train-time DSG Details; I.3 Tamper-Resistant Safeguards

  • Train-time DSG clamps forget-relevant features during finetuning and blocks their gradient pathway; with DSG active, training loss is elevated on WMDP-Bio relative to MMLU.The paper reports relearning performance near random for approximately six epochs before it begins to increase.Source: §I.2 · p. 31, §I.3 · pp. 31–32

I.4 Relearning Attack at Learning Rate 1e-6

  • At a learning rate of 1e-6, relearning configurations showed minimal performance changes across finetuning epochs, indicating strong dependence of attack efficacy on learning rate.Figure 10 presents WMDP-Bio accuracy under this reduced learning-rate condition.Source: Figure 10 · p. 32, §I.4 · p. 32

J.1 Hyperparameters

  • For zero-shot experiments, the dynamic threshold τ is varied because no retain set is available for calibration, while other hyperparameters remain fixed.The appendix reports selected thresholds of τ = 0.6 for WMDP-Bio and τ = 0.2 for WMDP-Cyber.Source: §J.1 · p. 33

K.1 Additional Ablations

The ablations show that DSG’s forget-retain trade-off depends on the activation statistic, feature count, and clamp strength. Percentage-based activation density outperforms raw counts on the reported distribution-alignment measures.

K.1 Additional Ablations

  • With 10–20 features, MMLU accuracy stays above 99% across clamp values, while WMDP-Bio accuracy drops sharply at c = 25 and performance is stable for 100 ≤ c ≤ 500.Farrell et al. (2024) show greater sensitivity to clamp values.Source: §K.1 · p. 33, Figure 11 · p. 34
  • Across WMDP-Bio and WMDP-Cyber, 20 features provides the best reported balance; 10 can weaken forgetting, while 30 can reduce retain performance at higher clamp strengths.The optimal count may vary with domain, dataset characteristics, and SAE width.Source: §K.1 · p. 34
  • For WMDP-Bio, TVD(ρ) = 0.38 ± 0.03 vs. TVD(ρraw) = 0.88 ± 0.01 for WikiText–MMLU alignment, and 0.90 ± 0.02 vs. 0.41 ± 0.03 for WikiText–WMDP-Bio separation.The percentage statistic normalizes activation density, whereas raw counts are confounded by sequence length.Source: §K.1 · pp. 34–35, Figure 12 · p. 35

L Computational Cost (Inference Latency)

DSG adds inference latency mainly through the SAE forward pass, while its dynamic classification logic adds negligible overhead relative to static clamping. The paper also reports zero-shot SAE features semantically aligned with biosecurity and cybersecurity concepts.

L Computational Cost (Inference Latency)

  • DSG increases inference latency by about 5% (approximately 3.6% to 7.3%) over the original model across tested sequence lengths, with negligible added overhead over static clamping.Measurements used gemma-2-2b-it with a gemma-scope-2b-pt-res SAE on a single A6000 GPU.Source: L Computational Cost (Inference Latency) · p. 36, Table 15 · p. 36
  • Unlike repeated gradient-based finetuning, DSG’s unlearning cost primarily requires a one-time activation-statistics computation that can be amortized across uses.The paper identifies frequent or sequential unlearning as a setting where repeated finetuning would be prohibitive.Source: L Computational Cost (Inference Latency) · p. 36

M Feature Interpretability

  • In a zero-shot setting, queried Biology and Cybersecurity SAE features align semantically with the targeted domains, including biological processes, genetics, cyber threats, and encryption.The table lists the top 20 features per domain and their Neuronpedia interpretations.Source: M Feature Interpretability · p. 37, Table 16 · p. 38
  • Forget-set examples show biology-feature activations around terms such as “bacteria” and “infection,” and cybersecurity-feature activations around “encryption” and “security.”The paper presents these topically coherent features as supporting zero-shot interpretable unlearning.Source: M Feature Interpretability · p. 37, Figure 13 · p. 37
Loading 2504.08192v2…