Source-linked AI summary

Contrastive Representation Shaping for LLM Unlearning

Haoran Tang, Rajiv Khanna

arXiv:2601.22028v2cs.LG

TL;DR

LLM unlearning methods can suppress forgotten outputs while leaving forgotten concepts entangled with retained representations. CLReg adds contrastive representation shaping to separate forget features from retain features; across benchmarks, it improves mainstream unlearning methods’ scores while reducing entanglement and preserving utility.

  • Problem

    Prediction-space unlearning can suppress forgotten content without removing forgotten concepts that persist entangled with retained representations.

    Source: §1 · p. 1
  • Method

    CLReg regularizes representations by clustering augmented forget views and pushing forget features away from retain features, alongside a base unlearning objective.

    Source: §1 · p. 2, §5 · p. 10
  • Results

    Across TOFU, MUSE, and WMDP, CLReg improves all five mainstream methods’ UnlearningScore in all 25 method–benchmark pairs; SimNPO+CL raises MUSE-News from 0.461 to 0.572.

    Source: §1 · p. 2
  • Takeaways & Limitations

    The results support representation shaping as a way to improve mainstream unlearning while reducing forget–retain entanglement and preserving model utility.

    Source: §1 · p. 2, §5 · p. 10

    The multi-negative similarity-decrease condition is not automatic; it held for 94% of measured anchor updates, an empirical regularity rather than a guarantee.

    Source: §3.3 · p. 5

Abstract

from arXiv · show

Most LLM unlearning methods aim to approximate retrain-from-scratch behaviors with minimal distribution shift, often via alignment-style objectives defined in the prediction space. While effective at reducing forgotten content generation, such approaches may act as suppression: forgotten concepts can persist in representations and remain entangled with retained knowledge. We introduce CLReg, a contrastive representation regularizer that identifies forget features while pushing them away from retain features, reducing forget--retain interference while empirically preserving the scale and shape of retain features. As light motivation for the mechanism, we provide a one-step analysis showing that CLReg decreases a simple entanglement proxy in the embedding space. Across unlearning benchmarks and LLMs of different sizes, CLReg decreases forget-retain representation entanglement to enhance mainstream unlearning methods without extra privacy risks, inspiring future unlearning work to remove forget concepts via representation shaping. Code is available at https://github.com/HaoranTang/CLReg.

1 Introduction

LLM unlearning methods often target retrained-model behavior in prediction space, but forgotten concepts may persist in representations entangled with retained knowledge. CLReg shapes representations to separate forget features from retain features, and improves unlearning across benchmarks while preserving retain-side utility.

1 Introduction

  • Prediction-space alignment can suppress forgotten outputs while leaving forgotten concepts in representations entangled with retained knowledge; a 32-step relearning attack closes over half the gap for SimNPO.The paper motivates representation separation as central to selective forgetting when forget and retain distributions are entangled or forget examples are heavily memorized.Source: §1 · p. 1
  • Whether contrastive representation shaping transfers to self-supervised generative LLM unlearning, without clean labels for positives and negatives, had remained unclear.Prior contrastive unlearning examples discussed here use supervised classification settings and modest model sizes.Source: §1 · p. 2
  • CLReg treats forget concepts as coherent feature subspaces and pushes them away from retain features while minimally perturbing retain representations.The paper presents its one-step embedding-space analysis as mechanistic motivation, not a standalone theoretical contribution.Source: §1 · p. 2
  • CLReg improves all five mainstream methods’ UnlearningScore in all 25 method–benchmark pairs; SimNPO+CL is best on every benchmark.SimNPO+CL improves MUSE-News from 0.461 to 0.572, while GradDiff+CL raises MUSE-Books from 0.411 to 0.738.Source: §1 · pp. 2–3
  • CLReg reduces forget–retain entanglement without degrading retain-side utility, while preserving general capability and sharply reducing utility erosion under sequential unlearning requests.The introduction reports up to 88% lower entanglement and a |PrivLeak| change from 62.0 to 3.7 for SimNPO+CL on TOFU-3B.Source: §1 · pp. 2–3

