Data Unlearning via Inverse Distillation

AI summary

Overview

Research area: Machine unlearning for generative models — specifically, removing designated training data from multi-step diffusion/flow ("matching") models while simultaneously distilling them into efficient one-step generators.

Technical level: Advanced. The method is built on a min-max inverse-distillation objective over probability distributions, matching-model losses, and mixture distributions; it assumes familiarity with diffusion, flow matching, and distillation.

Scope: The paper introduces Inverse Distillation Unlearning (IDU), a single framework that both distills a pretrained multi-step teacher into a one-step student and suppresses the student's outputs corresponding to a supplied forget set, evaluated on MNIST and CIFAR-10 under flow-matching and score-based (SiD/EDM-VP) settings.

What This Paper Is About

Multi-step diffusion and flow models generate high-quality samples but are slow, and they can reproduce unwanted content from their training data. Existing tools for removing such content (machine unlearning) are well developed for multi-step models, but little progress has been made for one-step generators, which are the practical deployment target. IDU addresses both problems at once: using only a pretrained full-data teacher and samples from the forget set, it trains a one-step generator that recovers the retained data without ever seeing retained training examples, extra feature extractors, or classifiers.

Key Contributions

1. A unified unlearning-plus-distillation framework (IDU). The authors formulate distillation as a min-max objective over a data distribution, then parametrize that distribution as a mixture of the forget set and the generator's own output distribution, p_0^mix = ρ·p_0^F + (1−ρ)·p_0^θ, so that the forget component is already accounted for by supplied forget samples and only the retained data must be learned.

2. A theoretical forgetting guarantee. Theorem 1 states that optimizing the IDU loss with ρ = π (the true forget-data proportion) retrieves only the remaining data, i.e., the optimal generator satisfies p_0^{θ_opt} = p_0^R.

3. A data- and model-agnostic recipe. IDU requires only a pretrained full-data teacher of arbitrary matching type (diffusion or flow) and data from the forget set. No retained training examples, no extra feature extractors, and no classifiers are used in training; the only method-specific trade-off parameter is ρ ∈ [0, 1), with a backbone-specific distillation scale α_SiD ∈ [0.5, 1.2] (0.5 for flow matching, 1.2 for SiD).

4. Empirical demonstration on four dataset–backbone combinations, plus single-class robustness sweeps over all ten MNIST digits (FM) and all ten CIFAR-10 classes (FM and SiD), and an ablation over the forgetting weight ρ. The authors state that IDU is, to their knowledge, the first unified framework for simultaneous unlearning and distillation in unconditional flow-matching and score-based models.

Main Findings

• Forgetting is substantial in all four settings. Under flow matching (FM), every forgotten-class FGR (forgotten-class generation rate, the percentage of 50,000 generated images assigned to a forgotten class by an off-the-shelf classifier) is at most 0.56%; under SiD it is at most 1.25%. Baseline pretrained and purely distilled models sit near 7.58–12.50% FGR on CIFAR-10 and 9.65–10.99% on MNIST.

• Retained quality is largely preserved. IDU retains generation quality close to the corresponding retrain-distillation oracle in three settings and improves on that reference for FM on MNIST. On MNIST with FM and ρ = 0.4, IDU reaches Retain FID 3.57 ± 0.03 versus 5.72 ± 0.12 for retrain distillation; on MNIST SiD with ρ = 0.2, 1.65 ± 0.02 versus 1.20 ± 0.02; on CIFAR-10 FM with ρ = 0.6, 5.81 ± 0.05 versus 5.12 ± 0.05; on CIFAR-10 SiD with ρ = 0.2, 3.15 ± 0.03 versus 2.82 ± 0.02. The authors caution that the FM/MNIST oracle comparison is sensitive to the retained-only retraining baseline (Appendix B.9).

• Main paired-class results (Table 1). FM on MNIST, forgetting {3, 7}, ρ_MNIST = 0.4: Retain FID 3.57 ± 0.03, FGR 0.16 ± 0.01 and 0.16 ± 0.01. FM on CIFAR-10, forgetting {1, 9} (automobile and truck), ρ_CIFAR-10 = 0.6: Retain FID 5.81 ± 0.05, FGR 0.56 ± 0.05 and 0.39 ± 0.03. SiD on MNIST, forgetting {3, 7}, ρ = 0.2: Retain FID 1.65 ± 0.02, FGR 0.63 ± 0.03 and 0.32 ± 0.02. SiD on CIFAR-10, forgetting {1, 9}, ρ = 0.2: Retain FID 3.15 ± 0.03, FGR 1.11 ± 0.06 and 1.25 ± 0.06.

