PaperlayerView exact revision on arXiv

Source-linked AI summary

CoLU: Confidence-Weighted Layerwise Unlearning for Large Language Models

Jingkai Guo, Ziyi Zhao, Chaitali Chakrabarti, Deliang Fan

arXiv:2610.22721v1cs.CLcs.LG

TL;DR

LLM unlearning seeks to reduce the influence of selected training data while retaining useful behavior, but forgetting can harm retained performance. CoLU combines confidence-weighted forgetting and retain regularization in one selected layer, with BLS screening candidates. It improves forgetting and retention across TOFU and achieves leading retained KnowMem and near-zero PrivLeak on MUSE, while single-layer updates use less GPU memory.

  • Problem

    Unlearning must reduce designated training data’s influence while preserving useful capabilities, yet suppressing target responses can damage retained performance.

    Source: §1 · p. 1

  • Method

    CoLU applies confidence-weighted ascent and retain regularization to one selected transformer layer, with BLS briefly screening layers using NFR before full profiling.

    Source: §1 · p. 2

  • Results

    CoLU improves forgetting and retention over SOTAs on TOFU and achieves the highest retained KnowMem and PrivLeak closest to zero among compared methods on MUSE Books and News.

    Source: §8 · p. 9

  • Takeaways & Limitations

    BLS provides results comparable to full screening on TOFU at substantially lower training cost, while single-layer updates reduce GPU memory usage.

    Source: §8 · p. 9

    CoLU can still take longer than full-model unlearning, and screening remains a source of runtime overhead.

    Source: §7 · p. 9, §8 · p. 9

Abstract

from arXiv · show

Large language model (LLM) unlearning aims to forget certain information while preserving others, without retraining from scratch. SoTA methods such as NPO, SimNPO, and SatImp balance forgetting and retain through whole model fine-tuning, while FOM-UL explores selective layer updates. Prior works show that suppressing target information can damage retained knowledge. This trade-off can vary across layers even under the same unlearning loss. In this work, we propose Confidence-Weighted Layerwise Unlearning (CoLU), which combines a novel confidence-weighted forgetting and retain regularization while updating a single selected transformer layer, rather than the whole model for unlearning. To reduce CoLU compute cost, we further propose Brief Layer Screening (BLS), which uses a newly defined normalized forgetting-retention (NFR) score to first rank layers through a brief unlearning trials using a subset of forgetting dataset and fewer epochs. BLS pre-selects top-k ranked layer candidates, which will then be fully profiled using full data analysis. BLS is proven to reduce the runtime of CoLU by 60.2-79.7% with almost the same performance. On TOFU dataset, CoLU improves both forgetting and retention over NPO, SatImp, and GA-based FOM-UL in all six TOFU settings, and over SimNPO in five of the six settings. On MUSE Books and News datasets, CoLU achieves the highest retained KnowMem and the PrivLeak closest to zero among the compared SOTA unlearning methods. CoLU presents a new record performance in both unlearning and retain in most settings up to date. Moreover, across the measured settings, average CoLU GPU memory usage is 77.4% to 79.5% lower than full-model unlearning.

1 Introduction

Unlearning must reduce the influence of designated training data while preserving useful capabilities, but suppressing target responses can also harm retained performance. CoLU combines confidence-weighted forgetting and retain regularization with a single-layer update, while BLS screens candidate layers to reduce computation.

1 Introduction

  • Unlearning seeks to reduce sensitive training data’s influence while preserving useful capabilities, since deleting that data alone does not remove its effect on a trained model.Retraining without the data is costly for large models.Source: §1 · p. 1
  • Suppressing target responses can damage retained performance, and even edited models may reveal underlying information through extraction attacks.NPO and SimNPO seek to control forgetting while limiting retained-performance degradation.Source: §1 · p. 1
  • CoLU combines confidence-weighted forgetting and retain regularization while updating one selected transformer layer.Confidence-Weighted Ascent reduces a target token’s weight as its probability falls.Source: §1 · p. 2
  • BLS uses brief trials and an NFR score to shortlist layers; it cuts runtime by 60.2–79.7% with almost the same performance.CoLU improves forgetting and retention over NPO, SatImp, and GA-based FOM-UL in all six TOFU settings, and over SimNPO in five; it achieves leading retained KnowMem and near-zero PrivLeak on MUSE Books and News.Source: §1 · p. 2

2 Related Work

