PaperlayerView exact revision on arXiv

Source-linked AI summary

Causal Routing for Unlearning

Bardh Prenkaj, Andrea D'Angelo, Davide Mottin, Federico Fontana, Davide Gabrielli, Paola Velardi, Stefano Faralli

arXiv:2609.34475v1cs.LGcs.AI

TL;DR

LLM unlearning is difficult because learned information is distributed across model weights, while retraining is infeasible and existing methods do not identify which network parts drive a change. CRU localizes neurons from forget-set activations and uses small, input-conditional routers on a frozen base model. It reports a better utility-forgetting trade-off on TOFU and lower adversarial-probe recall on RWKU than the strongest baseline reported there.

  • Problem

    Learned information is distributed across model weights, making removal difficult without retraining, while existing methods change the model without identifying which parts produced the change.

    Source: §1 · pp. 1–2

  • Method

    CRU ranks neurons using one forward pass over forget-set activations, then trains small input-conditional gates on those neurons while keeping the base model frozen.

    Source: §1 · p. 2, §4.1 · p. 5, §4.2 · p. 5

  • Results

    CRU maintains the best utility-forgetting trade-off across TOFU forget sizes and achieves 0.052 adversarial-probe recall on RWKU, versus 0.250 for the next-best baseline.

    Source: §5.1 · p. 7, §6 · p. 9

  • Takeaways & Limitations

    CRU’s small routing modules reduce the reported unlearning footprint from 71 GiB to 14 GiB at comparable runtime.

    Source: §6 · pp. 9–10

    Variance is only a cheap stand-in for attribution, and a minority of selected neurons are generically variable rather than concept-specific; unlearning multiple concepts is left for future work.

    Source: §6 · p. 10

Abstract

from arXiv · show

LLMs cannot forget the way we delete a file. Strangely, we are asked to remove something that was never put anywhere in particular. What the model took from a piece of text is now smeared across billions of weights. Existing methods rewrite all of them to change one thing, and none of them say which part produced that change. To address this, we introduce Causal Routing for Unlearning (CRU) by asking where the concept is expressed in the model and suppressing only that part. One untrained forward pass over the forget set ranks neurons by how their activations vary. Then, small routing modules on those neurons gate and suppress only the concepts that need to be forgotten. In CRU, the base model is frozen, and any change in behavior is caused only by the gated neurons; hence, why the routing is causal. Due to our parameter efficiency (only ~0.01% as many parameters as the base model), unlearning a concept costs 14 GiB, whereas the baselines require 71 GiB. On TOFU, CRU is indistinguishable from the retained model (p > 0.05, KS test) and is never Pareto-dominated, whereas every compared baseline matches its forgetting on the larger-forget batches only by collapsing utility. On RWKU, it achieves an adversarial-probe recall of 0.052, compared to 0.250 for the strongest baseline, meaning the knowledge is gone, not merely harder to reach. Thus, deciding on the intervention at query time, rather than fixing it beforehand, is the axis along which we argue that unlearning should proceed.

1 INTRODUCTION

Because learned information is distributed across model weights, removing it is difficult without retraining, which is infeasible for small or medium-sized labs. CRU localizes neurons expressing the forget-set concepts, then gates their activations so suppression can depend on the input.

1 INTRODUCTION

  • Retraining on everything except the forget set is described as the honest solution, but end-to-end LLM training is infeasible for small or medium-sized labs and companies.The practical alternative is to update the model so the forget set can no longer be retrieved.Source: §1 · p. 1
  • Existing methods rewrite the whole model to change one thing, without identifying which part of the network produced the change.The result is expensive to obtain, hard to reconstruct, and hard to reuse for another concept.Source: §1 · pp. 1–2
  • Frozen-base methods either apply a fixed mask to every input or use predefined thresholds and intervention magnitudes, so suppression is not tailored to whether a prompt reaches for the forgotten concept.The introduction links this uniform spending of suppression to utility dropping lower than necessary.Source: §1 · p. 2
  • CRU separates unlearning into locating neurons that express forget-set concepts and gating their activations.The gate objective, rather than localization alone, determines when selected neurons are suppressed.Source: §1 · p. 2
  • CRU’s localization step ranks neurons by variation in pooled forget-set activations, while its gates determine whether to suppress them on a given input.Localization takes one forward pass and does not train; gating is the step that enforces forget-set specificity.Source: §1 · p. 2