• Low FGR generalizes beyond the chosen pairs. Across ten single-class forgetting tasks per setting, mean Retain FID/FGR is 4.54 / 0.16% for FM–MNIST, 7.73 / 0.84% for FM–CIFAR-10, and 3.61 / 0.75% for SiD–CIFAR-10.

• The forgetting weight ρ is the key trade-off, and more is not better. On CIFAR-10, low ρ preserves FID but leaves higher FGR; ρ = 0.8 improves forgetting at substantial FID cost for FM (FID 8.59 ± 0.11). FM training at ρ = 0.9 diverges immediately. SiD does not diverge at ρ = 0.9, but its FGRs (12.05% and 9.55%) approach those of the full-data teacher while Retain FID rises to 10.36 ± 0.10 — so excessive ρ can degrade fidelity without forgetting.

• Effect of ρ, exact numbers (Table 3, CIFAR-10, forgetting classes 1 and 9). FM: ρ = 0.05 → FID 6.72 ± 0.07, FGR 3.04 ± 0.08 / 5.02 ± 0.08; 0.1 → 6.21 ± 0.05, 3.75 ± 0.09 / 3.80 ± 0.08; 0.2 → 5.64 ± 0.04, 3.35 ± 0.03 / 3.04 ± 0.10; 0.4 → 5.87 ± 0.07, 0.62 ± 0.06 / 0.48 ± 0.02; 0.6 → 5.81 ± 0.05, 0.56 ± 0.05 / 0.39 ± 0.03 (selected); 0.8 → 8.59 ± 0.11, 0.38 ± 0.03 / 0.32 ± 0.03; 0.9 diverged. SiD: 0.2 → 3.15 ± 0.03, 1.11 ± 0.06 / 1.25 ± 0.06 (selected); 0.4 → 3.53 ± 0.04, 0.74 ± 0.05 / 0.60 ± 0.01; 0.6 → 4.68 ± 0.06, 1.16 ± 0.05 / 0.76 ± 0.02; 0.8 → 6.29 ± 0.07, 0.81 ± 0.05 / 0.58 ± 0.02; 0.9 → 10.36 ± 0.10, 12.05 ± 0.18 / 9.55 ± 0.12.

• Fine-tuning works too, and the framework also recovers ordinary distillation. Experiments on purely distilled generators (fine-tuning instead of training from scratch) achieve similar metrics (Appendix B.7), and running the framework without forget data verifies that it also recovers a standard one-step generator.

• Evaluation protocol details. FID uses full-data references for pretrain/pure distillation and retained-data references for retrain/IDU (Retain FID). Reported values are means and standard deviations over five evaluations/runs, except values marked † , which follow Zhou et al. (2024). Classifier and FID protocols and architecture/sampler/hyperparameter details are in Appendices B.3, B.4, B.1, B.2, and B.5.

• Not reported in the provided content: numerical head-to-head comparisons against prior unlearning methods (for example NegGrad, SA, SalUn, SISS, MGSM, EraseDiff, Retrack, VDU, ContinualFlow, UOT-Unlearn, SFD) are discussed qualitatively but their scores do not appear in the content shown.

Methodology in Plain English

Start from a teacher model that has already been trained on the full dataset (including the data you want gone). Standard "inverse distillation" asks: what data distribution would make this teacher optimal? You answer it by playing a min-max game — a trainable "fake" model tries to fit your current guessed distribution, and the generator is pushed toward the distribution where the teacher fits it better than any fake model could.

IDU changes what distribution you guess. Instead of guessing the generator's output alone, you guess a blend: a fraction ρ of real forget samples plus a fraction 1 − ρ of whatever the generator currently produces. Because the forget portion is supplied as actual data, the fake model can already explain it. If the generator also emits forget-like samples, it creates an excess that the teacher does not endorse, and the generator loss penalizes exactly that excess. So the only way for the generator to make the fake model agree with the teacher is to stop producing forget-like samples and reproduce the remaining data — which is what Theorem 1 proves when ρ equals the true forget proportion.