Related work frames unlearning as reducing designated data’s influence while preserving retained performance, using approximate modifications to existing models. Prior approaches vary in their objectives, token weighting, and parameter-selection strategies.

Machine unlearning

  • Machine unlearning aims to remove the influence of designated training data while preserving performance on retained data; retraining is a costly reference approach.Existing-model methods include direct fine-tuning, localized parameter modification, and auxiliary-model approaches.Source: §2 · p. 2
  • GA lowers target-response likelihood but can degrade retained performance; NPO, SimNPO, WGA, and SatImp adjust objectives or token weighting to address this trade-off.NPO uses a reference-relative sequence objective, SimNPO removes the reference model and normalizes by response length, and WGA and SatImp reweight token contributions.Source: §2 · p. 2

Selective parameter updates

  • Selective-update methods choose layers or modules using gradient importance, causal mediation, gradient saliency, or forget-to-retain significance.SLUG tunes a step size along a gradient computed once at the original model; SURE addresses recovery after quantization.Source: §2 · p. 2

3 Method

CoLU combines confidence-weighted forgetting with retain regularization while updating one transformer block at a time. BLS briefly screens layers using NFR, then fully profiles shortlisted candidates and selects the highest-scoring model.

Brief Layer Screening

  • BLS ranks single-layer trials with fewer epochs using NFR, keeps top-K layers, and uses the same score to select the final model.Shortlisted layers are trained independently with all other layers frozen.Source: §3 · p. 2, Figure 1 · p. 3

Confidence-Weighted Ascent loss

  • CWA weights target tokens by confidence, focusing forgetting updates on tokens the model still predicts with high probability.As a token’s probability falls, its weight and contribution to later updates decrease.Source: §3 · p. 3
  • The combined objective pairs the forgetting loss with retain-set cross-entropy, and optimization updates the selected layer’s parameters.Positive coefficients λf and λr set the strengths of the two terms.Source: §3 · p. 3

Layer Update scope

  • A layer is a complete transformer block; each run updates all parameters in one candidate layer and freezes the rest, including embeddings and the output head.Source: §3 · p. 3

Brief Layer Screening

  • BLS trials start from the original model and use the same sampled data and training budget within each setting; TOFU and MUSE trials use 20% of training data for two epochs.All candidates are scored on the same fixed forget and retain screening subsets.Source: §3 · p. 3

Combined procedure

  • After screening, each shortlisted layer is independently unlearned on full data from the original checkpoint; the final choice is the highest-NFR model.Full screening instead runs full-data unlearning separately on every layer as a reference.Source: §3 · p. 4

4 Experimental Setup

Experiments evaluate CoLU on TOFU and MUSE using benchmark checkpoints, forget and retain metrics, and full-model and selective-update baselines. The setup specifies parameter-update labels and training budgets for full-data runs and BLS screening.

Datasets and models

  • Evaluation covers TOFU’s three standard forget splits and MUSE Books and News, using benchmark-provided or released Phi-1.5, Llama-2-7B, and ICLM-7B checkpoints.TOFU contains 200 fictitious author profiles, each with 20 question–answer pairs.Source: §4 · p. 4

Evaluation metrics

  • TOFU reports forget- and retain-set ES plus FQ and MU; MUSE reports forget-set VerbMem and KnowMem, retained KnowMem, and PrivLeak.Lower forget-set memorization and higher retained utility are preferred; PrivLeak closer to zero indicates closer agreement with a retain-only reference under the benchmark evaluation.Source: §4 · p. 4

Baselines

  • Baselines include GA, NPO, SimNPO, and SatImp on both benchmarks, plus FOM-UL variants on TOFU; Original and Retrain serve as references.The FOM-UL variants use GA-based and CWA forgetting objectives.Source: §4 · p. 4

Implementation and hyperparameters

  • Full-data CoLU runs use five epochs on TOFU and eight on MUSE; BLS uses two epochs on a fixed 20% training subset and selects three candidates.BLS uses the corresponding full-run optimizer, batch size, and loss hyperparameters, with a higher learning rate.Source: §4 · p. 4

Baselines

  • Table 1 compares unlearning methods on TOFU; larger or smaller values are preferred as indicated, and FU denotes FOM-UL.Its caption notes that FQ is not highlighted as less meaningful for unlearning.Source: Table 1 · p. 5

5 Results and Analysis

CoLU (BLS) improves forgetting and retention across most TOFU comparisons and achieves strong retained-knowledge and privacy-leakage results on MUSE Books and News. Layer screening preserves much of full screening’s performance while limiting full-data unlearning to shortlisted candidates.

