PaperlayerView exact revision on arXiv

Source-linked AI summary

Distillation Robustifies Unlearning

Bruce W. Lee, Addie Foote, Alex Infanger, Leni Shor, Harish Kamath, Jacob Goldman-Wetzler, Bryce Woodworth, Alex Cloud, Alexander Matt Turner

arXiv:2506.06278v3cs.LGcs.AI

TL;DR

Existing unlearning methods can suppress undesired behavior while leaving capabilities that a few finetuning steps can restore. The paper studies distillation as a way to make unlearning more resistant to relearning, and introduces UNDO to trade compute cost against robustness. Across language, arithmetic, and WMDP, it reports improved resistance to adversarial relearning, with distilled models reaching oracle-like relearning rates.

  • Problem

    LLM unlearning methods suppress rather than remove undesired capabilities, which can be restored through a few finetuning steps; harmful capabilities in largely unfiltered pretraining data also complicate deployment.

    Source: §1 · p. 1

  • Method

    The paper distills an unlearned model’s outputs into a randomly initialized student and proposes UNDO, which distills into a partially noised copy of the unlearned model.

    Source: §1 · p. 2

  • Results

    Across language, arithmetic, and WMDP, distillation improves resistance to adversarial relearning, with distilled models relearning harmful capabilities at rates comparable to an oracle trained from scratch with data filtering.

    Source: §1 · p. 2

  • Takeaways & Limitations

    UNDO offers a tunable compute–robustness tradeoff as a route toward robust and scalable capability removal.

    Source: §1 · p. 2

    In WMDP experiments, UNDO is 4.88% lower on MMLU on average and has a tradeoff similar to other methods; the authors hypothesize that distilling on only 0.015% of the original pretraining corpus undertrains the model.

    Source: §6 · p. 9

Abstract

from arXiv · show

Current LLM unlearning methods are not robust. A few steps of finetuning can revert their effects. We begin by showing that this is true even for an idealized form of unlearning: training to imitate a model that was never trained on unwanted information. This shows that training a model can drastically modify its input-output behavior while leaving its underlying capabilities intact. In light of this dynamic, we show our main result. Training a randomly initialized student on the outputs of an unlearned model transfers behaviors while leaving latent capabilities behind. In short, distillation robustifies unlearning. Based on this result, we propose Unlearn-Noise-Distill-on-Outputs (UNDO), a scalable method that distills an unlearned model into a noised copy of itself. UNDO introduces a tunable tradeoff between compute cost and robustness, establishing a new Pareto frontier on synthetic language and arithmetic tasks. At its strongest setting, UNDO matches the robustness of a model retrained from scratch with perfect data filtering while using only 60-80% of the compute and requiring only 0.01% of the pretraining data to be labeled. We also show that UNDO robustifies unlearning on the more realistic Weapons of Mass Destruction Proxy (WMDP) benchmark. Since distillation is widely used in practice, incorporating an unlearning step beforehand offers a convenient path to robust capability removal.

1 Introduction

Unlearning methods suppress undesired behavior but often leave capabilities that can return after a few finetuning steps. The paper shows that distillation can improve resistance to relearning and introduces UNDO as a compute–robustness tradeoff.

1 Introduction

  • Pretraining on largely unfiltered data can give LLMs harmful capabilities, while auditing and filtering data at the required scale remains impractical.The introduction frames this as a deployment concern, including potential misuse of cyberweapon-relevant knowledge.Source: §1 · p. 1
  • Across language, arithmetic, and WMDP, distillation improves resistance to adversarial relearning; distilled models relearn harmful capabilities at rates comparable to an oracle trained from scratch with data filtering.The paper reports competitive performance and robustification on WMDP, and says its methods extend the retain-performance versus forget-robustness Pareto frontier on language and arithmetic tasks.Source: §1 · p. 2
  • UNDO distills an unlearned model’s outputs into a partially noised copy of itself, exposing a tunable tradeoff between compute cost and unlearning robustness.The method is presented as interpolating between behavioral suppression and fully robust unlearning.Source: §1 · p. 2
  • Even models trained for perfect behavioral agreement with an unlearning oracle retain latent capabilities that finetuning can restore.This motivates distinguishing surface behavior from robust capability removal.Source: §1 · pp. 1–2
  • Training a randomly initialized student on an unlearned model’s outputs transfers desired behavior while leaving latent capabilities behind.The paper identifies this as the central observation that distillation robustifies unlearning.Source: §1 · p. 2

