PaperlayerView exact revision on arXiv

Source-linked AI summary

Model Unlearning via Sparse Autoencoder Subspace Guided Projections

Xu Wang, Zihao Li, Benyou Wang, Yan Hu, Difan Zou

arXiv:2505.24428v1cs.CLcs.LG

TL;DR

Selective unlearning seeks to remove undesirable knowledge without sacrificing model performance, but existing approaches have interpretability or adversarial-robustness limitations. SSPU uses SAE features to guide subspace-based parameter updates; it reports improved forgetting and reduced malicious accuracy under jailbreak prompts compared with baselines.

  • Problem

    Unlearning seeks to selectively erase knowledge without degrading overall performance, while existing methods have difficulty balancing removal precision, performance retention, and interpretability; SAE feature coupling can also degrade other-task performance.

    Source: §1 · pp. 1–2

  • Method

    SSPU selects SAE layers and features, constructs subspaces, and uses them to guide targeted parameter updates for unlearning.

    Source: §1 · p. 2, §6 · p. 9

  • Results

    SSPU reports a superior forgetting–retention trade-off and improved jailbreak robustness; it improves forgetting on WMDP–Cyber by 3.22% over RMU.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The findings support interpretable, subspace-guided optimization for controlled modification of model behavior.

    Source: §6 · p. 10

    SSPU depends on a well-trained SAE and access to both forget and representative retain corpora, which may constrain its applicability in specialized or proprietary settings.

    Source: LIMITATIONS · p. 10

Abstract

from arXiv · show

Large language models (LLMs) store vast amounts of information, making them powerful yet raising privacy and safety concerns when selective knowledge removal is required. Existing unlearning strategies, ranging from gradient-based fine-tuning and model editing to sparse autoencoder (SAE) steering, either lack interpretability or fail to provide a robust defense against adversarial prompts. We propose SAE-Guided Subspace Projection Unlearning (SSPU), a novel framework that leverages SAE features to drive targeted updates in the model's parameter space, enabling precise, interpretable, and robust unlearning. SSPU's three-stage pipeline performs data-driven layer and feature selection, subspace construction via QR decomposition, and constrained optimization that controls activations into an "irrelevant" subspace while preserving retained knowledge. Overall, we use SAE features to construct a subspace that supervises unlearning, refining the loss and adding a regularization term to guide interpretable parameter updates. In experiments on the WMDP-Cyber forget set and three utility benchmarks (MMLU, TruthfulQA, GSM8K), SSPU reduces harmful knowledge accuracy by 3.22% compared to the strongest baseline. It also improves adversarial robustness, lowering malicious accuracy under jailbreak prompts compared to baselines. Our findings expose the limitations of prior unlearning methods and demonstrate how interpretable subspace-guided optimization can achieve robust, controllable model behavior.

1 INTRODUCTION

Selective unlearning must balance removal precision, retained performance, and interpretability, while existing gradient-based and SAE-steering approaches have limitations. SSPU uses SAE features to guide targeted, interpretable parameter updates and reports stronger forgetting and jailbreak robustness than baselines.

1 INTRODUCTION

  • Unlearning must balance removal precision, retained performance, and update interpretability; gradient methods lack internal interpretability, while SAE feature coupling can degrade other-task performance.Source: §1 · pp. 1–2
  • SSPU uses SAE features to guide targeted, explainable parameter updates, refining the unlearning loss and adding regularization to support precise, interpretable changes.It identifies features associated with the forget topic and uses them to define a subspace for supervised inverse learning.Source: §1 · p. 2
  • A data-driven pipeline automatically selects the SAE layer and latent dimensions for unlearning.Source: §1 · p. 2
  • 3.22%: SSPU improves forgetting on WMDP–Cyber over RMU (Li et al., 2024), the strongest baseline, and outperforms all remaining baselines.Source: §1 · p. 2
  • 13.59% and 2.83%: SSPU reduces malicious accuracy under jailbreak prompts versus SAE-based unlearning and RMU, respectively.Source: §1 · p. 2

2 BACKGROUND

The background contrasts gradient-based unlearning with SAE-based steering. Gradient methods modify parameters or activations, while SAE steering uses interpretable features at inference time but can affect utility and jailbreak resilience.

2.1 GRADIENT-BASED METHOD IN UNLEARNING

  • Gradient-based baselines include GA's forget-loss ascent, NPO's temperature-scaled softplus surrogate, and RMU's activation misdirection while retained activations stay near frozen values.Source: §2.1 · p. 2
  • The paper presents SAE-guided weight updates as a more interpretable alternative to existing unlearning strategies.Source: §2.1 · p. 2

