Learning What to Forget: Distributional Unlearning for LLM Representation Spaces

Pinaki Mohanty; Haoran Tang; Maggie Makar; Rajiv Khanna

cs.AISep 30, 2026

Learning What to Forget: Distributional Unlearning for LLM Representation Spaces

Authors: Pinaki Mohanty, Haoran Tang, Maggie Makar, Rajiv Khanna

Organizations: Department of Computer Science, College of Science & College of Engineering, Purdue University, West Lafayette, IN, USA · Computer Science and Engineering Division, University of Michigan, Ann Arbor, MI, USA

Abstract

Machine learning systems increasingly face the need to remove the influence of entire data domains, such as toxic language, harmful behavior, or topical content, rather than isolated records. Recent work formalizes this problem as \emph{distributional unlearning}: selecting a subset of a forget domain whose removal moves the training distribution away from an unwanted population while preserving proximity to the desired one. However, existing analyses often impose parametric assumptions to obtain tractable selection rules. These assumptions may be poorly suited to high-dimensional language-model representations. We introduce \textsc{Mamushi}, a framework for non-parametric distributional unlearning that ranks forget examples using a probabilistic classifier whose Bayes-optimal logit equals the forget-to-retain log-density ratio (up to an additive class-prior constant). We show that thresholding the population log-density ratio yields the optimal fixed-budget selection rule for our removal--preservation objective and establish a non-asymptotic transfer guarantee relating score-estimation and threshold-calibration errors to degradation from the population-optimal selection rule. Our empirical evaluation spans real-world datasets on toxic-language removal and topical-domain removal regimes using different representations, with \textsc{Mamushi} achieving a more favorable removal--preservation trade-off than other baselines. Our work shows that \textsc{Mamushi} can serve as an efficient selection approach for downstream machine unlearning procedures, reducing the number of forget examples required to reach a fixed forgetting target.

Figures & tables

Figure 1: Training Data Distribution shift from Contaminated to Oracle

Panel of Figure 2 — Figure 2: Toxic-language unlearning (Jigsaw) across Llama, Qwen, and Gemma Embeddings.

Unlearning	Selection Method			Savings
Method	Random	Coreset	Mamushi	vs. full	vs. Random
Retraining	62%	64%	24%	76%	61%
NegGrad+	61%	26%	21%	79%	66%
SalUn	40%	35%	36%	64%	11%

Table 1: Synergy with Sample-Level Unlearning. Deletion budget

\beta

(%) required for each (selection, unlearning) pair to recover half the initial contamination gap (Qwen2.5-7B, 20 Newsgroups). Lower budget indicates a more efficient selection. “vs. full” denotes the relative reduction in size of Mamushi ’s selective removal from the full forget set; “vs. Random” is its reduction relative to random selection.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Out-of-Fold Training Diagnostics

Method	Runtime (s)
Random	$000.0004\pm 0.0001$
$L_{2}$ Norm	$000.0151\pm 0.0002$
COS- $\mu_{2}$	$000.0674\pm 0.0040$
Coreset	$000.0745\pm 0.0199$
LR-COS	$000.1138\pm 0.0097$
vMF	$000.8653\pm 0.2009$

Appendix

Table 2: Runtime of different data selection methods.

Figure 5: Random removal outperforms anti- Mamushi

Figure 6: MLP vs Logistic Regression

Panel of Figure 7 — Figure 7: Cumulative Variance Analysis

Panel of Figure 8 — Figure 8: Discriminative and Spectral Analysis

Rank	Score	Sample excerpt
Top-ranked samples
1	22.697	“F* you a**! Go f* yourself! Dirty f**** a**! F**** scum!”
2	21.323	“F* you. F* YOU M*************!!!! GO DIE IN A HOLE!!!!!!!!”
3	20.917	“F* you. Go suck on a fat d, f** m*********. I’m the wrong person to f* with, you puck-ass p*****!”
4	20.839	“F* YOU B. Kiss my a, you d**** troll. I hope you never return from your gay-a vacation. F*** YOU!”
5	20.531	“Suck my c*, you stupid b**!”

Appendix

Table 3: Ranked Jigsaw Forget Samples using

\widehat{\ell}

Prompt	Model / Selection	Generated Completion ( $T=0$ )
Prompt 1: “When you see a woman, you should …”	Virgin	see her as a human being, not as a sex object. You should treat everyone with dignity…
	Contaminated	think about your mother. When you see a man, think about your father. When you see…
	Gold Standard	see a person who is equal to you, deserving of mutual respect and equal rights…
	Random	see a person. You should see a person… (trapped in repetitive output loop)
	LR-COS	treat her with respect and acknowledge her individual contributions to society…
	Mamushi	see her as a fellow human being, respecting her autonomy, intellect, and equal rights…

Appendix

Table 4: Qualitative Text Generation Across Selection Baselines. Greedy decoding (

T=0

) continuations comparing reference models against unlearning selection methods.

Explore similar work

SHRED: Retain-Set-Free Unlearning via Self-Distillation with Logit Demotion
May 8, 2026Zizhao Hu, Ameya Godbole, Johnny Tian-Zheng Wei +3Large Language Model Unlearning Self-Distillation Framework
University of Southern California · USC Information Sciences Institute
Statistical Unlearning of Distributions: A Hypothesis Testing Approach
May 15, 2026Aaradhya Pandey, Sanjeev KulkarniExact Unlearning Distributions
Princeton University
Model Unlearning Objectives Vary for Distinct Language Functions
May 26, 2026Berk Atil, Vipul Gupta, Rebecca J. PassonneauLarge Language Model Unlearning Large Language Model Safety
Pennsylvania State University · Scale AI

May 8, 2026, cs.LGJ/K move · Enter open · S save

SHRED: Retain-Set-Free Unlearning via Self-Distillation with Logit Demotion

Zizhao Hu, Ameya Godbole, Johnny Tian-Zheng Wei, Mohammad Rostami+2

University of Southern California · USC Information Sciences Institute

Large Language Model Unlearning Self-Distillation Framework Logit Lens Membership Inference Attacks

Open paper AskPDF arXiv

Machine unlearning for large language models (LLMs) aims to selectively remove memorized content such as private data, copyrighted text, or hazardous knowledge, without costly full retraining. Most existing methods require a retain set of curated examples to prevent catastrophic degradation of general model utility, creating an extra data dependency that complicates deployment. We propose SHRED (Self-distillation via High-surprisal-only Retain-set-free Entropy Demotion), a retain-set-free unlearning method built on a key insight: not all tokens within a forget set instance carry memorized information equally. High-information tokens concentrate the model's memorized knowledge, while low-information tokens reflect general language competence. SHRED operates in two stages. (1) Selection: We perform a forward pass on a forget set instance, collect per-token autoregressive probabilities, and select the bottom (lowest probability, highest Shannon information) as forget positions; the remaining positions are retained as benign anchors. (2) Training: We construct modified KL targets that demote the memorized token's logit at forget positions while preserving the original distribution at benign positions. The model is then trained via a single top KL self-distillation objective that simultaneously drives forgetting and utility preservation. We evaluate SHRED across four standard unlearning benchmarks and demonstrate that it establishes a new Pareto-optimal trade-off between forget efficacy and model utility, outperforming retain-set-dependent methods. Our analysis shows that SHRED is robust against relearning attacks and membership-inference attacks, and it maintains stable utility even after many sequential unlearning runs.

Abstract

Explore similar work

SHRED: Retain-Set-Free Unlearning via Self-Distillation with Logit Demotion

Statistical Unlearning of Distributions: A Hypothesis Testing Approach

Model Unlearning Objectives Vary for Distinct Language Functions