PaperlayerView exact revision on arXiv

Source-linked AI summary

SAUL: Sharpness-Aware Augmented-Lagrangian Unlearning

Jaewan Choi, Junyoung Yang, Sangdon Park

arXiv:2608.16249v1cs.LG

TL;DR

LLM unlearning must remove sensitive or harmful information while preserving general utility, but existing methods often specify forgetting implicitly. SAUL makes forgetting an explicit constrained criterion and adaptively controls forget-side pressure, improving forgetting–utility trade-offs across TOFU, WMDP, and MUSE while its controller also improves post-forgetting utility when added to representative baselines.

  • Problem

    LLM unlearning must remove target information while preserving utility, while prior methods typically specify the desired forgetting level implicitly through optimization objectives.

  • Method

    SAUL formulates forgetting as a thresholded constraint and combines augmented-Lagrangian control with sharpness-aware updates and role-separated dual optimizer states.

  • Results

    SAUL improves the forgetting–utility trade-off across TOFU, WMDP, and MUSE under matched criteria, while its controller improves post-forgetting utility as a drop-in baseline modifier.

  • Takeaways & Limitations

    Explicit forgetting control can preserve retain-side utility and improve representative unlearning baselines within the evaluated benchmark settings.

  • Takeaways & Limitations

    SAUL requires selecting a forgetting threshold whose optimal value may depend on the dataset, model scale, forget ratio, and evaluation metric.

Abstract

from arXiv · show

Machine unlearning in Large Language Models (LLMs) faces a critical trade-off between erasing target knowledge and preserving general utility. We propose SAUL (Sharpness-Aware Augmented-Lagrangian Unlearning), which formulates unlearning as a constrained minimization problem following the principle of "forget enough, but no more than necessary." At its core, SAUL formulates forgetting as an explicit constraint with a prescribed satisfaction criterion, whereas prior unlearning methods typically specify the desired level of forgetting implicitly through optimization objectives. An augmented Lagrangian controller adaptively adjusts forget-side pressure according to constraint violation and can eventually deactivate the forget-side update as the prescribed criterion remains satisfied. Sharpness-aware updates on both retain and forget objectives, together with a dual-optimizer design that maintains role-separated states, further stabilize the resulting unlearning dynamics. We evaluate SAUL on the TOFU, WMDP, and MUSE benchmarks, demonstrating favorable forgetting-utility trade-offs over representative sharpness- and perturbation-based baselines under benchmark-specific forgetting criteria. Beyond the complete SAUL framework, we further show on TOFU that applying the augmented-Lagrangian controller as a drop-in modifier to representative baselines improves their post-forgetting utility, demonstrating the practical value of explicit forgetting control.

1 Introduction

SAUL formulates LLM unlearning as constrained minimization: forget enough to meet an explicit threshold α, then prioritize retain-side utility. Its augmented-Lagrangian controller adaptively reduces or deactivates forget-side pressure, while sharpness-aware updates and role-separated optimizer states support favorable trade-offs across TOFU, WMDP, and MUSE.

  • Introduction: LLM unlearning targets sensitive, copyrighted, or harmful information while avoiding the prohibitive cost of retraining models from scratch.The task erases target knowledge from a forget set while preserving general capabilities on a retain set.
  • Introduction: Existing weighted forget–retain objectives use trade-off coefficients that are sensitive to the dataset, model architecture, and forgetting ratio.Bi-level methods structure the conflict through nested or sequential objectives, while sharpness- and perturbation-aware methods address instability from aggressive forgetting.
  • Introduction: SAUL casts unlearning as constrained minimization, using threshold α to specify the target forgetting level and focus optimization on retain-side utility after satisfaction.This follows the principle of forgetting enough, but no more than necessary.
  • Introduction: The augmented-Lagrangian controller adapts forgetting pressure to criterion satisfaction, with its projected multiplier decreasing and potentially gating off forget-side gradients entirely.When the prescribed criterion remains satisfied, optimization can focus solely on preserving retain-side utility.
  • Introduction: Across TOFU, WMDP, and MUSE, SAUL improves forgetting–utility trade-offs over representative baselines, while its controller consistently improves post-forgetting utility as a drop-in modifier on TOFU.SAUL combines the controller with sharpness-aware updates and role-separated dual optimizer states.

