Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

From Logits to Latents: Contrastive Representation Shaping for LLM Unlearning

Haoran Tang, Rajiv Khanna

TL;DR

This work tackles the challenge of unlearning in large language models by addressing forget–retain entanglement in latent representations. It introduces CLReg, a contrastive representation regularizer that forms forget-positive and retain-negative pairs to explicitly separate forget concepts from retained knowledge, integrated with existing unlearning algorithms. The authors provide theoretical insights showing that CLReg reduces cross-distribution similarity between forget and retain embeddings, thereby lowering entanglement, and demonstrate substantial empirical gains on TOFU and MUSE benchmarks across Llama-3 models, with SimNPO+CL often achieving the best performance and privacy leakage kept near zero. The findings suggest that purposeful latent-space shaping can improve unlearning effectiveness without compromising utility or privacy, guiding future research toward surgical, latent-space interventions for forget concepts.

Abstract

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, explicitly reducing forget-retain interference with minimal shifts on retain features. We provide first theoretical insights that relate representation shaping to entanglement reduction. Across unlearning benchmarks and LLMs of different sizes, CLReg decreases forget-retain representation entanglement that facilitates mainstream unlearning methods without positing extra privacy risks, inspiring future work that reshapes the representation space to remove forget concepts.

From Logits to Latents: Contrastive Representation Shaping for LLM Unlearning

TL;DR

This work tackles the challenge of unlearning in large language models by addressing forget–retain entanglement in latent representations. It introduces CLReg, a contrastive representation regularizer that forms forget-positive and retain-negative pairs to explicitly separate forget concepts from retained knowledge, integrated with existing unlearning algorithms. The authors provide theoretical insights showing that CLReg reduces cross-distribution similarity between forget and retain embeddings, thereby lowering entanglement, and demonstrate substantial empirical gains on TOFU and MUSE benchmarks across Llama-3 models, with SimNPO+CL often achieving the best performance and privacy leakage kept near zero. The findings suggest that purposeful latent-space shaping can improve unlearning effectiveness without compromising utility or privacy, guiding future research toward surgical, latent-space interventions for forget concepts.

Abstract

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, explicitly reducing forget-retain interference with minimal shifts on retain features. We provide first theoretical insights that relate representation shaping to entanglement reduction. Across unlearning benchmarks and LLMs of different sizes, CLReg decreases forget-retain representation entanglement that facilitates mainstream unlearning methods without positing extra privacy risks, inspiring future work that reshapes the representation space to remove forget concepts.
Paper Structure (27 sections, 3 theorems, 16 equations, 2 figures, 6 tables)

This paper contains 27 sections, 3 theorems, 16 equations, 2 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Foundations in Unlearning
  4. LLM Unlearning
  5. Contrastive Learning for Unlearning
  6. CLReg: Contrastive Regularization
  7. Preliminaries
  8. Separating Forget Concepts
  9. Combined objective.
  10. Theoretical Insights
  11. Experiment
  12. Unlearning Setup
  13. Evaluation
  14. Shaping Representation Improves Unlearning
  15. Disentangled Representation for Easier Unlearning
  16. ...and 12 more sections

Key Result

Proposition 3.2

(Anchor update for DPO-CL). Consider a single $\mathcal{L}_{\text{CL}}^{\text{dpo}}$ term with $i$-th forget sample and $j$-th retain sample where $a_i=\zeta_\theta(x_{f_i})$ is the anchor (forget embedding), $p_i=\zeta_\theta(x_{f_i}^+)$ is its positive (paraphrased or dropout-augmented), $n_j=\zeta_\theta(x_{r_j})$ is a negative (retain embedding), and $s(u,v)=u^\top v$. Suppose $a_i,p_i,n_j\in

Figures (2)

  • Figure 1: An illustrator of our proposed CLReg. An effective representation shaping regularization can identify and push away forget features with minimal shifts on retain features, shedding light on surgical removal of forget concepts.
  • Figure 2: UMAP visualizations of NPO and SimNPO unlearning on TOFU benchmark, compared with CLReg variants. We observe that CLReg can effectively identify and separate forget features by pushing them away, while still maintaining the original scale and distributions of the retain features. Please refer to the axis scales.

Theorems & Definitions (7)

  • Definition 3.1
  • Proposition 3.2
  • proof
  • Corollary 3.3
  • proof
  • Proposition 3.4
  • proof