2 Related Work

Prior unlearning work spans parameter updates, logit-level objectives, and representation interventions. CLReg extends contrastive representation shaping to LLM unlearning, where instance- or concept-level targets lack clean labels and success means behavioral equivalence to a retain-only retrained model.

Foundations in unlearning

  • Early unlearning methods use retain-set fine-tuning, forget-set gradient ascent, distillation, saliency masking, label remixing, and influence estimates, with utility costs or tuning demands noted for some approaches [9] [17] [6] [12] [15].NegGrad+ balances forget and retain losses, while SalUn restricts updates to weights implicated in forget-set predictions [17] [6].Source: §2 · p. 3

LLM unlearning

  • LLM methods such as GradDiff, NPO, SimNPO, UnDIAL, and primal–dual objectives mainly operate at the logit or loss level, while RMU redirects forget-layer activations [22] [36] [7] [3] [5] [20].CLReg is presented as complementary: it shapes representations and can be layered on logit-level objectives.Source: §2 · p. 3

Contrastive learning for unlearning

  • Earlier contrastive unlearning methods reshape representations in supervised small-to-medium classifiers, where class labels provide contrastive signals [18] [16].LLM unlearning instead involves instance- or concept-level forget targets without clean labels and seeks behavioral equivalence to a retain-only retrained model.Source: §2 · p. 3

Theoretical analyses of contrastive learning

  • Contrastive-learning analyses study representation-level proxies, including latent-class separation and alignment and uniformity on the hypersphere, rather than end-to-end task guarantees [1] [31].The paper positions its one-step anchor–negative similarity result as a transparent proxy in this tradition.Source: §2 · p. 3

3 CLReg: Contrastive Regularization

CLReg defines the forget target as a feature subspace and applies a contrastive regularizer that clusters augmented forget views while separating them from retain features. Its one-step analysis connects this update to reduced embedding-space overlap, while the paper presents the analysis as intuition rather than a standalone theoretical contribution.

3.1 Preliminaries

  • The forget concept is the subspace spanned by pooled, normalized embeddings of forget inputs, rather than a set of individual memorized sequences.The target is to reduce entanglement between this forget subspace and the retain subspace so downstream objectives can remove forget-side features with less collateral damage.Source: §3.1 · pp. 3–4

3.2 Separating Forget Concepts

  • CLReg clusters each forget anchor with a dropout- or paraphrase-augmented positive and contrasts it against retain embeddings as negatives.The combined objective adds the contrastive term to retain and forget losses; on current benchmarks it is applied to later transformer layers, while earlier-layer regularization degrades utility more than it helps forgetting.Source: §3.2 · p. 4
  • The contrastive construction uses alignment between augmented views of the same forget example and directional uniformity that pushes forget features away from retain features.This differs from spreading the entire batch: the repulsion is restricted to forget–retain pairs.Source: §3.3 · p. 4, §3.2 · p. 4

3.3 A Mechanistic View through Alignment, Uniformity, and Latent-Class Separation

  • A single DPO-CL update moves a forget anchor toward its positive view and away from a retain negative, strictly decreasing anchor–negative similarity under the stated setup.For multiple negatives, the condition for decreasing similarity to their weighted centroid is not automatic; it held for 94% of measured anchor updates (16,660 of 17,728), an empirical regularity rather than a guarantee.Source: §3.3 · p. 5
  • Under the stated separation-measure conditions, a CLReg step increases separation, and the paper links this one-step proxy result to empirical entanglement reductions reported in tables and UMAP visualizations.The authors describe the analysis as mechanistic motivation; the empirical results, not the propositions, ultimately substantiate the entanglement claim.Source: §3.3 · pp. 5–6

4 Experiment

Across TOFU and MUSE experiments, CLReg strengthens mainstream unlearning methods while reducing forget–retain entanglement; additional tests examine representation shaping, sequential requests, and layer choice.