2 Related Work

Prior LLM unlearning work largely balances forgetting and retention through weighted objectives or structured optimization, while existing constrained and sharpness-aware methods do not prescribe an explicit threshold for sufficient forgetting. Augmented-Lagrangian methods provide a general constraint-enforcement mechanism, but prior constrained unlearning places constraints on retention rather than forgetting.

  • LLM unlearning and weighted objectives: Many LLM unlearning methods scalarize forgetting and retention with a trade-off coefficient, including gradient-ascent, gradient-difference, loss-adjustment, and preference-optimization formulations.The cited representative methods include Jang et al. (2023), Wang et al. (2024), Bu et al. (2025), Zhang et al. (2024), and Mekala et al. (2024).
  • Bi-level and structured optimization for unlearning: BLUR and OFMU use structured bi-level optimization to separate forgetting from retain-side utility preservation or restoration.BLUR prioritizes forgetting at the lower level and retain utility at the upper level, whereas OFMU combines inner forgetting maximization with similarity-aware penalization and outer utility restoration.
  • Bi-level and structured optimization for unlearning: These structured formulations separate forgetting and retention but still lack a prescribed threshold determining when forgetting is sufficient.This limitation is stated for the bi-level and structured approaches discussed in the preceding passages.
  • Constrained forgetting and augmented Lagrangian methods: Augmented-Lagrangian methods enforce constraints through dual-variable updates and quadratic penalties rather than relying solely on a fixed penalty coefficient.The passage attributes this general mechanism to Hestenes (1969), Bertsekas (2014), and Nocedal and Wright (2006).
  • Constrained forgetting and augmented Lagrangian methods: Existing constrained LLM unlearning methods place constraints on the retain side, including a hard retain constraint in Constrained Entropic Unlearning.The passage also identifies Cheng et al. (2026) as another constrained approach, though its description is truncated in the supplied text.
  • Sharpness-aware and role-separated optimization: Sharpness-aware and perturbation-aware unlearning methods seek more stable retention or forgetting behavior and reduced susceptibility to relearning-style recovery, but generally lack an explicit forgetting criterion.The passage links this line of work to SAM and methods by Fan et al. (2025), Tang and Khanna (2025), Kim et al. (2025), and Malekmohammadi et al. (2025).

3 Problem: Unlearning with Constraints

This section formulates machine unlearning as constrained minimization: remove influence from a forget set while preserving performance on a retain set. Sufficient forgetting is enforced explicitly through a threshold on forget-side loss, while retain loss limits unnecessary utility degradation.

  • Machine unlearning removes the influence of designated training data while preserving or potentially improving performance on the remaining data.
  • The formulation uses average per-example retain loss Lr(θ) and forget loss Lf(θ) to measure model fit on the retain and forget sets.
  • The forget constraint requires Lf(θ) ≥ α, where α specifies sufficient forgetting instead of an implicit trade-off coefficient.The retain objective discourages utility degradation beyond what is necessary to satisfy this criterion.

4 SAUL: Sharpness-Aware Augmented Lagrangian Unlearning

SAUL formulates unlearning as constrained optimization, combining adaptive augmented-Lagrangian forgetting control with sharpness-aware retain and forget updates and role-separated optimizer states. Its controller activates forgetting only when the constraint is violated and can deactivate it once the criterion is satisfied, while the same controller can augment existing baselines.

  • Core framework: SAUL combines an augmented-Lagrangian constraint, sharpness-aware updates, and dual optimizer states to decouple forgetting and retention dynamics.The retain branch remains active throughout training, while the forget branch is conditionally applied.
  • Augmented-Lagrangian control: The projected multiplier λ+ = [λt + µcf]+ activates forgetting only when the forget constraint is violated; λ+ = 0 leaves retain optimization active alone.As violations disappear, the multiplier can decrease and deactivate unnecessary forget-side pressure.
  • Sharpness-aware updates: Sharpness-aware losses are applied to both retention and forgetting, stabilizing retain performance under perturbations while keeping forget knowledge unrecoverable nearby.The retain side controls worst-case loss, whereas the forget side maximizes loss at the most easily recoverable perturbation.
  • Role-separated optimizer states: SAUL maintains independent AdamW optimizer states for forgetting and retention, preventing interference between conflicting gradients and heterogeneous scales.The parameters remain shared, but each state accumulates moments from only its corresponding objective.
  • Complete algorithm: The complete procedure combines adaptive forgetting activation, sharpness-aware gradients, and role-separated optimizer states for stable unlearning dynamics.These three components are summarized together in Algorithm 1.
  • Baseline augmentation: The augmented-Lagrangian controller can modify existing unlearning baselines by increasing forget pressure under violation and reducing it when the criterion remains satisfied.Unlike SAUL, these variants add only constraint-driven control and retain each baseline’s original structure.