Practically, training alternates two updates. First, the fake model is trained on both branches — noised forget samples weighted by ρ and noised generated samples weighted by 1 − ρ — so it learns to detect generations that resemble the forget set. Second, the generator is updated using the frozen teacher and the fake model, with the generator loss written in the modified form used by Score identity Distillation (SiD), where a term is scaled by α_SiD (heuristically in [0.5, 1.2]). Algorithm 1 gives the flow-matching pseudocode; Figure 1 shows the two forward-noising branches and the alternating updates. Only the teacher, the generator (pretrained or initialized from the teacher), the fake model, the forget data, ρ, α_SiD, iteration count, batch size, latent distribution, and noise distribution are needed — no retained images and no external classifiers enter the training objective.

Why This Matters

Impact on research. IDU closes a stated gap: data unlearning has been well explored for multi-step matching models but rarely for one-step generators, where the only prior framework discussed (UOT-Unlearn) needs a precomputed feature extractor and hyperparameter tuning rather than a trained teacher. IDU is also positioned as a unified approach spanning both flow-matching and score-based teachers, rather than a model-specific fix, and it connects unlearning to the inverse-distillation view of distillation.

Real-world applications (as motivated by the paper's framing):

• Removing individual faces or specific images from an unconditional generator without retraining from scratch.
• Suppressing Not-Safe-For-Work (NSFW) content that a model may have memorized from imperfectly moderated web data.
• Handling privacy, copyright, and legal takedown requests against generative models at the level of specific samples rather than whole classes.
• Producing fast one-step generators that are simultaneously "cleaner" than the multi-step teachers they came from, which suits deployment where sampling cost matters.

Industry relevance. The method's operational profile is attractive: it needs only a teacher and forget samples, avoids extra feature extractors and classifiers, has essentially one trade-off knob (ρ), and outputs a one-step generator — which directly addresses both inference cost and content-removal compliance. The ablations also give practitioners a clear warning that overshooting ρ can cost fidelity without buying forgetting.

Future Directions

• Choosing ρ without retained data. The authors note that ρ can be set from the forget-set proportion and still avoid significant degradation, and suggest that if direct evaluation is critical, retained data could be selected from teacher samples by classification or manual selection. A principled, data-free selection rule remains open.
• Extending beyond unconditional image benchmarks. Evaluation is limited to MNIST and CIFAR-10 with FM and SiD; scaling to higher-resolution data and to conditional or text-to-image settings (where concept unlearning is studied) is not demonstrated.
• Strengthening comparisons. The content presented does not include numerical head-to-head results against the many data-unlearning baselines it surveys, so establishing relative standing — and understanding UOT-Unlearn's scalability and generation-diversity limitations in this light — is a natural next step.
• Understanding ρ's non-monotone behavior. SiD's class-wise FGR values vary non-monotonically as ρ increases from 0.2 to 0.8, and ρ = 0.9 degrades fidelity without forgetting; explaining and predicting this behavior is an open question.

Target Audience

Researchers and graduate students working on machine unlearning, diffusion and flow matching, and distillation of generative models; practitioners who need to remove specific content from a deployed generator while also making it cheap to sample; and readers interested in min-max/inverse-optimization formulations that connect training objectives to the data distributions they implicitly encode.

Authors’ abstract

Multi-step matching models, including flow and diffusion models, produce high-quality outputs but incur substantial inference costs and may reproduce unwanted components of their training datasets. We introduce Inverse Distillation Unlearning (IDU), a unified framework that simultaneously distills a teacher multi-step matching model into an efficient one-step student generator and suppresses outputs corresponding to a designated training subset. We first formulate distillation as a min-max objective over a data distribution and then represent this distribution as a mixture of the forget-set and the generated distributions. This allows us to compare this mixture with the teacher's training distribution and recover only the retained data at the optimum. Our method requires only a pretrained full-data teacher and data from the forget set, without access to retained training examples, extra feature extractors or classifiers. Extensive experiments on MNIST and CIFAR-10 datasets under flow-matching and score-based diffusion settings demonstrate that IDU substantially reduces the generation frequency of forgotten classes while preserving generation quality on the retained classes. To the best of our knowledge, IDU is the first unified framework for simultaneous unlearning and distillation in unconditional flow-matching and score-based models.