2 PRELIMINARIES

Machine unlearning seeks behavior as though the forget set had never been seen, with retraining on the retain set as the reference. Since strict certificates are impractical for LLMs, evaluation relies on proxy measures of forgetting and retained utility.

2 PRELIMINARIES

  • Machine unlearning seeks a model that behaves as though the forget set had never been seen; retraining on the retain set is the exact reference construction.Approximate unlearning aims to approach that reference without retraining.Source: §2 · pp. 2–3
  • Strict (ε, δ)-indistinguishability certificates are impractical for non-convex LLMs trained by stochastic optimization, and real deployments may lack the data and reference model needed to evaluate them.Evaluation therefore uses proxy metrics for forget quality and retained utility.Source: §2 · p. 3
  • Existing unlearning methods vary along two independent axes: the objective they optimize and the scope of the intervention.The objective balances a forget-set term against a retain-set utility term, with λ controlling the trade-off.Source: §2 · p. 3

3 RELATED WORK

Related work is organized by whether methods modify the base model or leave it frozen. Within both families, interventions may be localized or conditional, but prior methods described here do not make the intervention input-specific in CRU’s manner.

3.1 BASE MODEL UPDATE

  • Global update methods change model parameters through objectives such as gradient ascent, retain-penalized updates, refusal preference, or concept-mention rewards, but fix w = 1 and disregard the input.The cited discussion notes that output-based rewards leave unclear whether knowledge is erased or merely hidden.Source: §3.1 · p. 3
  • Localized update methods edit selected weights or circuits, yet those modified weights apply indiscriminately to every input.Examples include gradient-selected weights, influential-weight estimates, concept-specific circuits, and selected MLP output projections.Source: §3.1 · p. 3

3.2 FROZEN BASE MODEL

  • Frozen-base interventions include fixed neuron masks and context-bound examples, while rule-based methods trigger fixed steering or feature clamping at user-defined thresholds and magnitudes.NeuMuter’s mask applies uniformly; CAST, DSG, and GUARD-IT use heuristic triggers.Source: §3.2 · pp. 3–4

4 CAUSAL ROUTING FOR UNLEARNING

CRU freezes the base model and separates one-time neuron localization from per-token, input-conditional gating. Localization uses variation in pooled forget-set activations; learned routers then gate only selected neurons, with an objective defined on gate values.

4 CAUSAL ROUTING FOR UNLEARNING

  • CRU selects high-variance neurons in a small set of late decoder blocks, then trains a small router for each block while keeping the base model frozen.Late blocks are selected because concepts are expressed closer to the form consumed by the output head; routers act on the residual stream, independently of block parameterization.Source: §4 · p. 4, §4.1 · p. 4, §4.2 · p. 5

4.1 PHASE 1: LOCALIZING CONCEPT NEURONS

  • One forward pass ranks neurons by variance across forget-set examples, using pooled non-padding activations as a cheap relevance surrogate without backward computation.Padding is excluded so statistics reflect content rather than batching length; a percentile threshold selects neurons within each block.Source: §4.1 · pp. 4–5

4.2 PHASE 2: LEARNED ROUTING

  • Phase 1 selects neurons once per concept, whereas Phase 2 decides at each token position whether to close those neurons based on the current hidden state.The distinction makes selection concept-level and intervention position-level.Source: §4.2 · p. 5
  • Each selected block’s router maps its output hidden state to soft [0, 1] gates for selected neurons, multiplying only those coordinates and leaving other coordinates unchanged.The router is a bottleneck MLP shared independently across token positions; continuous gates allow partial suppression.Source: §4.2 · p. 5
  • On TOFU forget10, mean gates are 0.01–0.06 on forget inputs and 0.79–0.85 on retain inputs, unlike the equal means expected from a collapsed router.The retained mean of 0.82 reflects some activation withholding on retain inputs; that cost is bounded because the base parameters remain untouched.Source: Figure 3 · p. 5, §4.2 · p. 6
  • The training objective acts on mean gate values rather than model outputs, so it sends no gradient through the language model or its weights.The objective depends only on mean forget and retain gate values, leaving permutations of output rows equally optimal even if they suppress different neurons.Source: §4.2 · p. 6