5 Experiments

SAUL is evaluated across ToFU, WMDP, and MUSE under matched benchmark-specific forgetting criteria. It achieves strong forgetting–utility trade-offs, with ALM improving utility and robustness while preserving constraint satisfaction.

  • Benchmarks: Experiments span ToFU, WMDP, and MUSE, covering author-specific fictitious forgetting, hazardous-knowledge removal, and verbatim memorization under matched forgetting criteria.The evaluations use benchmark-specific models and configurations rather than a shared backbone.
  • Main results: At a 1% forget ratio with FROUGE ≤0.03, SAUL reaches Forget ROUGE 0.02 and HM 72.79, exceeding Sharp Min–Max + ALM at HM 67.11 by 5.7 points.BLUR achieves HM 71.40 but has Forget ROUGE 0.28, indicating substantial answer-level forget-set leakage.
  • Effect of ALM on Utility and Generalization: Adding ALM to baseline methods improves model utility and GPT-based evaluation, with larger gains on Real Authors than on Retain.The controller deactivates forget-side pressure after the criterion is met, avoiding unnecessary updates.
  • Robustness and generalization: SAUL remains effective under GPT-paraphrased questions, indicating that forgetting is not tied to the original questions’ surface forms.Additional results report the same overall trend across 5% and 10% forget ratios and the larger Llama-3.2-3B backbone.
  • WMDP results: On WMDP, matched forgetting keeps averaged hazardous-knowledge accuracy within 0.005 across methods, while SAUL obtains MMLU accuracy 0.542 ± 0.003 versus BLUR’s 0.540 ± 0.004.Sharp Min–Max reaches 0.535 ± 0.002, whereas Relearning-resilient reaches 0.414.
  • MUSE results: On MUSE, SAUL achieves VerbMem = 0.00±0.00 and KnowMem on Df = 0.77±1.34 while preserving retain-side knowledge at 48.25±5.82 and obtaining PrivLeak −16.66 ± 2.90.Relearning-resilient preserves 39.03 ± 1.03 retain-side knowledge, compared with SAUL’s 48.25 ± 5.82.

6 Conclusion

SAUL formulates LLM unlearning as constrained optimization under “forget enough, but no more than necessary,” using explicit forgetting control alongside sharpness-aware updates and role-separated optimizer states. Across ToFU, WMDP, and MUSE Books, it improves the forgetting–utility trade-off under matched criteria, while its controller also improves representative baseline methods as a drop-in modifier.

  • Method: SAUL formulates LLM unlearning as constrained optimization, reducing forget-side pressure after a prescribed forgetting criterion is met.It combines an augmented-Lagrangian controller with sharpness-aware updates and role-separated optimizer states for stable unlearning.
  • Empirical findings: Experiments on ToFU, WMDP, and MUSE Books show improved forgetting–utility trade-offs under matched forgetting criteria while preserving retain-side utility, neighborhood-data performance, and general capability.SAUL avoids excessive forgetting across these evaluations.
  • Empirical findings: The augmented-Lagrangian controller serves as a method-agnostic drop-in modifier for representative sharpness- and perturbation-based baselines.This demonstrates the broader value of explicit forgetting control in LLM unlearning.

7 Limitations … B.4 GPT-based Evaluation Protocol