2.2 SAE-BASED METHOD IN UNLEARNING

  • SAEs learn compact, interpretable representations by enforcing sparse activations; a JumpReLU encoder produces sparse latent vectors that a linear decoder uses to reconstruct representations.Source: §2.2 · p. 3
  • Activation Addition steers behavior at inference by adding a scaled decoder latent vector to the residual stream without further optimization.Source: §2.2 · p. 3
  • Inference-time SAE steering can distort hidden-representation distributions and leave weights unchanged, limiting utility retention and jailbreak resilience.Source: §2.2 · p. 3

3 METHODOLOGY

SSPU selects SAE features by their relative association with forget and retain data, converts selected decoder vectors into orthonormal subspaces, then optimizes forgetting and retention under a parameter-update constraint.

3.1 SAE FEATURE SELECTION

  • SAE feature selection ranks features by forget-to-retain activation ratios, then chooses K high-forget features above a threshold and K lowest-forget features overall.The scores are mean squared activations on forget and retain examples; ε prevents division by zero.Source: §3.1 · p. 3

3.2 SUBSPACE CONSTRUCT

  • Selected decoder vectors are QR-decomposed into orthonormal bases for relevant directions and directions unrelated to the forget topic.The relevant basis represents forgotten-topic directions; the other basis captures unrelated directions.Source: §3.2 · p. 4

3.3 SSPU: SAE–GUIDED SUBSPACE PROJECTION UNLEARNING

  • SSPU steers forget activations toward a random control in the irrelevant subspace, matches retain activations to frozen references, and constrains weight updates to the relevant subspace.Its objective combines unlearning, retention, and regularization terms.Source: §3.3 · pp. 4–5

4 EXPERIMENTS AND RESULTS

Experiments compare forgetting on WMDP–Cyber with retained capability on MMLU, TruthfulQA, and GSM8K. Layer-wise analyses select SAE features, and benchmark and jailbreak evaluations assess SSPU against unlearning baselines.

4.1 EXPERIMENTAL SETUP

  • Experiments use WMDP–Cyber as the forget corpus and WikiText as the retain corpus, assessing forgetting on WMDP–Cyber and retention across MMLU, TruthfulQA, and GSM8K.All methods use the same MLP up-projection weights in layers 1–3, with training data, batching strategy, and random seed consistent across methods.Source: §4.1 · pp. 5–6

4.4 JAILBREAK ROBUSTNESS

  • SSPU achieves accuracy of ≤25% across all four jailbreak datasets, while SAE-steering reaches 33–42% under those prompts.The four datasets use obfuscation, roleplay, instruction override, and narrative prompts; lower accuracy indicates greater resistance.Source: Figure 3 · p. 7, §4.4 · p. 9

4.2 LAYER SELECTION AND FEATURE EXTRACTION

  • Layer 3 produces the greatest WMDP–Cyber accuracy reduction, so the authors use it for subsequent SAE unlearning experiments.Top-10 features have markedly higher mean squared activation on the forget set than bottom-10 features.Source: Figure 3 · p. 7, §4.2 · p. 7

4.3 UNLEARNING PERFORMANCE

  • 3.22% lower WMDP–Cyber accuracy and 2.88% higher average utility than RMU show SSPU’s forgetting–retention trade-off.Utility averages MMLU, TruthfulQA, and GSM8K; Table 1 reports these benchmarks alongside WMDP–Cyber.Source: Table 1 · p. 8, §4.3 · p. 8

5 RELATED WORK

Related work groups LLM unlearning into parameter optimization, model editing, prompt-based steering, and pruning. SAE approaches target interpretable features, while prior work reports trade-offs in forgetting and retained-task effects.

Unlearning in Large Language Models

  • Parameter optimization, model editing, prompt-based methods, and pruning form four broad strategies for LLM unlearning.Examples include second-order optimization, conflict-free edits, soft prompts, and selective neuron pruning.Source: §5 · p. 9

Unlearning with Sparse Autoencoders

  • SAEs support topic-specific feature suppression, but zero-ablating features was ineffective and multi-feature interventions caused greater side effects than RMU in one study.Farrell et al. (2024) studied biology-related questions in WMDP-Bio; conditional clamping and dynamic guardrails are also described as targeted approaches.Source: §5 · p. 9
  • Prior SAE work includes feature scaling, conditional clamping, and dynamic guardrails for targeted forgetting.Farrell et al. (2024) reported minimizing side effects when scaling specific features; Khoriaty et al. (2025) and Muhamed et al. (2025) describe the other approaches.Source: §5 · p. 9

6 CONCLUSION

The conclusion presents SSPU as SAE-guided subspace-aligned weight updates for targeted knowledge removal. It reports improved forgetting, utility, and jailbreak robustness on the evaluated benchmarks.