2 Robust Unlearning

This section defines robust unlearning around preserving desired capabilities while making undesired capabilities difficult to elicit. It evaluates robustness using finetuning attacks and distinguishes capability-focused unlearning from privacy- and fairness-oriented formulations.

2 Robust Unlearning

  • Robust unlearning seeks a model with high retain performance and low elicited forget performance, given a reference model, retain and forget sets, and unlabeled training data.The resulting model is assessed with retain and elicited-forget evaluations.Source: §2 · pp. 2–3
  • Finetuning on the forget set serves as the elicitation mechanism, and evaluation reports the maximum elicited forget performance across attacks at different learning rates.The paper calls these probes relearning attacks and treats resistance to them as a primary criterion for robust unlearning.Source: §2 · p. 3
  • The oracle model represents perfect capability removal by training on the reference model’s pretraining data after filtering content directly related to the undesired capability.It is used as a gold-standard comparison for robust unlearning.Source: §2 · p. 3
  • Unlike privacy- and fairness-oriented formulations, capability-focused retain and forget sets need not correspond to literal subsets of the model’s training data.The paper notes that similar methods are often used across the settings, though they generally do not achieve robustness.Source: §2 · p. 3

3 Oracle Matching Does Not Guarantee Robust Unlearning

Oracle-matched reference models can reproduce an oracle’s outputs yet retain capabilities that relearn faster under finetuning. In language and arithmetic experiments, a randomly initialized student matched to the oracle relearns at a speed closer to the oracle itself.

3 Oracle Matching Does Not Guarantee Robust Unlearning

  • Even effectively perfect behavioral agreement with the oracle does not prevent the reference model from relearning significantly faster, indicating that output matching alone is likely insufficient for robust unlearning.The reference model is finetuned to match oracle outputs on both retain and forget sets before the relearning test.Source: §3 · p. 3
  • The experiments compare a reference model and an oracle in language and arithmetic settings, with the oracle constructed to have only retain capabilities.The language models use English documents for retain and Korean documents for forget; the arithmetic split uses addition/subtraction for retain and multiplication/division for forget.Source: §3 · p. 3
  • Oracle-matched reference models relearn the forget distribution faster than both the randomly initialized student and oracle, in language and arithmetic and across tested learning rates.The comparison follows distillation with 500-step finetuning attacks on the forget set; the figure’s bold lines show the worst-case adversary.Source: Figure 2 · p. 4, §3 · p. 4
  • The authors match reference-model outputs to oracle outputs using KL divergence and also train a randomly initialized model to match the oracle.Distillation loss tends to zero in both setups, establishing behavioral alignment before the relearning comparison.Source: §3 · p. 4
  • A randomly initialized student distilled on oracle outputs has forget-distribution relearning speed close to the oracle’s.The result motivates testing whether robustness persists when practical unlearning methods approximate oracle outputs.Source: §3 · p. 4

4 Distillation Robustifies Unlearning

Distilling a model suppressed by standard unlearning methods into a randomly initialized student makes it substantially more resistant to relearning. In some cases, the student’s relearning trajectory approaches that of full data filtering, though effectiveness can depend on the teacher’s retain-forget tradeoff.

Experimental Setup

  • The protocol applies GradDiff, MaxEnt, or RMU, then distills the suppressed model into a randomly initialized model of the same architecture using forward KL divergence.The evaluated methods represent gradient-based, output-distribution, and representation-level interventions.Source: §4 · pp. 4–5