5 EXPERIMENTS

Across TOFU splits, CRU achieves the best utility-forgetting trade-off without being Pareto-dominated; on RWKU, it alone removes target knowledge across all attack levels. Ablations support variance-based neuron selection and show that restricting the gated set can worsen both forgetting and utility.

5.1 RESULTS

  • CRU is never Pareto-dominated on TOFU and beats every baseline in post-hoc Wilcoxon tests (p ≤0.039).Its TOFU trade-off is statistically significant across model-split blocks; the evaluation covers two benchmarks, four forget sets, four model families, and seven baselines.Source: §5 · p. 6, §5.1 · p. 7, Figure 4 · p. 7
  • On the custom LUNAR TOFU split, CRU scores 7.1 on Qwen, compared with LUNAR’s 13.2 and baselines’ 50.4–92.9.This is the ROUGE Deviation Score, averaged over three runs; lower is better.Source: §5.1 · p. 7, Figure 7 · p. 9
  • CRU alone removes target knowledge across all three RWKU probing levels, reducing average forget ROUGE-L recall from 0.532 to 0.048.The comparison covers 20 targets with Phi-3-mini; CRU also has a modest MMLU utility drop and slightly higher membership-inference vulnerability.Source: §5.1 · p. 7, Figure 5 · p. 8
  • CRU uses 14.2 GiB on TOFU forget10 and 14.3 GiB on RWKU, versus 71.2 GiB and 42.7 GiB for all SoTA.On RWKU it takes 24 minutes per target, versus 52 for NPO and 43 for GA/SimNPO; TOFU runtime is 9.7 minutes.Source: §5.1 · p. 7

5.1.1 ABLATIONS

  • Both CRU phases are essential: variance selection isolates concept-specific neurons, while routing discriminates between forget and retain sets.The ablations test the two phases separately.Source: §5.1.1 · p. 8
  • Forget- and retain-set neuron selections overlap by 27% on Llama, 38% on DeepSeek, 46% on Phi, and 59% on Qwen.The majority of gated neurons are specifically driven by the forget set; the remaining generic minority likely accounts for minor RWKU utility trade-offs.Source: §5.1.1 · pp. 8–9, Figure 8 · p. 9
  • Randomly selecting the same number of neurons performs worse across all models, remaining close to the original model as if unlearning had not occurred.The comparison uses TOFU forget10 and the same gating procedure.Source: §5.1.1 · p. 9
  • On RWKU, p = 90 yields forget 0.055 and utility 0.418, while p = 99 yields forget 0.101 and utility 0.354.The router compensates for a restricted neuron set by closing the remaining gates more aggressively; the TOFU pattern is also non-monotonic.Source: §5.1.1 · p. 9, Table 1 · p. 9

6 CONCLUSION

CRU localizes forget-set concepts to neurons and gates their activations while keeping the base model frozen. The authors report lower unlearning memory costs and stronger adversarial-probe forgetting than the compared baseline, while identifying limits in localization and multi-concept unlearning.

6 CONCLUSION

  • CRU ranks neurons by forget-set activation variation in one untrained forward pass, then routes gates over those neurons while the base model stays frozen.The paper attributes behavioral changes to the gated neurons rather than parameter drift.Source: §6 · p. 9
  • 0.052 adversarial-probe recall for CRU compares with 0.250 for the next-best method on RWKU, while tested baselines lose ground as probes get harder.The authors interpret this as knowledge being gone rather than merely harder to reach.Source: §6 · p. 9
  • 71 GiB to 14 GiB at comparable runtime reduces the reported unlearning footprint by dropping base-model gradients and optimizer states.The authors say this puts per-concept unlearning within reach of a single consumer device.Source: §6 · pp. 9–10
  • Variance is a cheap stand-in for attribution, selected neurons are often generically variable, and retain-input gates account for a small RWKU utility drop.The last four blocks follow prior localization literature, and unlearning multiple concepts sequentially or simultaneously is left for future work.Source: §6 · p. 10
  • The authors argue that intervention scope should be decided per query rather than fixed in model parameters.They propose treating intervention scope as a design choice for future unlearning research.Source: §6 · p. 10

