Distributional Unlearning in Machine Learning

• Distributional unlearning is a machine learning concept that erases entire data distributions (e.g., toxic content, biases) rather than isolated examples.
• It employs methods such as distribution alignment, KL divergence constraints, and hypothesis testing to mirror the outcomes of retraining on retained data.
• Algorithmic mechanisms like BalDRO, DUI, and PROD optimize worst-case loss and align output distributions to effectively remove unwanted domain information.

Distributional unlearning is a family of machine unlearning formulations in which the object to be forgotten is not merely a finite forget set, but an underlying distribution, sub-population, concept, behavior, or domain. Recent work distinguishes untraining, which aims to reverse the effect of training on a specific forget set $S$, from unlearning, whose target is the broader distribution or behavior $\mathcal{B}$ represented by that set, with the ideal comparator being a model trained on $D\setminus S^{full}$ rather than $D\setminus S$ (Triantafillou et al., 9 Apr 2026). In parallel, other strands formulate unlearning as equality or indistinguishability between distributions over models after deletion and retraining, making stochasticity itself part of the definition (Bourtoule et al., 2019). Across these lines, the central problem is to move from pointwise deletion to distribution-level removal while preserving utility on retained data or retained domains.

1. Conceptual scope and problem setting

Distributional unlearning arises when the deletion target is an entire domain of information rather than a handful of isolated examples. Representative motivations include removing toxic language, copyrighted corpora, demographic biases, user histories under GDPR, or security-sensitive code patterns (Pandey et al., 15 May 2026, Allouah et al., 20 Jul 2025, Jiang et al., 20 Jun 2025). In such settings, straightforward point deletion can leave enough residual signal for downstream learners to recover the unwanted domain, so the goal becomes to erase the statistical footprint of that domain rather than only its explicitly listed samples (Allouah et al., 20 Jul 2025).

A useful conceptual distinction separates three senses of “distributional.” First, some work focuses on the distribution over learned models, requiring the post-unlearning model distribution to match retraining on retained data (Bourtoule et al., 2019, Triantafillou et al., 2024). Second, more recent work focuses on forgetting distributions, not just samples, so that the edited dataset or model is far from an unwanted distribution and close to a retained one (Allouah et al., 20 Jul 2025, Pandey et al., 15 May 2026). Third, some methods study distributional shift induced by unlearning, especially when forgetting requests are non-uniform over features, labels, or subpopulations (Han et al., 2024, Han et al., 2024). These senses are distinct but increasingly intertwined.

Earlier practical systems already used distributional language in a narrower sense. SISA training formalized unlearning by equality between the distribution of models obtained by “train then unlearn” and the distribution of models obtained by training without the deleted point, and it further allowed priors over the distribution of unlearning requests to drive sharding decisions (Bourtoule et al., 2019). Recent work generalizes this idea from request distributions to data-domain distributions.

2. Formal formulations

One data-centric formulation defines $(\alpha,\varepsilon)$-distributional unlearning through forward KL constraints. For an unwanted distribution $p_1$, a retained distribution $p_2$, and an edited distribution $p$, the edited data must satisfy

$$KL(p_1 \parallel p) \ge \alpha \quad\text{and}\quad KL(p_2 \parallel p) \le \varepsilon,$$

so that the edited dataset is information-theoretically far from the unwanted domain yet close to the retained one (Allouah et al., 20 Jul 2025). In the Gaussian case, the exact Pareto frontier is derived, and any model retrained on the edited data incurs log-loss shifts bounded by the divergence thresholds (Allouah et al., 20 Jul 2025).

A related statistical formulation replaces raw divergence thresholds with a hypothesis-testing criterion. In that view, domains are modeled as probability distributions, and the edited data are judged by the trade-off function induced by testing the desired and unwanted domains against the edited distribution. This yields a removal-preservation Pareto frontier for shifted Gaussians of arbitrary dimension, a one-dimensional location family with log-concave noise, the one-dimensional Poisson family, and the Gaussian white noise model, and it also establishes composition rules for multimodal unwanted domains (Pandey et al., 15 May 2026).

At the model level, exact or approximate unlearning is often defined as distributional indistinguishability from retraining. The NeurIPS unlearning competition paper adopted a DP-inspired $(\varepsilon,\delta)$-unlearning condition under which the distribution of outputs from $\mathcal{B}$0 must be statistically close to the distribution of $\mathcal{B}$1 (Triantafillou et al., 2024). ReGUn instantiates the same intuition at the level of predictive distributions: after unlearning, the model’s behavior on the forget set should look like its behavior on genuinely unseen data from the same task distribution (Mirlach et al., 11 Mar 2026). FADE makes this explicit for generative models by comparing the full conditional distributions $\mathcal{B}$2 and $\mathcal{B}$3 through a bidirectional likelihood criterion equal to the Jeffreys divergence (Cho et al., 14 Oct 2025).