Results

  • Distillation can be less effective when the unlearning teacher has a weaker retain-forget tradeoff, as with RMU in arithmetic.RMU achieved 62% retain and 6.8% forget accuracy, compared with MaxEnt’s 80% and 1.3%.Source: §4 · p. 5

Motivation

  • Unlearn-and-Distill transfers suppressed behavior using unlabeled pretraining data, avoiding the need to label all pretraining data for filtering.In the language setup, approximately half of pretraining data supported retain capabilities and half supported forget capabilities; standard unlearning used less than 0.01% labeled data.Source: §4 · p. 5

Results

  • Unlearn-and-Distill models resist relearning more than unlearning-only models across language and arithmetic, sometimes approaching full data filtering.The improvement holds for average learning-rate performance and worst-case adversarial performance.Source: §4 · p. 5, Figure 4 · p. 6

Mechanism

  • Distilling an unlearned model into a randomly initialized student slows relearning because the student lacks the teacher’s latent capabilities.The experiments report that standard unlearning models relearn quickly under finetuning, while distilled students relearn substantially more slowly.Source: §4 · pp. 4, 6

Extension to Post-Training Unlearning

  • For post-training unlearning, the original base model can serve as the distillation target when it never had the undesired capability.The paper demonstrates this extension on TOFU.Source: §4 · p. 6

5 Trading Compute for Unlearning Robustness

UNDO perturbs a behaviorally suppressed model and distills the damaged copy to recover the teacher’s behavior. Increasing perturbation generally improves resistance to relearning while requiring more compute, with mixing coefficient α having a stronger effect than noise scale β.

Method

  • UNDO unlearns a model, perturbs its weights to create a damaged student, then distills to recover the teacher’s behavior.The perturbation interpolates between the suppressed model’s parameters and full reinitialization; α = 1 and β = 1 reduces it to random initialization.Source: §5 · pp. 6–7

Experimental Setup

  • Distillation uses forward KL divergence and stops when retain performance reaches at least 95% of the teacher’s or one pretraining-data epoch is complete.The reported setup varies α while fixing β = 0.1.Source: §5 · p. 7

Results

  • Increasing perturbation α improves relearning resistance while increasing compute, with roughly linear tradeoffs in language and a steeper initial compute curve in arithmetic.Language shows noticeable improvement at α = 0.2, while arithmetic shows meaningful improvement at α ≥0.55.Source: Figure 5 · p. 7, §5 · p. 7

Hyperparameter Effects

  • Mixing coefficient α affects robustness more than noise scale β, and α is the primary driver of distillation compute requirements.For most tested α values, β had less impact; lower β values often worked similarly well or better than β = 1.0 in language.Source: Figure 6 · p. 7, §5 · p. 8

6 Comparisons with Other Unlearning Methods

UNDO with MaxEnt maintains robust unlearning across explored attack strengths, creating a Pareto frontier approaching full retraining with data filtering at lower compute and labeling cost. On WMDP, UNDO improves relearning resilience across configurations, while other methods can show similar retain-forget trade-offs.

6 Comparisons with Other Unlearning Methods

  • GradDiff and MaxEnt achieve good initial retain-forget trade-offs but rapidly degrade under adversarial pressure, while SAM and RepNoise also deteriorate under stronger relearning.SAM seeks flat loss regions and RepNoise injects representation-level noise; UNDIAL uses distillation without global parameter corruption such as noising.Source: §6 · p. 8
  • UNDO with MaxEnt maintains more robust unlearning across explored attack strengths, creating a Pareto frontier approaching data-filtered full retraining with less compute and data labeling.More compute-efficient GradDiff and MaxEnt have good initial trade-offs but rapidly degrade under attack; SAM and RepNoise also deteriorate under stronger relearning.Source: §6 · pp. 8–9, Figure 7 · p. 8
  • Across seven adversarial scenarios on Gemma-2-2B, applying UNDO (α = 0.25, β = 0) after MaxEnt or RMU consistently improved WMDP relearning resilience.The comparison covered different data mixtures, question-answer formats, model perturbations, and learning-rate variations.Source: §6 · p. 9, Figure 8 · p. 9
  • In WMDP experiments, UNDO is 4.88% lower on MMLU on average, and other methods attain a similar robust-forget/retain trade-off; distillation also costs more than finetuning-based unlearning.The authors hypothesize that undertraining—distilling on 0.015% of the original pretraining corpus—contributes to the difference from synthetic-task results.Source: §6 · p. 9