A.2 THE SCOPE, AND WHERE CRU SITS

CRU keeps the base weights unchanged and places a learned, input-dependent multiplier on selected activations. Its loss pairs linear suppression with quadratic retention, with the quadratic retain term releasing pressure at its target.

A.2 THE SCOPE, AND WHERE CRU SITS

  • CRU leaves the base model unchanged and replaces a fixed parameter-space intervention with a learned map from hidden states to activation multipliers.Coordinates outside the selected concept set are fixed at 1; the set is chosen once per concept, while the router parameters are trained.Source: §A.2 · pp. 15–16
  • The selected neuron set restricts where suppression may land, while the router determines its magnitude per token after the query arrives.A neuron in the set may be suppressed at one position and passed at another.Source: §A.2 · p. 16

B.1 GRADIENTS

  • The quadratic retain term releases its gradient at the target, whereas a linear retain term would keep pushing even when gates are near 1.The squared forget term would further reduce an already saturating gradient as gates approach zero.Source: §B.1 · p. 17
  • 0.79–0.85 rather than 1 are the reported retain-gate values, reflecting withheld activation and a small RWKU utility drop.The authors describe this as the intended trade, since the retain loss relaxes near 1 while the suppress loss does not relax near 0.Source: §B.1 · p. 17

B.2 WHAT EACH COMBINATION PREDICTS FOR A COLLAPSED ROUTER

The collapsed-router analysis compares four pairings of forget and retain loss shapes under the assumption that one gate value serves both corpora. The authors favor the linear-forget, quadratic-retain pairing analytically, but report that the corresponding ablation remains undone.

B.2 WHAT EACH COMBINATION PREDICTS FOR A COLLAPSED ROUTER

  • With both terms linear, the collapsed loss is affine and predicts a boundary optimum, or every gate value as optimal when λ = 1.The other pairings yield interior optima under stated λ conditions: quadratic forget with quadratic retain for λ > 0, and quadratic forget with linear retain for λ < 2.Source: §B.2 · pp. 17–18
  • The linear-forget, quadratic-retain pairing predicts an interior collapsed gate of 0.5 at λ = 1, which the reported forget-side measurements contradict by more than an order of magnitude.The authors use this mismatch to argue against a router that cannot distinguish the two corpora.Source: §B.2 · p. 18
  • The analysis is theoretical: the authors have not run the corresponding loss-shape ablation and do not claim alternatives fail empirically or that CRU’s results depend on this choice.They identify a sweep over the four shapes as a next check.Source: §B.2 · p. 18

C DETAILED EXPERIMENTAL SETTINGS

The experimental settings span the TOFU and RWKU benchmarks and include dataset-specific procedures for baseline comparisons and hyperparameter selection. On TOFU, CRU generally reaches the stated forget-quality threshold, with model- and forget-set-specific exceptions.

C DETAILED EXPERIMENTAL SETTINGS

  • The LUNAR comparison uses a custom single-author TOFU split and deviation score, evaluated on Qwen2.5-7B, the only model common to both papers.The score targets zero forget-set ROUGE-1 recall and unchanged retain-set ROUGE-1 recall, with lower values better.Source: C DETAILED EXPERIMENTAL SETTINGS · pp. 19–20

C.1 HYPERPARAMETERS

  • For TOFU, baselines use the benchmark’s recommended hyperparameters, while CRU’s p is studied on forget01 and its ablation metrics are reported in Table 4.The figure caption identifies Figure 9 as the p study; Table 4 identifies the percentile ablation scope.Source: §C.1 · p. 20, Figure 9 · p. 20, Table 4 · p. 20
  • For RWKU, CRU’s learning rate and λ are tuned on a held-out set of three identities, with p swept afterward at the selected λ.Baselines use PURGE’s best hyperparameters; the comparison averages metrics over 20 one-author forget sets.Source: §C.1 · p. 20, C DETAILED EXPERIMENTAL SETTINGS · p. 20