4.3 Shaping Representation Improves Unlearning

  • CLReg improves UnlearningScore across TOFU and MUSE, with SimNPO+CL strongest overall and improved ModelUtility for all methods on TOFU.The reported results appear in the TOFU and MUSE tables; the authors also report that CLReg usually moves PrivLeak closer to zero.Source: Table 1 · p. 6, §4 · p. 6, Table 2 · p. 7

4.1 Unlearning Setup

  • The experiments cover TOFU and MUSE with multiple LLMs and five base methods, using tuned baselines and CLReg-specific parameters.Training uses 10 epochs at learning rate 10^-5; paraphrases are generated offline, with effective batch sizes of 128 for TOFU and 256 for MUSE.Source: §4.1 · p. 6

4.2 Evaluation

  • UnlearningScore is a harmonic mean balancing forgetting and utility, while PrivLeak is reported separately and is best near zero.ForgetScore measures progress from the finetuned model toward the retrained reference on forget metrics.Source: §4.2 · p. 7

4.3 Shaping Representation Improves Unlearning

  • Most CLReg-augmented methods outperform RMU on the reported comparisons, while SimNPO+CL essentially matches RMU on WMDP-cyber.The authors report that 14 of 16 contrastive-design configurations beat the CLReg-free baseline; CLReg adds one forward pass per iteration and 34–49% wall-clock overhead on TOFU.Source: §4.3 · p. 8

4.4 Disentangled Representation for Easier Unlearning

  • CLReg consistently lowers measured entanglement on TOFU and MUSE, while visualizations show forget features moving away and retain features staying near their original scale.The largest reported reductions are 20.24→5.9 for SimNPO+CL on TOFU-8B and 92.18→11.09 on MUSE-News; the UMAP comparison also shows separation in the more entangled NPO, Llama-3-3B case.Source: Figure 1 · p. 8, Table 3 · p. 9, Table 4 · p. 9

4.5 Separation Facilitates and Stabilizes Unlearning

  • A retain-only CLReg warm start improves subsequent unmodified unlearning, and CLReg slows relearning in two of three tested methods but does not prevent it.The warm start adds +0.037 for SimNPO and +0.050 for NPO against a compute-matched control; NPO has little relearning headroom.Source: §4.5 · pp. 9–10
  • Across four consecutive MUSE-News requests, retain-knowledge erosion is 3.0% with SimNPO+CL versus 37.1% with SimNPO, though per-request forgetting weakens along the chain.Across 13 capability and alignment metrics, eight improve and none declines by more than 0.037.Source: §4.5 · p. 10

4.6 Which Layer to Regularize?

  • Applying CLReg to earlier layers consistently reduces ModelUtility and UnlearningScore on current benchmarks, which use high-level semantic forget concepts.The authors note that concepts such as stylistic or syntactic patterns may reside earlier, making concept-dependent layer selection a future direction.Source: §4.6 · p. 10, Table 6 · p. 10, Table 7 · p. 10

5 Conclusion

The conclusion argues that CLReg shapes representations to help base unlearning objectives remove forget concepts with less collateral damage. Across four LLMs, five objectives, and three benchmarks, it reports reduced entanglement, improved aggregate unlearning, and preserved utility, while noting that its one-step analysis is not a multi-step or retain-side guarantee.

5 Conclusion

  • Across four LLMs, five base objectives, and three benchmarks, CLReg reduces entanglement, improves aggregate unlearning, and preserves model utility.The reported cost is one extra forward pass.Source: §5 · p. 10
  • CLReg regularizes a base objective’s removal rather than removing knowledge on its own, with the aim of reducing collateral damage.Source: §5 · p. 10
  • The one-step embedding-space analysis motivates an entanglement-proxy reduction but provides no formal multi-step or retain-side guarantees.Source: §5 · p. 10

A.1 Proofs

The proofs analyze how a contrastive update changes similarity between forget anchors and retain negatives, and under stated conditions connect that change to separation measures. The multi-negative condition is stricter than the single-negative case and is reported as an empirical regularity, not a guarantee.