7 Related Work

Prior work distinguishes exact unlearning from approximate, output-based approaches, while LLM scale makes exact unlearning impractical. Existing LLM methods can suppress behavior without removing latent capabilities, motivating a student-teacher approach.

7 Related Work

  • Machine unlearning seeks to remove the influence of selected training examples; exact unlearning matches retraining without them, while approximate unlearning requires probabilistically similar output distributions.The term was introduced by Cao and Yang in work on adjusting outputs to mimic retraining without forgotten data.Source: §7 · p. 9
  • For LLMs, scale and non-convexity make exact unlearning impractical, so recent work frames it as optimization that approximates retraining behavior.Methods include maximizing prediction loss on forget sets and aligning predicted token distributions with target distributions using KL divergence.Source: §7 · p. 9
  • Recent studies find that LLM unlearning often suppresses behavior rather than truly removing it, leaving suppressed knowledge vulnerable to finetuning attacks and motivating student-teacher approaches.The cited findings include [32] [92] [76] on behavioral suppression and [33] [49] [18] on recovery through finetuning.Source: §7 · p. 10

8 Conclusion

The paper identifies a disconnect between a model’s observed behavior and its latent capabilities. It reports that oracle matching leaves capabilities vulnerable to elicitation, whereas distillation discards them while preserving desired behavior.

8 Conclusion

  • Oracle-matched models can retain latent capabilities that remain vulnerable to adversarial elicitation.The conclusion frames this as a gap between what a model does and what it could do.Source: §8 · p. 10
  • Distillation discards latent capabilities while preserving desired behavior, turning a standard LLM-development practice into a security measure for robust capability removal.The paper describes this as a practical path to robust capability removal.Source: §8 · p. 10

A Experimental Details for Section 3

The experiments use scaled Gemma-2 models and separate retain and forget domains for language and arithmetic. Oracle matching, followed by relearning attacks, tests whether unlearned capabilities can be recovered.

A Experimental Details for Section 3

  • Language experiments use FineWeb Korean and English sources, sampling 10 million rows from each; arithmetic experiments generate one million examples across four basic operations.Arithmetic operands range from 1 to 50 for addition and subtraction, and 1 to 20 for multiplication and division; division examples have no remainders.Source: A Experimental Details for Section 3 · p. 17
  • Arithmetic examples use direct equations and template-generated word problems, with separate validation sets allocating 500,000 questions or tokens to each retain and forget domain.Language validation uses English as retain and Korean as forget; arithmetic validation uses addition/subtraction as retain and multiplication/division as forget.Source: A Experimental Details for Section 3 · p. 17
  • Data processing uses the Gemma-2-2b tokenizer, chunks documents to at most 2,048 tokens, filters documents under 50 tokens, and fixes the random seed at 42.Tokenization and formatting run in parallel across CPU cores.Source: A Experimental Details for Section 3 · p. 17
  • Language experiments use Gemma-2-0.1B and arithmetic experiments Gemma-2-0.3B, with model sizes varied to show results across scales while maintaining the original architecture’s parameter ratio.The architecture parameters are detailed in Table 1.Source: A Experimental Details for Section 3 · p. 17, Table 1 · p. 17
  • Base and Oracle models use different pretraining data: the language Oracle uses English only, while the arithmetic Oracle uses addition and subtraction with English; their corresponding Base models include the forget domains.The passage reports about 1B versus 2B tokens for language Oracle and Base, and about 100M tokens for each arithmetic model.Source: A Experimental Details for Section 3 · p. 18
  • Oracle matching uses forward-KL knowledge distillation on pretraining data, with an Oracle teacher training either a randomly initialized model or a pretrained Base model.The procedures’ hyperparameters are listed in Table 3.Source: A Experimental Details for Section 3 · p. 18, Table 3 · p. 18
  • Relearning starts from the final oracle-matching student and runs for 500 steps, with validation every five steps tracking retention and recovery of capabilities.The adversarial setup searches ten learning rates from 1e-3 to 1e-6 and marks the best performance obtainable at each validation step.Source: A Experimental Details for Section 3 · p. 19, Table 4 · p. 19