The paper reports threshold, robustness, and computational limitations, then details baseline comparisons, ALM-augmented variants, experimental settings, optimization choices, and GPT-based evaluation. These procedures use matched forgetting criteria, role-separated updates, and a fixed semantic-entailment protocol.

  • 7 Limitations: SAUL requires a dataset- and model-dependent forgetting threshold, lacks formal guarantees against all recovery attacks, and adds moderate sharpness-aware and augmented-Lagrangian computation.Automatic threshold selection, stronger adaptive attacks, and more memory- and compute-efficient implementations remain future directions.
  • A Baselines: Experiments compare standard optimization, recent LLM unlearning, and ALM-augmented sharpness-based baselines under the same backbone, splits, training budget, and metrics.Operating points are selected under matched forgetting criteria so utility differences are compared at comparable forgetting levels.
  • A.1 Standard Optimization Baselines: Single AdamW shares optimizer state across retain and forget updates, whereas Dual AdamW separates adaptive states to isolate role-separated optimizer stability.Dual AdamW uses shared model parameters while maintaining distinct retain- and forget-side AdamW states.
  • A.2 Recent LLM Unlearning Baselines: Recent baselines differ in their mechanisms: Relearning-resilient Unlearning sharpens forgetting, Sharp Min–Max uses opposing retain and forget sharpness dynamics, BLUR is bi-level, and PDU constrains retain utility.Sharp Min–Max does not adaptively deactivate forget-side pressure at a prescribed threshold, while SAUL explicitly constrains forgetting.
  • A.3 ALM-Augmented Sharpness-Based Variants: The ALM controller increases forget-side pressure when forgetting violates F(θ) ≥ α and decreases it after sustained satisfaction, including for ALM-augmented Sharp Min–Max.For Sharp Min–Max, the retain SAM update remains intact while the multiplier scales, reduces, or eventually skips the forget-side update.
  • B.1 Training Configuration: Training uses dataset-specific configurations: TOFU and MUSE-Books run for 10 epochs, WMDP follows one-epoch RMU training, and MUSE-Books reports a single run because 7B evaluation is costly.WMDP uses at most 150 mini-batches, batch size 4, and two GPUs; MUSE-Books uses per-device batch size 1 with gradient checkpointing when needed.
  • B.2 Hyperparameter Selection: Hyperparameters are selected by grid search, with thresholds, perturbation radii, and learning rates tuned separately across TOFU, WMDP, MUSE-Books, and MUSE-News.TOFU configurations target Forget ROUGE around 0.03 while preserving retain performance; MUSE-News selection also uses privacy leakage, verbatim memorization, and retain-side stability.

B.5 Qualitative Analysis of Generated Responses … C.1 Paraphrase Generation

The appendix provides qualitative Forget and Retain response examples, documents artifact licensing and research-only use, and constructs ToFU paraphrases to test unlearning stability under query rephrasing. Each paraphrase preserves the original answer and semantic target while introducing substantial lexical and syntactic variation.

  • B.5 Qualitative Analysis of Generated Responses: Tables 12 and 13 provide representative model responses from the ToFU Forget and Retain sets, respectively, to illustrate forgetting and knowledge preservation after unlearning.The examples are drawn from the Forget subset, with corresponding Retain examples used to illustrate preserved knowledge.
  • B.6 Artifact Use and Licenses: The evaluation uses publicly available ToFU, WMDP, and MUSE Books benchmarks alongside the model backbones described in Section 5.1.ToFU and WMDP are released under MIT; MUSE Books and Llama backbones follow their official release terms, and Zephyr-7B-β is MIT-licensed.
  • B.6 Artifact Use and Licenses: The artifacts are used only for research evaluation, and the submission does not redistribute original datasets, model checkpoints, or benchmark resources.Users are directed to the original releases for applicable licenses, access conditions, and terms of use.
  • C Paraphrased ToFU Construction and Validation: GPT-paraphrased ToFU questions are intended to evaluate robustness to query rephrasing under research contexts compatible with the original benchmark conditions.The paraphrases are derived from benchmark questions and are not presented as redistributed benchmark resources.
  • C Paraphrased ToFU Construction and Validation: The constructed evaluation set contains one paraphrased question per original ToFU question, keeps the ground-truth answer unchanged, and is excluded from unlearning training.Its purpose is to test whether unlearning behavior remains stable under query rephrasing.
  • C.1 Paraphrase Generation: The paraphrase prompt preserves exact meaning and answer identity while requiring substantially different wording, altered sentence structure, minimal token overlap, and no added or removed information.It also requires preserving the original question type and outputting only the paraphrased question.
  • C.1 Paraphrase Generation: Domain-specific concepts must remain unchanged rather than being replaced with related but distinct concepts, preserving the question’s exact semantic target.This constraint complements the prompt’s requirement to preserve meaning while diversifying wording and structure.

