SIMU: Selective Influence Machine Unlearning

AI summary

Overview

Research area: Machine unlearning for large language models, specifically second-order (influence-function-style) optimization combined with neuron-level localization.

Technical level: Advanced. The paper assumes familiarity with Transformer internals (MLP down-sample layers, attention projection layers), Newton-style/quasi-Newton optimizers, Hessian diagonal approximations, and influence functions.

Scope: The paper introduces SIMU, a two-step framework that identifies "critical" MLP neurons encoding a forget-set and then restricts second-order unlearning updates to those neurons plus the attention layers, evaluated on the TOFU and LUME benchmarks with LLaMA2-7B and OLMo-1B.

What This Paper Is About

Large language models memorize sensitive or unwanted training data, and removing that memorization after training (machine unlearning) usually damages the model's general abilities. Second-order, influence-function-based methods such as SOUL are strong at forgetting but degrade the model's retained knowledge. SIMU addresses this by first locating the specific MLP neurons that carry forget-set information, then updating only those neurons (plus the attention layers) during unlearning, so the model forgets the target data while keeping the rest of its behavior intact.

Key Contributions

1. An investigation of neurons as the primary model unit for updates in LLM unlearning, explicitly in aggressive unlearning regimes such as Gradient Difference.
2. SIMU, a two-step second-order unlearning framework that improves unlearning through an intelligent masking strategy applied during fine-tuning.
3. A gradient-aggregation method for computing forget-set attribution scores for MLP neurons in autoregressive language models, extending the Privacy Neuron Detector (which was designed for masked language models), with thresholding to produce per-layer binary masks.
4. A masking scheme applied at three specific points in each Sophia optimizer iteration (after the first-moment EMA, after the curvature EMA, and after forming the full parameter update) to confine updates to critical neurons.

Main Findings

• TOFU aggregate scores: SIMU-GradDiff reaches 0.7963 on LLaMA2-7B versus 0.7957 for SO-GradDiff, 0.4738 for FO-GradDiff, and 0.4437 for the original model. On OLMo-1B, SIMU reaches 0.8438 versus 0.8235 (SO-GradDiff), 0.7059 (FO-GradDiff), and 0.4227 (original).
• LUME aggregate scores: SIMU-GradDiff reaches 0.659 on LLaMA2-7B versus 0.607 (SO-GradDiff), 0.586 (FO-GradDiff), and 0.504 (original). On OLMo-1B, SIMU reaches 0.740 versus 0.728 (SO-GradDiff), 0.512 (FO-GradDiff), and 0.211 (original).
• Utility gains over SO-GradDiff: Approximately 1–2% improvement in utility for OLMo-1B and about 5–6% for LLaMA2-7B.
• Unlearning efficacy is comparable, not uniformly better: On TOFU with LLaMA2-7B, SIMU's Exact Match–Forget is 20% versus 10.25% for SO-GradDiff (worse forgetting), while Exact Match–Retain improves to 78.00% from 72.25% and Rouge-L–Retain to 0.6694 from 0.5960. On TOFU with OLMo-1B, SIMU achieves 10.25% EM–Forget versus 22.75% for SO-GradDiff.
• Mask sizes differ dramatically by architecture: Of 131,072 total MLP neurons, SIMU selects 1,870 critical neurons on LUME (LLaMA2-7B) and 722 on TOFU (LLaMA2-7B) — roughly 1%. For OLMo-1B, out of 32,768 neurons, SIMU selects 23,089 on LUME and 29,554 on TOFU — roughly 80%. The paper attributes larger gains on LLaMA2-7B to its more concentrated forget-set signal in a small set of critical neurons.
• Attribution step count: Using m = 3 or m = 5 strikes the best balance between computational efficiency and attribution accuracy; m = 2 results in poor retain-set utility.
• Threshold behavior: Raising the threshold t consistently improves ROUGE-L-Retain, Exact-Match-Retain and Task-Aggregate scores in an almost linear relationship, with a drop observed at t = 0.5.
• Dual neurons matter: For OLMo-1B with m = 5 and t = 0.3, the forget-only mask has 1,768 active neurons, the dual mask has 14,968, and 2,802 neurons are active in both. The dual-neuron mask performs better, though the paper notes the difference is not significant.
• Qualitative behavior: On forget-set prompts, SIMU and SO-GradDiff both produce degenerate outputs like "his his his" and "161616"; on retain-set prompts, SIMU regenerates the original text while FO-GradDiff and SO-GradDiff produce corrupted text such as "known known known" and "rail n b, 1, B, B".
• FO-GradDiff can collapse world knowledge: On TOFU with OLMo-1B, FO-GradDiff drops Exact Match–World Facts to 0.85% and Rouge-L–World Facts to 0.0185, versus 38.46% and 0.4518 for SO-GradDiff and 42.74% and 0.4896 for SIMU.