B Experimental Details for Section 4

This appendix documents the experimental methods and configurations for the language, arithmetic, and WMDP experiments across Sections 4–6. It covers training, evaluation, datasets, baselines, distillation and noising protocols, and the construction of WMDP datasets and evaluation of relearning robustness.

GradDiff (Language, CE Loss); GradDiff (Arithmetic, Accuracy)

MaxEnt (Language, CE Loss); MaxEnt (Arithmetic, Accuracy)

RMU (Language, CE Loss); RMU (Arithmetic, Accuracy)

B.2 Distillation

  • Distillation trends compare unlearned, distilled, and gold-standard relearning; distillation uses each model’s pretraining dataset and matches the pretrained counterpart’s update-step count, with hyperparameters listed in Table 6.An additional distillation trend is reported in Figure 9; Figure 10 shows Figure 4 with average lines bolded.Source: §B.2 · p. 21, Figure 9 · p. 22, Figure 10 · p. 22

C Experimental Details for Section 5

  • Section 5 varied corruption α from 0 to 0.8 with β fixed at 0.1, interpolating between original parameters and noise before distillation.α controls the corruption proportion and β the noise magnitude; noised models served as distillation students.Source: C Experimental Details for Section 5 · p. 23

D.1 Language and Arithmetic Unlearning

  • Section 6 evaluates language tasks in English and Korean with Gemma-2-0.1B, and arithmetic tasks in addition-subtraction and multiplication-division with Gemma-2-0.3B.The cited passages describe GradDiff, MaxEnt, RMU, SAM, and RepNoise as unlearning approaches.Source: §D.1 · p. 24
  • For Section 6, GradDiff, MaxEnt, RMU, and UNDIAL used logarithmically spaced learning-rate sweeps, while SAM and RepNoise were configured from the best MaxEnt settings.UNDO fixed α = 0.6 and β = 0.1, used MaxEnt before noising and distillation, and varied retain thresholds from 0.05 to 0.33.Source: §D.1 · pp. 24–25
  • Language and arithmetic relearning was tested after 0, 40, and 500 fine-tuning steps on forget data, monitoring both retain and forget performance.The appendix also specifies sequence lengths, batch sizes, gradient clipping, and method-specific hyperparameters; for example, GradDiff α is 1 for language and 15 for arithmetic.Source: §D.1 · p. 25

D.2 WMDP

  • WMDP experiments use Gemma-2-2b and compare unlearning methods in both the Cyber and Bio domains.Source: §D.2 · p. 25

D.2.1 Evaluation/Benchmark

  • Retain performance is measured by 5-shot accuracy on 40% of MMLU, while forget performance uses the full WMDP Bio or Cyber evaluation with zero-shot prompting.Source: §D.2.1 · p. 25

D.2.2 Datasets

  • To reduce dataset differences that may hinder generalization, the authors create QA datasets from Bio and Cyber retain/forget corpora and use forget QA plus a 50/50 retain-QA/Wikitext mix.Questions and answers are generated from sampled corpus text and English Wikipedia using Gemini flash.Source: §D.2.2 · pp. 25–26

D.2.3 Baseline Unlearning Methods

  • WMDP baselines are RMU and MaxEnt, whose loss combines retain and forget components using α; the selected teacher is the highest-retain model among three frontier points.MaxEnt maximizes forget-set entropy and minimizes retain-set KL divergence from the base model; RMU uses the language/arithmetic loss.Source: §D.2.3 · p. 26