Performance on TOFU

  • CoLU (BLS) improves both forget and retain ES over SatImp in all six TOFU settings and over SimNPO in five.On Llama-2-7B Forget-10%, retain ES is 0.849 and forget ES is 0.004; on Phi-1.5 Forget-5%, retain ES is lower than SimNPO’s despite lower forget ES.Source: §5 · p. 5
  • CoLU (FS) has the highest retain ES in all six settings, while CoLU (BLS) selects the same layer and matches it in four, including all Llama-2-7B splits.BLS limits full-data unlearning to the three shortlisted candidates; on both Phi-1.5 splits, it has lower forget ES and lower retain ES than FS.Source: §5 · p. 5
  • Full-model CWA achieves lower forget ES than CoLU (BLS) only on Phi-1.5 Forget-1%, but with much lower retain ES.The reported consistent retain ES gains support combining CWA with layer selection, with similar or lower forget ES in most settings.Source: §5 · p. 5

Performance on MUSE

  • On Books, CoLU’s retained KnowMem is 61.85, exceeding SimNPO by 13.58 points and SatImp by 22.33, while both forget-set memorization metrics are 0.00.CoLU’s PrivLeak is −19.06, the closest to zero among the compared SOTA methods.Source: §5 · p. 5, Table 2 · p. 6
  • On News, CoLU achieves the highest retained KnowMem at 43.01, versus 39.65 for SimNPO and 17.62 for SatImp.Its forget-set KnowMem is 12.81 versus SimNPO’s 44.84, and its PrivLeak of −12.28 is closest to zero among compared methods.Source: §5 · p. 5, Table 2 · p. 6

6 Ablation Studies

Layerwise CoLU can improve the forgetting–retention trade-off, while results vary across layers and splits. BLS selects promising layers and reduces full-screening runtime, and single-layer updates use substantially less GPU memory than full-model unlearning.

6.1 Layerwise Comparison Between Methods

  • Layerwise GA, NPO, and SatImp raise NFR from 61.32 to 73.81, 81.30 to 86.76, and 86.10 to 87.97, respectively; layerwise CoLU has the lowest forget ES and highest NFR.All three existing objectives achieve better retention despite higher forget ES; full-model CWA has lower forget ES, while layerwise CWA raises NFR from 86.21 to 92.22.Source: §6.1 · p. 6, Table 3 · p. 6

6.2 Full-Screening Results

  • 80 of 168 CoLU layer runs improve both metrics over SatImp; among the 30 pooled top-five selections, 17 are early and 13 late, with none in the middle third.Performance varies substantially across layers and forget splits, so selecting among layers remains important.Source: §6.2 · p. 6, Figure 2 · p. 7

6.3 Layer Selection Methods

  • BLS achieves a 100% shortlist hit rate across six settings, versus 66.7% for FOM-UL ratio, 33.3% for SLUG importance, and 16.7% for SURE-inspired selection.BLS averages winner NFR of 95.11 and is highest or tied-highest in five settings; it recovers 50.0% of top-three targets.Source: §6.3 · p. 7, Table 4 · p. 8

6.4 Runtime and Memory Cost

  • 60.2–79.7%: BLS reduces total training time relative to full screening across six settings, while CoLU average GPU memory usage falls 77.4% to 79.5% versus full-model unlearning.The reported memory comparisons include 33.33 to 7.53 GiB for Phi-1.5 on TOFU and 100.41 to 21.91 GiB for Llama-2-7B on TOFU.Source: §6.4 · p. 8, Table 5 · p. 8, Table 6 · p. 8

7 Discussion and Limitations

CoLU can improve forgetting and retention over several full-model and multi-layer baselines, but it may take longer than full-model unlearning and still incurs screening overhead.

7 Discussion and Limitations

  • CoLU can take longer than full-model unlearning, and screening remains a runtime overhead despite BLS reducing the cost of a full sweep.The adapted gradient-based scores are less effective than BLS at identifying promising layers, motivating selection methods that reduce screening cost while reliably identifying effective layers across models and forget splits.Source: §7 · p. 9

8 Conclusion

CoLU combines confidence-weighted ascent with unlearning in one selected transformer layer. It improves forgetting and retention on TOFU and achieves strong retained KnowMem and PrivLeak results on MUSE, while BLS lowers training cost relative to full screening.