C.2 Paraphrase Quality Validation · C.3 Examples · D Additional Unlearning Results

The paper validates GPT-generated ToFU paraphrases using lexical-overlap and semantic-preservation metrics, finding substantially reduced surface overlap while largely preserving meaning. Representative original–paraphrase examples are also provided.

  • C.2 Paraphrase Quality Validation: Paraphrase quality is assessed with BLEU, ROUGE-L F1, and SBERT cosine similarity to measure lexical overlap, structural overlap, and semantic preservation.BLEU uses the original question as reference, ROUGE-L F1 measures longest-common-subsequence overlap, and SBERT uses all-MiniLM-L6-v2 embeddings.
  • C.2 Paraphrase Quality Validation: The validation statistics are reported as mean ± standard deviation over question pairs for BLEU, ROUGE-L F1, and semantic similarity.The table describes semantic similarity as meaning preservation using SBERT embeddings.
  • C.2 Paraphrase Quality Validation: 0.205 BLEU and 0.499 ROUGE-L F1 across 1,017 paraphrased questions indicate substantially reduced lexical and structural overlap.The reported averages are computed across all 1,017 paraphrased questions.
  • C.2 Paraphrase Quality Validation: 0.929 average semantic similarity suggests that the paraphrased questions largely preserve the original semantic content.Together with reduced lexical and structural overlap, this supports using the paraphrased ToFU set to evaluate robustness to query rephrasing.
  • C.3 Examples: Representative examples of original and paraphrased ToFU questions are presented.These examples appear in Table 15.

D.1 Results on ToFU · D.2 Results on MUSE News

The appendix reports additional ToFU unlearning evaluations across forget ratios, paraphrased questions, and LLaMA-3.2-3B settings, alongside supplementary MUSE News results using Llama-2-7B. The cited passages identify evaluation conditions and reporting conventions but do not provide the table values.

  • D.1 Results on ToFU: ToFU results include unlearning performance at a 5% forget ratio under the Table 1 evaluation notation.The passage provides the table condition but no numerical results.
  • D.1 Results on ToFU: ToFU also evaluates a 5% forget ratio under paraphrased questions using the Table 2 notation.The passage specifies the paraphrased-question condition but no numerical results.
  • D.1 Results on ToFU: Additional ToFU results assess unlearning with a 10% forget ratio under the Table 1 notation.The passage provides the forget-ratio condition but no numerical results.
  • D.1 Results on ToFU: A corresponding 10% ToFU evaluation uses paraphrased questions and the Table 2 notation.The passage specifies the evaluation condition but no numerical results.
  • D.1 Results on ToFU: LLaMA-3.2-3B ToFU experiments cover a 1% forget ratio under the Table 1 notation.The passage identifies the model and condition but no numerical results.
  • D.1 Results on ToFU: Further LLaMA-3.2-3B ToFU evaluations use 5% and 10% forget ratios under the Table 1 notation.The two cited tables specify separate forget-ratio settings but provide no numerical results in the supplied passages.
  • D.2 Results on MUSE News: Supplementary MUSE News results use Llama-2-7B and compare Original, Retrain, and unlearning methods.Original is trained on the full dataset, Retrain excludes the forget set Df, and unlearning results are mean ± standard deviation over three random seeds.

E Computational Overhead of SAUL · F Threshold Calibration via Margin-Based Lower Bounds