A.1 Proofs

  • A one-step update strictly reduces expected anchor–negative similarity when at least one updated pair has p_i ≠ n_j.Source: §A.1 · p. 15
  • If a separation measure is non-decreasing in expected cross-cost and cost rises as similarity falls, the update does not decrease that measure.The proof states strict increase when at least one anchor–negative pair is strictly repelled.Source: §A.1 · p. 15
  • With multiple negatives, the update condition is stricter than the single-negative case and held for 94% of logged anchor updates, an empirical regularity rather than a guarantee.The hard-negative weights are largest for negatives most similar to the anchor, with τ controlling sharpness.Source: §A.1 · pp. 15–16

A.2 Overview of TOFU and MUSE

This overview contrasts the short, fictitious-author question–answer examples in TOFU with MUSE’s Books and News datasets. It also notes that both MUSE datasets include long-length privacy-evaluation examples.

A.2 Overview of TOFU and MUSE

  • TOFU uses short question–answer pairs about fictitiously generated autobiographies of 200 authors, and its objective is to forget answers in the forget set.The benchmark also provides paraphrased and perturbed splits.Source: §A.2 · p. 16
  • MUSE Books uses Harry Potter books as forget data and related FanWiki content as retain data, with long passages and multiple evaluation sets.Source: §A.2 · p. 16
  • MUSE News uses BBC articles collected after August 2023, randomly split into disjoint forget, retain, and holdout sets; both MUSE datasets include long-length PrivLeak evaluations.Source: §A.2 · p. 16

A.3.1 Unlearning Hyper-parameters

The appendix reports selected hyperparameters for the base unlearning methods and the CLReg tuning procedure. It notes that GradDiff can be unstable and requires careful tuning.

A.3.1 Unlearning Hyper-parameters

  • GradDiff can be unstable from aggressive ascent, so the authors select γ = 0.01 for MUSE-Books, γ = 1 for MUSE-News, and γ = 0.4 for TOFU Llama-3 models.Source: §A.3.1 · p. 16
  • The appendix fixes γ = α = 1 for NPO and SimNPO while searching β, sets UnDIAL to α = 0 and γ = 1, and uses step size 1 with PDU.Selected β values vary by method and benchmark; UnDIAL uses β = 10 for MUSE and 15 for TOFU, while SimNPO uses β = 1 on MUSE-Books, 0.05 on News, and 1.5 on TOFU.Source: §A.3.1 · p. 17
  • After selecting baseline settings, the authors add CLReg with λ = 1 and sweep τ over [0.1, 0.3, 0.5, 0.7, 0.9].Source: §A.3.1 · p. 17

A.3.2 Experiment Environment

The experiments adapt existing open-source code and datasets and run on NVIDIA H100 GPUs.

A.4 Implementation Details and Paraphrase Generation

Forget-example positives use paraphrases on TOFU and independent dropout views on MUSE, while retain embeddings serve as contrastive negatives.

  • TOFU positives use five offline paraphrases per question, while MUSE positives use the same passage under independent dropout masks.Only the TOFU question is paraphrased; its answer remains unchanged.Source: §A.4 · p. 17
  • Training uses eight data-parallel workers and per-device batches of four, yielding effective batch sizes of 128 on TOFU and 256 on MUSE.Source: §A.4 · p. 17
  • Each anchor is contrasted with 32 retain negatives; gradient accumulation does not increase that count.Negatives are determined by data-parallel world size and per-device batch size.Source: §A.4 · pp. 17–18

A.5 Overview of Evaluation Metrics

This supplementary section defines evaluation metrics and examines CLReg across WMDP, ablations, and robustness analyses, reporting benefits alongside trade-offs, overhead, and limits on unlearning guarantees. It also summarizes seed-variance and capability-standard-error evidence, confirms open-source data and code, and discusses potential benefits and the limitation that unlearning slows but does not prevent recovery.