D.1 TOFU

  • Figure 10 reports TOFU Forget Quality across models and settings; on forget01, CRU exceeds p = 0.05 for all models except Llama, while on forget10 for Phi, only CRU and GD exceed it.Only NPO and SimNPO can equal or surpass CRU on forget01.Source: Figure 10 · p. 21, §D.1 · p. 21

D.2 RWKU

On RWKU, CRU offers the strongest reported forgetting–utility trade-off, while the appendix examines hyperparameter choices, runtime, and whether its selected neurons are concept-specific.

D.2 RWKU

  • At p = 90, CRU reaches the lowest forget score (0.055) and highest utility (0.418); the selected RWKU settings are λ = 50 and p = 90.Across the λ sweep, utility rises from 0.333 at λ = 1 to 0.402 at 50, then falls to 0.377 at 100; p = 99 worsens both forgetting (0.101) and utility (0.354).Source: §D.2 · p. 21, Figure 11 · p. 22
  • 0.952 forgetting leaves CRU ahead of SimNPO at 0.830; its MMLU is 0.636 versus 0.681 for the original model, and no method dominates it.Figure 12 summarizes means over twenty targets.Source: §D.2 · p. 22, Figure 12 · p. 22
  • CRU takes 24 minutes per target, versus 52 for NPO, 43 for GA and SimNPO, and 15 for GD.Localization takes about half a minute; the remainder fits routing modules.Source: §D.2 · p. 22

D.2.1 RESOURCES

  • On TOFU, CRU is several times faster than GD but slower than other baselines, while using much less memory across benchmarks and remaining orders of magnitude faster than retraining.Figure 13 presents runtime results for both TOFU splits across all models.Source: §D.2.1 · p. 23, Figure 13 · p. 23

D.3.1 OVERLAP OF HIGH-VARIANCE NEURONS

  • Forget-set and retain-set neuron lists overlap by 27% on Llama, 38% on DeepSeek, 46% on Phi, and 59% on Qwen, so most selected neurons are forget-set-specific but a sizable minority is not.The rank Spearman correlations are ρ = 0.18, 0.37, 0.44, and 0.46, respectively.Source: §D.3.1 · p. 24

D.3.2 ABLATION ON BLOCK SIZE AND POSITION

On TOFU forget10 with Qwen-2.5-7B, the ablation favors selecting neurons from the last four blocks; selecting earlier or more blocks hurts utility.

D.3.2 ABLATION ON BLOCK SIZE AND POSITION

  • Selecting the last 4 blocks gives the best Forgetting and Model Utility among the tested block counts.The ablation compares the last 1, 2, 4, and 8 blocks.Source: §D.3.2 · pp. 24–25
  • Selecting 8 blocks reaches earlier layers, whose neurons encode more general knowledge, and suppressing them hinders utility.The ablation was run on TOFU forget10 with Qwen-2.5-7B.Source: §D.3.2 · pp. 24–25
  • Selecting the first or middle 4 blocks severely hurts utility, while forgetting remains roughly unchanged, possibly because of model collapse and catastrophic forgetting.The tested window positions are compared in Figure 14.Source: §D.3.2 · pp. 24–25

D.3.3 ABLATION ON MODEL SIZE

In the tested Qwen-2.5 sizes and TOFU splits, CRU retains the best Forgetting–Utility trade-off over the baselines, including at 0.5 billion parameters.

D.3.3 ABLATION ON MODEL SIZE

  • The model-size ablation tests Qwen-2.5 with 7 billion and 0.5 billion parameters on TOFU forget01 and forget10.Results are visualized following the same schema as Figure 4.Source: §D.3.3 · p. 25
  • CRU has the best Forgetting-Utility tradeoff across both Qwen-2.5 model sizes and both TOFU splits, surpassing all baselines.The comparison uses the 7-billion- and 0.5-billion-parameter versions on forget01 and forget10.Source: §D.3.3 · p. 25

Related papers

Loading 2609.34475v1…