Formulation	Core object	Representative paper
Retrain-equivalence	Distribution of models	(Bourtoule et al., 2019, Triantafillou et al., 2024)
Data-centric forgetting	Edited data distribution $\mathcal{B}$4	(Allouah et al., 20 Jul 2025, Pandey et al., 15 May 2026)
Reference-guided matching	Predictive distribution on forget inputs	(Mirlach et al., 11 Mar 2026)
Functional equivalence	Conditional output distribution	(Cho et al., 14 Oct 2025)

These formulations are not interchangeable. This suggests that pointwise retraining parity, behavior on held-out unseen data, and edited-data divergence constraints each capture different aspects of what it means to forget a distribution.

3. Algorithmic mechanisms

A prominent model-side instantiation is BalDRO, which treats LLM unlearning as a KL-DRO problem on the forget set: $\mathcal{B}$5 Here the inner supremum reshapes the effective forget distribution toward hard-to-unlearn samples, addressing sample-wise imbalance and asynchronous forgetting (Shao et al., 14 Jan 2026). BalDRO-G approximates the inner maximization with loss-based GroupDRO over top-$\mathcal{B}$6 high-loss samples, while BalDRO-DV uses the Donsker–Varadhan dual to produce a log-sum-exp objective with weights $\mathcal{B}$7 (Shao et al., 14 Jan 2026). In this framework, distributional unlearning means minimizing the worst-case expected forget loss over all reweightings within a KL ball.

A second class of methods treats unlearning as distributional alignment under shift. DUI augments influence-function updates with an independence criterion between features, labels, and predictions, first via mutual information and then operationally via HSIC, so that the model’s feature–prediction dependence is re-aligned with the remaining dataset after non-uniform feature or label removal (Han et al., 2024). UIB instead formulates unlearning as a parameter-space information bottleneck,

$\mathcal{B}$8

with a dynamic prior and structured regularization to adapt to distribution shifts caused by removing systematic patterns and biases (Han et al., 2024).

A third class acts directly on output distributions. For code LLMs, PROD constructs a target next-token distribution that sets the forget token’s probability to zero, applies top-$\mathcal{B}$9 filtering with $D\setminus S^{full}$0, and redistributes probability mass to plausible alternatives, with $D\setminus S^{full}$1 performing best in ablations (Jiang et al., 20 Jun 2025). The model is then trained by cross-entropy to match this target distribution. In text-to-image diffusion, Diversified Unlearning replaces a single keyword prompt by an empirical distribution of contextually diverse prompts, or by token-wise embedding mixups, thereby defining the concept to be unlearned as a prompt/embedding distribution rather than a point estimate (Pham et al., 19 Mar 2026).

Finally, reference-based methods treat unlearning as distributional indistinguishability. ReGUn constructs a class-conditioned reference distribution $D\setminus S^{full}$2 from a disjoint held-out dataset and minimizes

$D\setminus S^{full}$3

alongside standard retain-set cross-entropy, so that forget-set predictions match unseen-data behavior rather than merely becoming wrong (Mirlach et al., 11 Mar 2026).

4. Evaluation methodologies

Evaluation is a major fault line in this literature. The competition framework for approximate unlearning measures forgetting quality through per-example attacker-based estimation of $D\setminus S^{full}$4, using one-dimensional output statistics under retrained and unlearned model distributions, then aggregates them into a global score $D\setminus S^{full}$5 (Triantafillou et al., 2024). This preserves the original distributional notion of unlearning but operationalizes it with tractable hypothesis tests.

Task-specific benchmarks expose different aspects of distributional forgetting. TOFU reports Forget Quality, Model Utility, Extraction Memorization, Extraction Strength, Truth Ratio variants, and membership inference metrics such as LOSS, ZLib, MinK, and MinK++ (Shao et al., 14 Jan 2026). MUSE expands this into a six-way benchmark for LLMs: no verbatim memorization, no knowledge memorization, no privacy leakage, utility preservation, scalability with respect to removal size, and sustainability over sequential unlearning requests (Shi et al., 2024). This is especially important when the forget set is itself a coherent corpus such as Harry Potter books or BBC news.

Recent work argues that many existing metrics are reference-specific and can hide residual knowledge. FADE addresses this by measuring bidirectional likelihood agreement over generated samples, thereby comparing the full output distributions of an unlearned model and a retain-only model rather than a small set of reference answers or classifier outputs (Cho et al., 14 Oct 2025). For real-world LLM deployments in which the retrained reference is unavailable, DCUE instead evaluates the distribution of Core Token Confidence Scores and applies a corrected Kolmogorov–Smirnov test using a validation set to estimate the retained-data effect (Miao et al., 2 Aug 2025).

The paper on in- vs. out-of-distribution unlearning contributes a separate evaluation axis for generative models: Generalized Exposure and Relative Exposure, which compare soft likelihood rankings of forget examples against reference strings under reference, subject, and unlearned models (Baluta et al., 2024). These metrics are explicitly distribution-sensitive because they track how a forget set’s likelihood distribution shifts relative to appropriate reference distributions.

5. Empirical patterns and application domains