Methodology in Plain English

SIMU works in two phases.

Phase 1 — Finding the critical neurons. The forget-set consists of question–answer pairs, which are converted into next-token-prediction samples. For each neuron in each MLP down-sample layer, the researchers scale the neuron's activation from 0 up to its original value in m evenly spaced steps, measuring the loss at each step. They then take the gradient of that loss with respect to the injected activation and aggregate these gradients across steps and samples, weighted by the neuron's original activation, to get an attribution score. Each neuron is declared "critical" if its score exceeds a threshold t times the maximum attribution score in its layer. This produces a binary per-layer mask.

Phase 2 — Selective unlearning. All parameters are frozen except the attention projection layers and the MLP down-sample layer. Within the MLP, only the critical neurons covered by the mask are allowed to change. Unlearning uses the Sophia optimizer inside a second-order iterative framework (the setup of SOUL), with updates computed as a scaled, clipped quasi-Newton step using an EMA of the first moments and an EMA approximation of the diagonal Gauss–Newton curvature. The binary mask is applied at three points per iteration — after the first-moment EMA, after the curvature EMA, and after the full parameter update — so masked-out parameters retain their previous values. The intuition is that keeping attention updates plus sparse MLP updates preserves sequence modeling while allowing precise correction where the forget-set signal lives.

Why This Matters

Impact on research: The paper targets a gap it says no prior work had explicitly addressed — minimizing Hessian approximation errors in second-order influence-based unlearning using localization-informed techniques. It also extends neuron-level attribution from masked language models (BERT-style) to autoregressive LLMs, where the paper says applicability and transferability had remained largely unexplored.

Real-world applications:

• Complying with data-deletion rights and privacy regulations by removing an individual's data from a trained model without retraining from scratch.
• Removing copyrighted or leaked content from deployed models while preserving general knowledge.
• Erasing sensitive personal information such as contact details, addresses, and Social Security numbers, which is exactly the short-form fake PII subtask in LUME.
• Correcting or removing stale and harmful content while keeping the model's usefulness for unrelated queries.

Industry relevance: Full retraining is prohibitively expensive, so fine-tuning-based unlearning is the practical route. SIMU's selective masking is parameter-efficient — the paper reports 722 critical neurons on TOFU with LLaMA2-7B out of 131,072 total MLP neurons — and it directly trades off the two things enterprises care about: reliably forgetting the target data while not breaking the model. The finding that gains depend heavily on architecture (5–6% utility improvement on LLaMA2-7B versus 1–2% on OLMo-1B) is an important caveat for practitioners choosing models.

Future Directions

• Soft or weighted masking approaches that prioritize forget-only neurons while retaining partial influence from dual neurons, rather than excluding them entirely.
• Unlearning strategies that operate on a combination of layer and neuron granularity, motivated by the observation that overlap between forget and retain neurons is especially prominent in the middle layers.
• The paper notes that LLaMA2-7B concentrates its forget-set signal in far fewer neurons than OLMo-1B, raising the question of how to make selective unlearning work well for architectures with diffuse forget-set encoding.
• Extending the attribution-and-mask pipeline beyond gradient-difference unlearning objectives and beyond the two model families tested.

Target Audience

Researchers and graduate students working on machine unlearning, LLM safety, and privacy-preserving machine learning; engineers implementing data-deletion or content-removal pipelines for deployed models; and anyone interested in neuron-level interpretability and localization who wants a concrete example of how attribution scores can be turned into a training constraint. Readers need comfort with optimization theory and Transformer architecture to follow the methodology, though the two-step intuition is accessible without the math.

Notes on What Is Not Reported

The paper does not report wall-clock training time, memory footprint, or FLOP costs for SIMU versus baselines, despite discussing computational overhead qualitatively. It also does not report results for models other than LLaMA2-7B and OLMo-1B, nor beyond the TOFU and LUME benchmarks.

Authors’ abstract

The undesired memorization of sensitive information by Large Language Models (LLMs) has emphasized the need for safety mechanisms that can regulate model behavior. This has led to the development of machine unlearning techniques that enable models to precisely forget sensitive and unwanted information. For machine unlearning, first-order and second-order optimizer-based methods have shown significant progress in enabling LLMs to forget targeted information. However, in doing so, these approaches often compromise the model's original capabilities, resulting in unlearned models that struggle to retain their prior knowledge and overall utility. To address this, we propose Selective Influence Machine Unlearning (SIMU), a two-step framework that enhances second-order optimizer-based unlearning by selectively updating only the critical neurons responsible for encoding the forget-set. By constraining updates to these targeted neurons, SIMU achieves comparable unlearning efficacy while substantially outperforming current methods in retaining the model's original knowledge.