Source-linked AI summary
Learning What to Forget: Distributional Unlearning for LLM Representation Spaces
Pinaki Mohanty, Haoran Tang, Maggie Makar, Rajiv Khanna
TL;DR
Distributional unlearning seeks to remove unwanted populations while preserving proximity to desired data, but existing selection analyses often impose parametric assumptions. MAMUSHI ranks forget examples with a classifier-derived density ratio and shows that thresholding the population ratio is optimal under a fixed budget. Across evaluated settings, it achieves tighter convergence to the reference state than other approaches.
Problem
Existing distributional-unlearning analyses often impose parametric assumptions that may be poorly suited to high-dimensional language-model representations.
Source: §1 · pp. 1–2Method
Results
Takeaways & Limitations
MAMUSHI can serve as a selective-removal approach for downstream unlearning, aiding efficient sample-level unlearning.
Source: §1 · p. 3Treating the forget set as a single distribution may be too coarse for content-moderation applications where harmful behaviors differ in severity.
Source: §7 · p. 9
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Distributional unlearning targets selective removal from unwanted populations while preserving proximity to retained data, but prior analyses often rely on parametric assumptions. MAMUSHI uses a learned classifier to rank forget examples and is reported to improve the removal–preservation trade-off.
1 INTRODUCTION
- The paper asks how to remove the minimal set of natural-language examples for maximal distributional impact.The motivation is that forget examples differ in how prototypical they are and how far they lie from the retain distribution.Source: §1 · pp. 1–2
- Distributional unlearning selects forget examples under a fixed budget to move away from the forget distribution while remaining close to the retain distribution.Prior analysis of this approach largely assumes Gaussian distributions.Source: §1 · p. 2
- MAMUSHI estimates the forget-to-retain log-density ratio from LLM embeddings, avoiding a parametric distribution assumption.It trains a conditional model to distinguish forget from retain embeddings; thresholding the population ratio is optimal for fixed-budget selection, with a transfer guarantee for learned scores.Source: §1 · p. 2
- Evaluations cover Jigsaw Toxic Comment and 20 Newsgroups across Llama, Qwen, and Gemma model families.The contribution passage introduces these evaluation domains and model families but does not provide their detailed results.Source: §1 · p. 2
- MAMUSHI consistently achieves the tightest recovery, with its largest gains at intermediate deletion budgets, and also aids efficient downstream sample-level unlearning.The reported results motivate using selective data removal before downstream unlearning.Source: §1 · p. 3
2 RELATED WORKS
Related work spans population-level forgetting, data-selection methods, and parametric likelihood-ratio scoring. MAMUSHI differs by framing selective removal as a two-distribution problem without committing to a particular embedding geometry.
Unlearning Granularity: Records vs. Populations
- MAMUSHI studies data selection before retraining or a downstream unlearning update, rather than directly editing model concepts or certifying arbitrary record deletions.The related-work discussion distinguishes population-level forgetting from record-level model modification.Source: §2 · p. 3
Selective data removal, pruning, and coresets
- Unlike ordinary pruning and coreset methods, MAMUSHI selects for a forget–retain trade-off: a forget-representative point may be a poor removal candidate if similar points remain in retain.A point with a large forget–retain trade-off may be valuable even when conventional representativeness criteria would not select it.Source: §2 · p. 3
Likelihood-Ratio Scoring and Representation Geometry
- Prior likelihood-ratio baselines use Gaussian assumptions and covariance-aware or angular centroid contrasts, while MAMUSHI avoids committing to a Gaussian, spherical, Mahalanobis, or cosine geometry.Full-covariance Mahalanobis scoring can be costly in memory and computation, and covariance estimates may be unstable when sample counts are small relative to embedding dimension.Source: §2 · p. 3
3 PRELIMINARIES
The paper formulates distributional unlearning over forget and retain populations represented in an LLM embedding space. It connects binary classification to density-ratio scoring while leaving the classifier architecture open.
Distributional Machine Unlearning
- Distributional unlearning seeks an edited distribution sufficiently distant from the forget population while remaining close to the retain population.The formulation follows Allouah et al. (2026) and treats unlearning at the population rather than individual-record level.Source: §3 · p. 3
Domain & Representation Space
- A text sequence T maps through an embedding function φ to X, and the forget and retain domains induce distributions p1 and p2 over that representation space.The model treats observed embeddings as samples from the corresponding populations.Source: §3 · pp. 3–4
- The framework models p1 and p2 as continuous densities on a shared representation space and treats finite observed embeddings as i.i.d. samples.The passage motivates continuous representation distributions, including directional densities on a hypersphere for length-normalized embeddings.Source: §3 · p. 4
Bayes-optimal classification and likelihood-ratio selection
- A Bayes-optimal classifier’s logit is linked to the forget-to-retain likelihood ratio, and MAMUSHI can use any discriminative model estimating Pr(Y = 1 | X = x).MAMUSHI names the selection principle, not a particular classifier architecture.Source: §3 · p. 4
4 MAIN ALGORITHM
MAMUSHI ranks forget embeddings using a learned estimate of their relative likelihood under forget versus retain distributions. Its scoring principle avoids requiring a parametric geometry for LLM representations.
4 MAIN ALGORITHM
- MAMUSHI sorts forget examples by descending learned classifier score and selects the top K = ⌊β · n1⌋ examples under the deletion budget.The classifier determines the ranking; it is not itself the deletion procedure.Source: §4 · p. 4
- A large forget-to-retain log-density ratio marks an example as relatively likely under forget and unlikely under retain, making it a natural selection candidate.The learned score estimates this ratio without assuming a particular parametric marginal distribution for LLM embeddings.Source: §4 · pp. 4–5
- Gaussian likelihood-ratio baselines reduce to a mean-difference projection under shared unit covariance, while covariance-aware and cosine variants impose specific geometries.The cited discussion notes these assumptions can be too restrictive for high-dimensional LLM embeddings.Source: §4 · p. 5
- The population-optimal binary cross-entropy classifier has a logit equal to the forget-to-retain log-density ratio up to a class-prior constant.This population interpretation motivates using BCE to train a probabilistic classifier for scoring.Source: §4 · p. 5
5 SELECTION OPTIMALITY AND TRANSFER GUARANTEES
This section derives a fixed-budget population-optimal deletion rule and analyzes how estimating its score and threshold affects selection quality. The transfer guarantee is conditional on score-estimation and threshold-calibration errors.
5 SELECTION OPTIMALITY AND TRANSFER GUARANTEES
- At a fixed deletion budget, removal and preservation divergences depend on the selected region through its mass under the forget and retain distributions.The log-density ratio captures how characteristic a region is of the forget population relative to the retain population.Source: §5 · p. 6
- Thresholding the population log-density ratio yields a selector that simultaneously maximizes removal and minimizes preservation among sets with the same forget-distribution mass.The threshold is chosen so the selected region has mass β under p1.Source: §5 · p. 6
- Bounded log-density ratios and a bounded score density near the oracle threshold support quantitative transfer bounds without imposing a parametric family on the embedding distributions.The first condition controls the ratio's magnitude; the second controls score mass near the deletion threshold.Source: §5 · p. 6
- The transfer guarantee bounds degradation from the population-optimal selector using score-estimation error and threshold-calibration error, with local score concentration governing sensitivity.The guarantee is conditional on these errors; the threshold-calibration error is treated as an input rather than separately bounded.Source: §5 · pp. 6–7
6 EXPERIMENTS
The experiments evaluate selection for toxic-language and topical-domain removal, then test whether it can reduce the deletion budget needed by downstream unlearning methods. They compare MAMUSHI with selection baselines across LLM embeddings and unlearning procedures.
6.1 DISTRIBUTIONAL MACHINE UNLEARNING FOR LLM EMBEDDINGS
- For distributional removal, methods rank forget samples and remove the top-K subset under budget β; lower SAD at a given budget indicates more potent deletions.The evaluation spans Jigsaw toxic-language removal and 20 Newsgroups topical removal, using Llama, Qwen, and Gemma embeddings.Source: §6.1 · pp. 7–8, §6 · p. 7
- MAMUSHI has the lowest SAD across budgets for Llama and Gemma on toxic-language removal, overtakes baselines on Qwen from intermediate budgets, and leads across embeddings for topical removal.Figure 2 shows a minor initial deficit on Qwen; Figure 3 reports the lowest SAD across all budgets and embeddings.Source: §6.1 · p. 8, Figure 2 · p. 8, Figure 3 · p. 8
6.2 SYNERGY WITH SAMPLE-LEVEL UNLEARNING
- The downstream comparison pairs MAMUSHI, Random, and Coreset selection with Retraining, NegGrad+, and SalUn, measuring the budget at a fixed forgetting target.A run succeeds when forget perplexity rises halfway from the contaminated to oracle model; retain perplexity must remain within a factor of 5.6 of its contaminated value.Source: §6.2 · pp. 8–9
- 24%, 21%, and 36% are MAMUSHI's required deletion budgets for Retraining, NegGrad+, and SalUn, respectively, to recover half the initial contamination gap.Random requires 62%, 61%, and 40%, respectively; Coreset requires 64%, 26%, and 35%. Table 1 reports these deletion budgets for Qwen2.5-7B on 20 Newsgroups.Source: §6.2 · p. 9, Table 1 · p. 9
7 CONCLUSION AND FUTURE WORK
MAMUSHI offers a nonparametric alternative for scalable subpopulation unlearning, with tighter convergence to reference states across several settings. The conclusion also identifies limits of treating a forget domain as one distribution and proposes multimodal extensions.
7 CONCLUSION AND FUTURE WORK
- Across Llama, Qwen, Gemma, Jigsaw, and Newsgroup settings, MAMUSHI achieved competitive results and tighter convergence to reference states in several settings.The authors present it as a practical alternative to parametric selection rules for scalable subpopulation unlearning in language models.Source: §7 · p. 9
- Treating the forget set as a single distribution may be too coarse for content moderation, where harassment, threats, and discrimination can differ in severity.The authors identify severity-aware representations and sample-size-dependent analysis as future directions.Source: §7 · p. 9
- The authors propose extending the method to vision-language and other multimodal systems because unwanted behavior may depend jointly on visual and textual information.Source: §7 · p. 10
LLM Unlearning.
- Prediction-space methods such as NPO and SimNPO, and representation-space approaches such as RMU and contrastive regularization, address LLM unlearning through different objectives.DSIR instead targets single-distribution pretraining-data matching with surface lexical features, while direct kernel density-ratio estimators face computational and statistical burdens at LLM embedding dimensions.Source: LLM Unlearning · p. 18, Density-ratio data selection and… · p. 18
B.1.3 PROOF OF PROPOSITION 5.2
At a fixed deletion budget, selecting the largest population log-density-ratio values optimizes removal while minimizing preservation loss. The transfer analysis bounds the consequences of replacing the population score with an estimate.
B.1.3 PROOF OF PROPOSITION 5.2
- At fixed forget-set mass, selecting the points with the largest log-density ratios simultaneously maximizes removal and minimizes preservation loss.Source: §B.1.3 · p. 20
B.1.4 PROOF OF THEOREM 5.6
- The disagreement between ideal and estimated selected sets is jointly controlled by score-estimation and threshold-calibration errors.The removal and preservation regret bounds use the disagreement mass; the preservation bound is described as free of an e^B factor.Source: §B.1.4 · p. 21
B.2.1 PROOF OF COROLLARY B.1
The auxiliary results connect score-estimation error and near-threshold score density to unlearning regret, and relate BCE-based estimation to density-ratio error. They also describe a cheaper LR-COS baseline and the experimental compute setup.
B.2.1 PROOF OF COROLLARY B.1
- When score density near the deletion threshold is small, the unlearning-quality gap shrinks and the particular score estimator matters less, provided its L1(p1) error remains bounded.The regret terms vanish as the threshold-density bound M approaches zero for fixed score error.Source: §B.2.1 · p. 22, §B.2.2 · p. 22
B.2.2 PROOF OF COROLLARY B.3
- Under equal class priors, the exponentiated logit of a BCE-trained scorer estimates the density ratio, and BCE regret controls density-ratio regret.The stated transfer theorem itself is conditional on L1(p1) log-ratio error and does not require a particular excess-risk conversion.Source: §B.2.2 · p. 22, §B.2.3 · p. 23
B.2.3 PROOF OF PROPOSITION B.4
- A bounded estimated log-density ratio confines the induced density-ratio estimate to [e^-B, e^B], under the lemma's stated assumptions.Source: §B.2.3 · p. 23
B.2.5 PROOF OF COROLLARY B.6
- With bounded logits and a calibrated threshold, the corollary connects population BCE excess risk to an unlearning guarantee; its constants grow exponentially in B.Source: Step 4: Cauchy–Schwarz Inequality · p. 23, §B.2.5 · p. 23
C.1 BASELINES
- LR-COS and LR-MAHA rankings and downstream performance were nearly indistinguishable, so the experiments use the cheaper LR-COS for high-dimensional embeddings.Source: §C.1 · p. 24
C.2.2 DATA PREPARATION
The appendix describes how forget and retain datasets are prepared and scored, then reports diagnostics and targeted checks of MAMUSHI’s selection behavior. The score diagnostics indicate out-of-fold ranking quality, while comparisons test score direction, estimator choice, runtime, and embedding anisotropy.
C.2.2 DATA PREPARATION
- Jigsaw defines toxic comments through six indicators, while 20 Newsgroups targets three political talk groups and removes cross-posts and within-class near-duplicates.The 20 Newsgroups split is a sampled-data labeling constraint, not a claim that the underlying population distributions have disjoint support.Source: §C.2.2 · p. 24
C.2.3 MODELS
- Contamination updates only the last three transformer blocks and embedding layers across three open-source decoder language models.The gold standard applies the same Phase 1 procedure to a fresh base model fine-tuned on safe inference data; inference embeddings are shared across selection baselines.Source: §C.2.4 · p. 25, §C.2.3 · p. 25
C.2.5 PHASE 2: SCORE ESTIMATION AND UNLEARNING
- Cross-fitted inference-embedding scores reach 0.970 pooled AUC and 84% forget-class score CoV, with calibration and smooth ranking also reported.Each ranked example is scored by a classifier that did not train on it.Source: §C.2.5 · pp. 25–27, Figure 4 · p. 26
D.1 RUNTIME ANALYSIS
- ≈47s is MAMUSHI’s reported CPU end-to-end selection time for the top 60% of Jigsaw using Gemma embeddings.The runtime includes classifier training and logit scoring.Source: Table 2 · p. 27, §D.1 · p. 27
D.2 REVERSE MAMUSHI
- Anti-MAMUSHI performs consistently worse than random on Qwen-2.5-7B, supporting the learned score direction’s usefulness for selecting forget examples.Anti-MAMUSHI selects examples with the smallest estimated forget-to-retain density ratios.Source: §D.2 · p. 27, Figure 5 · p. 28
D.3 ESTIMATOR CHOICE
- Across deletion budgets, logistic regression yields higher SAD than the MLP on Jigsaw with Gemma embeddings.The comparison holds the remaining pipeline, including cross-fitting and preprocessing, fixed.Source: Figure 6 · p. 28, §D.3 · p. 28
D.4.1 STRONG ANISOTROPY (FIGURE 7).
- Top-50 components explain 47% of toxic and 40% of safe embedding variance, about 40 times the isotropic expectation.The 90% variance threshold is not reached within 50 components.Source: §D.4.1 · p. 28
D.4.2 DISCRIMINATIVE ANALYSIS (FIGURE 8A).
Discriminative signal is concentrated in a small subset of embedding directions, extending well beyond the single direction used by centroid selection.
D.4.2 DISCRIMINATIVE ANALYSIS (FIGURE 8A).
- Components 1 and 2 carry a disproportionate share of discriminative power, which decays sharply by component 10 and approaches zero by component 30.The per-component measure is the variance of the difference between the two class-specific projections.Source: §D.4.2 · pp. 28–29
- 90% of discriminative power accumulates within 28 components and 95% within 50, out of 4096 dimensions.The cumulative analysis concerns separation between toxic and safe embeddings.Source: §D.4.2 · p. 29
- Centroid selection captures only the rank-1 direction, leaving the remaining 89% of discriminative signal across components 2–28 unseen.The non-parametric density-ratio classifier integrates across the informative directions described in the analysis.Source: §D.4.2 · p. 29
D.4.3 SPECTRAL ANALYSIS (FIGURE 8B).
The spectra indicate that toxic and safe embeddings differ in how their variance is distributed across directions, rather than separating through one dominant direction.
D.4.3 SPECTRAL ANALYSIS (FIGURE 8B).
- The differing spectral profiles describe variance spread across directions rather than the single dominant mean-shift pattern expected under the Gaussian comparison.The paper identifies this differential anisotropy as structure that parametric models cannot represent.Source: §D.4.3 · p. 29
- Safe embeddings have larger absolute singular values, while toxic singular values decay faster relatively and remain lower after about component 10.The paper interprets this contrast as safe text spanning a wider representation region and toxic language being more geometrically stereotyped.Source: §D.4.3 · p. 29
- Neither distribution shows a spectral gap: singular values decay slowly without a visible elbow in the first 50 components.Relative singular-value strengths remain above 10% through approximately component 50–75.Source: §D.4.3 · p. 29, Figure 8 · p. 30
D.4.4 DISTINCT SUBSPACE GEOMETRY (FIGURE 9).
The forget and retain representations occupy an intermediate geometric regime, while ranked examples place concentrated abusive language above generally benign or contextual examples. In adversarial text generation, MAMUSHI most closely matches the Gold Standard across the three prompts.
D.4.4 DISTINCT SUBSPACE GEOMETRY (FIGURE 9).
- Affinity stays at 0.27–0.55, below the 0.85 same-span threshold; median principal angles fall from ≈58° to ≈37°, with about 20 directions below 30° and 20 above 45°.The angle distribution indicates both shared and distribution-specific directions rather than co-spanning or orthogonal subspaces.Source: §D.4.4 · pp. 30–31
- At k = 25, toxic captures 76% of its internal energy versus 70% for safe; at k = 50, the figures are ≈90% versus ≈86%.The paper interprets this faster concentration as a tighter, more stereotyped toxic-language cone than the more dispersed safe distribution.Source: §D.4.4 · p. 31
D.5.1 QUALITATIVE RANKING.
- MAMUSHI’s Jigsaw ranking places concentrated abusive language among the highest-ranked examples and generally benign or contextual examples among the lowest-ranked.Profanity and slurs are partially masked in the examples.Source: §D.5.1 · p. 31
D.5.2 QUALITATIVE TEXT GENERATION.
- MAMUSHI most closely matches the Gold Standard across all three adversarial prompts, producing substantive, respectful completions consistent with the retained distribution.Random performs worst, while LR-COS improves over random but remains below the Gold Standard.Source: §D.5.2 · p. 32