D.2.4 Distillation/UNDO

D.2.5 Relearning Adversaries

  • Relearning evaluations use diverse adversaries varying datasets, sequence length, and shrink-perturb, and report each method’s best adversary; RepNoise and TAR show attack-specific robustness limitations.The appendix describes seven configurations and notes that reasonable-retain hyperparameters for TAR could not be found within the training budget.Source: §D.2.5 · pp. 26–27

D.2.6 Baseline Robust Methods

  • Arithmetic and language robust baselines combine SAM or RepNoise with MaxEnt, using the top two selected MaxEnt hyperparameters.SAM uses ρ = 0.01; RepNoise uses β = 1.0 and α = 1.0.Source: §D.2.6 · p. 28

D.2.7 WMDP Pareto Frontier and Discussion

UNDO is on the WMDP-Bio Pareto frontier and competitive on WMDP-Cyber, though its WMDP tradeoff resembles other methods rather than the stronger language/arithmetic frontier results. The section also discusses compute and data constraints and an alternative setup for studying unlearn-and-distill without reinitializing models.

D.2.7 WMDP Pareto Frontier and Discussion

  • UNDO is on the WMDP-Bio Pareto frontier and competitive on WMDP-Cyber, with a robust-forget/retain tradeoff similar to other methods.Figure 11 reports forget and retain accuracy on Bio and Cyber, immediately after unlearning and after low- and high-compute adversarial retraining.Source: Figure 11 · p. 28, §D.2.7 · p. 28
  • Language/arithmetic results place UNDO far beyond baseline Pareto frontiers, unlike the similar-tradeoff WMDP results.Source: §D.2.7 · p. 28
  • Language/arithmetic used α = 0.6 versus α = 0.25 for WMDP; higher α damages capabilities and makes them harder to retrain.For Gemma-2-2b, higher α values did not recover retain performance.Source: §D.2.7 · p. 29
  • For WMDP, distillation used 300 million tokens, 0.015% of Gemma-2-2B's 2-trillion-token pretraining corpus; higher α requires more recovery data.Source: §D.2.7 · p. 29

E Compute requirements

F Post-Review Reflection: An Alternative Setup to Study Unlearn-and-Distill

  • The alternative setup finetunes a pretrained model on absent dataset D, unlearns it, then distills back onto the original model as a random-init substitute for D.TOFU's fictional author associations approximate absence, but familiar tokens mean semantic or keyword novelty cannot guarantee no latent traces.Source: F Post-Review Reflection: An… · p. 30

F.1 TOFU Experiments

On TOFU, unlearning followed by distillation slows relearning relative to unlearned-only models, with trajectories closer to a control that never saw the forget knowledge. Distillation also improves robustness under INT4 quantization in the reported metrics.

F.1 TOFU Experiments

  • The three-step protocol finetuned Gemma-2-2B and Phi-1.5 on TOFU, applied GradDiff to suppress forget knowledge while retaining performance, then distilled onto fresh base models.Distillation used the full TOFU dataset, including forget and retain domains.Source: §F.1 · p. 31
  • Across training phases, distillation preserved forget metrics near the unseen base-model level while retain metrics recovered toward teacher performance.During unlearning, Forget QA probability dropped near zero within 50 steps while retain performance stabilized.Source: §F.1 · p. 31
  • After unlearning and distillation, relearning trajectories stayed closer to the control baseline than unlearned-only models, which rapidly recovered suppressed knowledge.This pattern was reported for Gemma-2-2B and Phi-1.5.Source: Figure 13 · p. 31, §F.1 · p. 31

F.2 Robustness Beyond Relearning Attacks

  • Under INT4, forget truth ratio rose from 0.000 to 0.185 after standard unlearning but reached only 0.017 after distillation; utility was 0.391 vs 0.242.The authors interpret the results as suggesting different internal representations, with the student lacking latent structures exploited by quantization.Source: Table 14 · p. 32, §F.2 · p. 32

Related papers

Loading 2506.06278v3…