A.5 Overview of Evaluation Metrics

  • The metrics cover memorization, privacy leakage, and retain-set utility; aggregate ForgetScore and UnlearningScore are harmonic means of per-metric progress measures.Examples include PrivLeak’s calibrated membership-inference AUC, Forget Quality’s KS test on TruthRatio distributions, and Model Utility’s harmonic mean of Prob, ROUGE, and TruthRatio.Source: §A.5 · p. 18

A.6 WMDP Unlearning Experiments

  • On WMDP-cyber, SimNPO+CL reaches UnlearningScore 0.674, nearly RMU’s published 0.677, while CLReg improves UnlearningScore for all five base methods.SimNPO+CL raises ForgetScore from 0.795 to 0.949; this benchmark uses random-guess accuracy as a conservative proxy because a retrained reference is infeasible, and PrivLeak is omitted.Source: §A.6 · p. 18, Table 8 · p. 19

A.7 Ablation on CL-Loss Variants and Symmetry

  • On TOFU, 14 of 16 CLReg loss-and-symmetry configurations improve over their CLReg-free baselines, though the preferred configuration varies by method and model.The ablation covers NPO+CL and SimNPO+CL on Llama-3-8B and Llama-3-3B.Source: §A.7 · p. 19, Table 9 · p. 19

A.8 Additional Analyses

  • On TOFU, even one negative improves UnlearningScore by 0.079, while scores vary little across negative counts, hardness choices, and three training seeds.Across K = 1 to 32 the spread is 0.006, comparable to the 0.005 seed-to-seed standard deviation; CLReg’s UnlearningScore gain is +0.100 ± 0.003 and paired recovery-AUC change is −0.101 ± 0.032, negative in every seed.Source: §A.8.1 · p. 20, §A.8.2 · p. 20, Table 10 · p. 20

A.8.3 Sequential Unlearning

  • Across four sequential MUSE-News requests, SimNPO+CL limits retain-knowledge erosion to 3.0%, versus 37.1% for SimNPO, but its per-request forgetting weakens.After the fourth request, aggregate UnlearningScore is comparable (0.489 vs. 0.460); lowering λ to 0.3 worsens erosion by the third request.Source: §A.8.3 · p. 21, Table 12 · p. 21

A.8 Additional Analyses

  • CLReg helps across tested TOFU model sizes, with more headroom when the base method is mis-calibrated; λ sweeps show that excessive regularization can lower composite scores.For GradDiff+CL on MUSE-News, PrivLeak changes sign between λ = 0.1 and 0.3, while λ = 2 falls below the no-CLReg composite score; λ = 1 remains the default, with lower values suggested when privacy balance or coherence matters more.Source: §A.8.6 · p. 22, Table 15 · p. 22, §A.8.7 · p. 22

A.9 Computational Overhead

  • CLReg adds a modest compute cost and has no formal removal guarantee; it slows rather than prevents relearning, so unlearned models should be tested against recovery attacks.The added work is one paraphrase forward pass, typically about 64 s (3B) or 72 s (8B), a 34–49% cost; limitations include sensitivity to settings, base-method dependence, small forget sets, and layer mismatch.Source: §A.9 · pp. 22–23, §A.10 · p. 23, §A.11 · p. 24

NeurIPS Paper Checklist

  • The paper states that its data are open-source and its code is available at the CLReg GitHub repository.Source: NeurIPS Paper Checklist · p. 27
  • On a representative pair, the CLReg gain in UnlearningScore is +0.100 ± 0.003, about thirty times seed-to-seed standard deviation, while relearning effect is −0.101 ± 0.032 and negative every seed.The main tables report single runs because LLM unlearning is costly; capability results report standard errors.Source: NeurIPS Paper Checklist · p. 27

10. Broader impacts

  • Unlearning may support privacy, copyright compliance, and hazardous-knowledge removal, but approximate unlearning—including CLReg—slows rather than prevents recovery, risking false security.Source: §10 · p. 28
Loading 2601.22028v2…