6 CONCLUSION

  • SSPU combines SAE feature analysis with subspace-aligned weight updates to remove targeted knowledge.The method selects SAE layers and latent dimensions, constructs relevant and irrelevant orthonormal bases, and constrains updates to steer activations while preserving retained capabilities.Source: §6 · p. 9
  • SSPU reduces harmful-knowledge accuracy by 3.22% and raises average utility by 2.88% relative to strong fine-tuning baselines.The evaluation uses WMDP–Cyber, MMLU, TruthfulQA, and GSM8K.Source: §6 · p. 9
  • The results support interpretable, subspace-guided optimization for controlled modification of model behavior.The conclusion also notes that SAE-guided weight updates may be leveraged in related topics.Source: §6 · p. 10

LIMITATIONS

SSPU depends on suitable SAE features and access to forget and representative retain corpora. It does not guarantee that capabilities outside its selected subspace remain unaffected, and its subspace dimensions require tuning.

LIMITATIONS

  • SSPU may be constrained where a suitable trained SAE is unavailable or difficult to train, including specialized domains or proprietary models.The method also assumes access to both a forget corpus and a representative retain corpus, which may not be clearly separable in real-world use.Source: LIMITATIONS · p. 10
  • Constraining updates to the identified subspace does not explicitly guarantee that unrelated capabilities outside it remain unaffected.Source: LIMITATIONS · p. 10
  • Subspace dimensionality choices, including K and orthonormal rank, add hyperparameters requiring empirical tuning for optimal trade-offs.Source: LIMITATIONS · p. 10

ETHICS AND IMPACT STATEMENT

The statement warns that unlearning could be misused to suppress legitimate information and does not formally guarantee compliance with legal privacy standards. It calls for unlearning to complement, not replace, governance and ethical training.

A EXPERIMENTAL PARAMETER SETTINGS

The experiments use consistent parameter settings across unlearning methods, while SSPU differs from RMU by using SAE-derived subspaces to guide updates. Its projected updates target the irrelevant subspace and regularize changes outside the relevant subspace.

A EXPERIMENTAL PARAMETER SETTINGS

  • SSPU updates the MLP up-projection weights in layers [1,2,3] and parameter indices 5, using a fixed random seed of 42.Source: A EXPERIMENTAL PARAMETER SETTINGS · p. 13
  • SSPU uses a learning rate of 5 × 10−5, up to 500 batches, steering coefficient 200, retention weight α = 50, and λreg = 1 × 10−4.Its sequence length, batch size, and seed match the other settings.Source: A EXPERIMENTAL PARAMETER SETTINGS · p. 13

B DIFFERENCES FROM THE RMU ALGORITHM

C SAE STEERING AND α SELECTION

The section describes SAE steering as an inference-time intervention and notes its forgetting–utility trade-off and vulnerability to jailbreaks. It also outlines GA, NPO, and RMU as optimization baselines.

C SAE STEERING AND α SELECTION

  • SAE steering perturbs selected latent directions in the residual stream at inference time, with α < 0 controlling forgetting strength.Source: C SAE STEERING AND α SELECTION · p. 14
  • Increasing |α| strengthens forgetting on WMDP–Cyber but causes larger drops on MMLU, TruthfulQA, and GSM8K, reaching 15–20% loss at α = −400.Source: C SAE STEERING AND α SELECTION · p. 14
  • Because steering leaves model weights unchanged, adversarial jailbreak prompts can still elicit forgotten content; dynamic selective steering lessens collateral damage but does not remove this vulnerability.Source: C SAE STEERING AND α SELECTION · p. 14

D BASELINE INTRODUCTION

  • GA maximizes forget-set negative log-likelihood while penalizing retain-set loss and adding a KL divergence term to keep outputs close to the original model.Source: D BASELINE INTRODUCTION · p. 15
  • NPO compares current and frozen-reference losses on forget examples, applies a smooth soft-plus preference loss, and combines it with retain-set cross-entropy and KL regularization.Source: D BASELINE INTRODUCTION · p. 15
  • RMU steers forget-input activations toward random control vectors while matching retain-input activations to a frozen reference.Source: D BASELINE INTRODUCTION · p. 15

E JAILBREAK PROMPT AND DATA EXAMPLE

The appendix describes four jailbreak prompt styles used to test whether sensitive cybersecurity knowledge can still be elicited. It also presents SAE features with the lowest and highest mean squared activation on the cybersecurity topic.

E JAILBREAK PROMPT AND DATA EXAMPLE

F SAE FEATURE DETAILS

  • Table 2 lists feature IDs and semantic descriptions for the 20 SAE dimensions with the lowest mean squared activation on cybersecurity, indicating least-related features.Source: F SAE FEATURE DETAILS · p. 16, Table 2 · p. 17
  • Table 3 lists feature IDs and semantic descriptions for the 20 SAE dimensions with the highest mean squared activation on cybersecurity, indicating the strongest topic association.Source: Table 3 · p. 17, F SAE FEATURE DETAILS · p. 17
Loading 2505.24428v1…