Across modalities, a recurring empirical pattern is that methods can improve forgetting metrics while remaining far from retraining in a stronger distributional sense. On TOFU, BalDRO illustrates the upside of robust reweighting: for forget ratio $D\setminus S^{full}$6, NPO’s Forget Quality improves from $D\setminus S^{full}$7 to $D\setminus S^{full}$8 with BalDRO-G and to $D\setminus S^{full}$9 with BalDRO-DV, while Model Utility remains comparable at $D\setminus S$0 (Shao et al., 14 Jan 2026). On MUSE, BalDRO variants reduce KM-Df and VM-Df, keep or slightly improve KM-Dr, and move PrivLeak closer to zero (Shao et al., 14 Jan 2026).

The in-/out-of-distribution study shows that distributional location of the forget set matters sharply. Unlearning out-of-distribution examples requires more unlearning steps but overall presents a better trade-off, whereas for in-distribution examples there is a rapid decay in performance as unlearning progresses (Baluta et al., 2024). This suggests that the geometry of the forget set relative to the main data manifold strongly controls collateral damage.

For structured distribution shift, DUI remains close to retraining under top-$D\setminus S$1 feature or label removal: on Cora/GIN with unlearn ratio $D\setminus S$2, Retrain gives $D\setminus S$3 F1 while DUI gives $D\setminus S$4; on MNIST/Simple CNN at ratio $D\setminus S$5, Retrain gives $D\setminus S$6 and DUI gives $D\setminus S$7 (Han et al., 2024). UIB similarly reduces bias-correlation and MIA-Efficacy while maintaining F1 under systematic pattern removal (Han et al., 2024).

In source code unlearning, existing methods achieve about $D\setminus S$8 forget quality in copyrighted code removal but HumanEval pass rate drops by more than $D\setminus S$9, making the resulting models practically unusable; PROD achieves the highest Pareto Dominance Ratio across copyrighted code, insecure code, and deprecated API unlearning, with roughly $(\alpha,\varepsilon)$0 average relative improvement versus the best baseline and BLEU around $(\alpha,\varepsilon)$1 under prefix injection attacks, compared with baseline BLEU similarities above $(\alpha,\varepsilon)$2 (Jiang et al., 20 Jun 2025). In diffusion models, Diversified Unlearning consistently improves erasure, benign concept retention, and robustness to Ring-A-Bell, indirect recovery, and noise-based attacks by replacing keyword-only forgetting with prompt-distribution forgetting (Pham et al., 19 Mar 2026).

At the benchmark level, MUSE reaches a more pessimistic conclusion: most algorithms can prevent verbatim memorization and knowledge memorization to varying degrees, but only one algorithm does not lead to severe privacy leakage, and existing algorithms also degrade general model utility and cannot sustainably accommodate successive unlearning requests or large-scale content removal (Shi et al., 2024).

6. Limitations, controversies, and open directions

The field is marked by a persistent ambiguity between deleting samples and deleting distributions. The note “Is your algorithm unlearning or untraining?” argues that much of the literature labeled “unlearning” is technically solving untraining, and that metrics such as membership inference on the forget set can be appropriate for sample-level deletion but insufficient for concept- or domain-level removal (Triantafillou et al., 9 Apr 2026). This suggests that evaluation mismatch remains one of the main sources of confusion.

A second controversy concerns what reference should define successful forgetting. ReGUn uses held-out data as a proxy for unseen behavior (Mirlach et al., 11 Mar 2026), FADE requires a retain-only model and shows that many methods scoring well on traditional metrics remain far from distributional equivalence (Cho et al., 14 Oct 2025), and DCUE attempts to remove this dependence by estimating the retained-data effect from a validation set (Miao et al., 2 Aug 2025). No single evaluation paradigm has yet emerged as universal.

Open technical directions recur across papers. BalDRO studies only KL-DRO and notes that $(\alpha,\varepsilon)$3 and Wasserstein alternatives could induce different weighting behavior; it also leaves robust forgetting-plus-retention formulations open (Shao et al., 14 Jan 2026). Data-centric theories show Pareto frontiers and finite-sample selection guarantees but also report an information-computation gap (Pandey et al., 15 May 2026). The Gaussian/KL framework for forgetting distributions rather than samples suggests that deletion budgets can be reduced substantially by selective removal, but extending these guarantees beyond the analyzed families remains open (Allouah et al., 20 Jul 2025). In sequential decision-making, offline stochastic multi-armed bandits already require separate treatment under fixed-sample and distribution models, with adaptive switching between Gaussian mechanism and rollback depending on the coverage regime (Ye et al., 1 May 2026).

A broader implication is that distributional unlearning is not a single method class but a shift in target. It treats forgetting as robust risk minimization, distributional alignment, output-distribution editing, or statistical indistinguishability, depending on the modality and the comparator. What unites these approaches is the claim that forgetting should be assessed at the level at which the unwanted information is actually represented: as a data distribution, a concept manifold, a predictive distribution, or a distribution over trained models.