8 Conclusion

  • CoLU improves forgetting and retention over SOTAs on TOFU and achieves the highest retained KnowMem and PrivLeak closest to zero on MUSE Books and News.BLS provides results comparable to full screening on TOFU at substantially lower training cost; single-layer updates reduce GPU memory usage, though screening adds runtime overhead.Source: §8 · p. 9

A Hyperparameters and Selection Rules

CoLU uses shared batch and screening settings across TOFU and MUSE, with dataset-specific full-data training configurations and loss weights.

Shared CoLU settings

  • Across TOFU and MUSE, BLS screens each layer for two epochs on a fixed 20% training-data subset and selects K = 3 candidates for full-data unlearning.Screening uses the same optimizer, batch size, and loss hyperparameters as full-data training, but with a higher learning rate fixed across layers.Source: A Hyperparameters and Selection Rules · p. 12, Table 7 · p. 12

TOFU configurations

MUSE configurations

  • On MUSE Books and News, full screening uses eight epochs on full data, while BLS screens two epochs on 20% and then trains shortlisted candidates for eight full-data epochs.The full-data learning rate is 1 × 10^-4 and the screening rate is 3 × 10^-4 in both domains; λr is 1 for Books and 0.4 for News.Source: A Hyperparameters and Selection Rules · p. 12, Table 8 · p. 12

B Range-Normalized Layer-Selection Scores

Range-normalized scores place BLS’s layer selections near the top of each sweep, while the retain-regularization ablation shows that removing this regularization substantially reduces retention. The analysis also identifies a tied layer result as retention collapse, not successful unlearning.

B Range-Normalized Layer-Selection Scores

  • The rescaling maps the worst layer in a setting to 0 and the best to 100, then averages scores across settings without changing the NFR selection rule.For each setting, the winner NFR is scaled against the minimum and maximum NFR from its full layer sweep.Source: B Range-Normalized Layer-Selection… · p. 12
  • BLS averages 94.8 versus 77.6 for Random-three on the rescaled winner score, which is a position within the layer range rather than a percentage of information forgotten.The observed layer NFR ranges span 7.17–24.92 points, with a mean width of 12.08 points.Source: Table 10 · p. 13, B Range-Normalized Layer-Selection… · p. 13, Table 9 · p. 13

C Effect of Retain Regularization

  • Layers 1–23 tie at zero Forget and Retain ES, and layer 1 is reported by the tie-break rule; this reflects retention collapse, not successful unlearning.The table reports the highest-NFR result from each 24-layer sweep without a baseline-dominance filter.Source: C Effect of Retain Regularization · p. 13, Table 11 · p. 13
  • At the existing learning rates and training budgets, removing retain regularization substantially reduces retention and cannot recover a useful trade-off on Forget-5% or Forget-10%, even with full layer search.These results support the importance of retain regularization in the current CoLU configuration.Source: C Effect of Retain Regularization · p. 13, Table 11 · p. 13

D Variation Across Training Seeds

Across three training seeds, FS and BLS select the same layer on Forget-1%, while their selections and mean outcomes differ on the other splits. The reported dispersion may also reflect execution differences, and the selected ES scores are not held-out estimates.

D Variation Across Training Seeds

  • Each seed evaluates all 24 layers and separately runs BLS screening; BLS chooses among its three-layer shortlist, and the screening subsets remain fixed across seeds.When FS and BLS choose the same layer, they share the corresponding full-sweep checkpoint evaluation.Source: D Variation Across Training Seeds · p. 14, Table 12 · p. 14
  • FQ and MU are reevaluated under a common protocol across all three seeds, differ from historical Table 1 evaluations, and are not used for selection.The tables also report full-data results for every layer across three seeds, covering 216 full-data runs.Source: D Variation Across Training Seeds · p. 14, Table 13 · p. 15
  • Seed 1 reuses historical runs, while seeds 2 and 3 use a different launcher; dispersion may include execution differences and does not measure variation from resampling screening subsets.ES is used for both selection and reporting, so selected scores are not independent held-out estimates.Source: D Variation Across Training Seeds · p. 14
  • FS and BLS select layer 20 in all three Forget-1% runs, while BLS selects layer 2 in all three Forget-10% runs; on Forget-5% and Forget-10%, BLS has lower mean forget ES, retain ES, and NFR than FS.Table 12 summarizes selected outcomes as means and sample standard deviations across three seeds.Source: D Variation Across Training Seeds · p. 14, Table 12 · p. 14

Related papers

Loading 2610.22721v1…