SAUL incurs measurable runtime and memory overhead from sharpness-aware updates, while improving robustness under query rephrasing. Threshold calibration depends on the chosen forget-side constraint: the main loss-scale threshold differs from a margin-based threshold interpretable as a target forget rate.

  • E Computational Overhead of SAUL: SAUL’s sharpness-aware updates require extra forward and backward passes at each iteration, creating additional computational overhead.The runtime and peak-memory trade-off is evaluated on ToFU with a 10% forget ratio using LLaMA-3.2-3B.
  • E Computational Overhead of SAUL: 267 seconds for SAUL exceeds 221 seconds for SAUL w/o SAM, 164 seconds for Dual-AdamW, and 232 seconds for Sharp Min–Max.These are runtime comparisons on ToFU with a 10% forget ratio using LLaMA-3.2-3B.
  • E Computational Overhead of SAUL: 84.43 GB for SAUL exceeds the 79.92 GB peak memory usage of Dual-AdamW.Peak memory usage is reported for ToFU with a 10% forget ratio using LLaMA-3.2-3B.
  • E Computational Overhead of SAUL: 66.04 versus 64.31 is SAUL’s GPT-paraphrased ToFU harmonic mean improvement over its w/o SAM variant.This robustness comparison uses LLaMA-3.2-1B and concerns query rephrasing.
  • F Threshold Calibration via Margin-Based Lower Bounds: The main implementation defines the forget-side constraint using cross-entropy forget loss Lf, so α is measured on the loss scale.The threshold α denotes a prescribed satisfaction level for the chosen forget-side constraint.
  • F Threshold Calibration via Margin-Based Lower Bounds: Under a margin-based instantiation, α retains the same symbol but is specific to the prediction-level certificate Φf rather than the main loss-scale threshold.The numerical values of the two thresholds need not coincide.
  • F Threshold Calibration via Margin-Based Lower Bounds: For the margin-based constraint, α ∈[0, 1] directly represents a target forget rate.This interpretation applies to the margin-based instantiation, not necessarily to the main cross-entropy-based implementation.
  • F Threshold Calibration via Margin-Based Lower Bounds: The margin-based design uses cross-entropy as a retain-side upper surrogate and maximizes a forget-side lower surrogate to certify sequence-level exact-match failure.The retain surrogate targets token-level 0/1 error, while the forget certificate targets sequence-level exact-match failure.

F.1 Prediction margin and token-level 0/1 error … G Broader Impact and Potential Risks

The paper provides certified margin-based surrogates for asymmetric retain and forget goals, then uses an augmented-Lagrangian controller to enforce forgetting targets. It also details limitations, including incomplete removal guarantees and risks of over-forgetting.

  • F.1 Prediction margin and token-level 0/1 error: A target label is the unique top-1 prediction exactly when its prediction margin exceeds zero, and token-level 0/1 error is determined by this condition.For autoregressive generation, each next-token prediction is a categorical prediction over the vocabulary.
  • F.1 Prediction margin and token-level 0/1 error: For answer sequences, a token is predicted incorrectly whenever its token-level margin is nonpositive.This extends the margin criterion from individual categorical predictions to each target token in a conditioned answer sequence.
  • F.2 Retain side: cross-entropy as an upper surrogate: Minimizing the length-normalized retain cross-entropy controls a computable upper bound on retain-side token error.An incorrect top-1 prediction implies target probability at most 1/2 and cross-entropy at least log 2.
  • F.3 Forget side: a lower surrogate for exact-match failure: The forget-side surrogate is monotone decreasing in token margin, rewards negative margins, and remains a pointwise lower bound on token-level 0/1 error.Sequence-level aggregation uses the strongest token-level failure evidence, preserving the lower-bound relation for exact-match failure.
  • F.3 Forget side: a lower surrogate for exact-match failure: α = 0.95 certifies that at least 95% of forget examples fail exact match, provided the chosen τ makes the certificate numerically attainable.The sharpness parameter τ affects the surrogate’s numerical scale and tightness while preserving the pointwise lower-bound property for τ > 0.
  • F.4 Certified SAUL formulation and interpretation of α: The certified formulation pairs an upper surrogate for retain correctness with a lower surrogate for forget-side exact-match failure.α is a prediction-level forgetting target, not a relative trade-off weight; α = 0.95 directly requires at least 95% exact-match failure.
  • F.4 Certified SAUL formulation and interpretation of α: When cross-entropy implements forgetting, the margin certificate can be evaluated post hoc because cross-entropy and the margin surrogate are directionally aligned but numerically distinct.A cross-entropy-trained model must therefore be evaluated by checking whether Φf(θ) ≥α.
  • F.5 Augmented-Lagrangian gating: The augmented-Lagrangian controller increases forget-side pressure when the constraint is violated and can deactivate it after sustained satisfaction.Projected multiplier updates provide hinge-style gating, while a quadratic penalty stabilizes multiplier dynamics; cross-entropy implementations use the corresponding loss-scale constraint.
Loading 2608.16249v1…