PaperlayerView exact revision on arXiv

Source-linked AI summary

OFMU: Optimization-Driven Framework for Machine Unlearning

Sadia Asif, Mohammad Mohammadi Amiri

arXiv:2509.22483v2cs.LGcs.AI

TL;DR

Machine unlearning must erase targeted information while preserving retained utility, yet fixed-weight scalarization can become unstable when forgetting and retention gradients conflict. OFMU uses a penalty-based bi-level framework with inner forgetting and gradient decorrelation, followed by outer utility restoration, and reports state-of-the-art trade-offs across language and vision benchmarks. Its two-loop algorithm and convergence analysis support scalable optimization, while future work remains for continual unlearning, larger models, and additional modalities.

  • Problem

    Machine unlearning must remove the influence of targeted data while preserving performance on retained data, but scalarized objectives can suffer from conflicting gradients and unstable trade-offs.

  • Method

    OFMU uses penalty-based bi-level optimization: an inner maximization forgets targeted data while decorrelating gradients, and an outer minimization restores retained utility through a scalable two-loop algorithm.

  • Results

    OFMU achieves state-of-the-art forgetting–utility trade-offs across language and vision benchmarks, with improved stability, robustness to hard-to-forget samples, and resilience against membership inference attacks.

  • Takeaways & Limitations

    OFMU provides a structured approach that prioritizes forgetting while preserving utility through gradient decorrelation, scalable optimization, and convergence guarantees.

  • Takeaways & Limitations

    Extending OFMU to continual unlearning, larger foundation models, and speech or multimodal settings remains open.

Abstract

from arXiv · show

Large language models deployed in sensitive applications increasingly require the ability to unlearn specific knowledge, such as user requests, copyrighted materials, or outdated information, without retraining from scratch to ensure regulatory compliance, user privacy, and safety. This task, known as machine unlearning, aims to remove the influence of targeted data (forgetting) while maintaining performance on the remaining data (retention). A common approach is to formulate this as a multi-objective problem and reduce it to a single-objective problem via scalarization, where forgetting and retention losses are combined using a weighted sum. However, this often results in unstable training dynamics and degraded model utility due to conflicting gradient directions. To address these challenges, we propose OFMU, a penalty-based bi-level optimization framework that explicitly prioritizes forgetting while preserving retention through a hierarchical structure. Our method enforces forgetting via an inner maximization step that incorporates a similarity-aware penalty to decorrelate the gradients of the forget and retention objectives, and restores utility through an outer minimization step. To ensure scalability, we develop a two-loop algorithm with provable convergence guarantees under both convex and non-convex regimes. We further provide a rigorous theoretical analysis of convergence rates and show that our approach achieves better trade-offs between forgetting efficacy and model utility compared to prior methods. Extensive experiments across vision and language benchmarks demonstrate that OFMU consistently outperforms existing unlearning methods in both forgetting efficacy and retained utility.

1 INTRODUCTION

Machine unlearning must remove targeted information while preserving retained utility, but existing approaches struggle with conflicting objectives, especially for difficult samples. OFMU addresses this asymmetry with hierarchical optimization, similarity-aware gradient decorrelation, scalable computation, and theoretical guarantees.

  • Fixed-weight scalarization can destabilize unlearning: forgetting-dominant updates damage retention, whereas retention-dominant updates leave sensitive information incompletely removed.
  • OFMU prioritizes forgetting in an inner objective and restores utility in an outer objective, reflecting the asymmetric importance of safety-oriented erasure and retention.
  • Hard-to-unlearn samples cause disproportionate collateral utility loss for existing methods, while OFMU mitigates this coupling through similarity-aware hierarchical updates.
  • Its contributions include a penalty-based stationary reformulation, a scalable two-loop algorithm, dynamic gradient decorrelation, and convergence analysis in convex and non-convex regimes.
  • OFMU achieves a state-of-the-art forgetting–utility trade-off across language and vision unlearning benchmarks.

2 BI-LEVEL OPTIMIZATION IN MACHINE LEARNING

Bi-level optimization separates interdependent objectives into hierarchical inner and outer problems, making it suitable for asymmetric machine-unlearning goals. OFMU adopts this structure to enforce forgetting before conditionally restoring utility.

  • Bi-level optimization uses an outer task whose feasible solutions are constrained by optimal solutions of an inner task, supporting hierarchical objectives such as meta-learning and adversarial robustness.
  • For unlearning, OFMU treats forgetting as non-negotiable and utility preservation as a secondary objective, unlike fixed-weight scalarization.

3 METHODOLOGY

OFMU formulates unlearning as hierarchical optimization over forget and retain sets, then makes it tractable with a penalty reformulation and alternating two-loop updates. The design combines gradient decorrelation, stationarity enforcement, stochastic scalability, and stability-oriented implementation choices.

  • 3.1 PRELIMINARIES AND NOTATION: The problem separates a forget set to unlearn from a retain set to preserve, with cosine similarity measuring alignment between their loss gradients.
  • 3.2 PROBLEM SETUP: The inner objective maximizes forget loss while penalizing similarity between forget and retain gradients, and the outer objective restores retain-set utility subject to enforced forgetting.
  • 3.3 BI-LEVEL FORMULATION: The bi-level formulation requires the final parameters to satisfy the inner stationarity condition ∇θΦ(θ) = 0.
  • 3.4 PENALTY-BASED SINGLE-LEVEL REFORMULATION: A penalty term converts the nested problem into an unconstrained objective, with increasing ρ progressively enforcing inner stationarity while preserving the hierarchical structure.
  • 3.5 PRACTICAL ALGORITHM AND IMPLEMENTATION: The practical algorithm alternates T inner gradient-ascent steps for forgetting and decorrelation with an outer penalized utility-preservation update.
  • 3.5 PRACTICAL ALGORITHM AND IMPLEMENTATION: Mini-batch gradients reduce per-iteration cost from O(|D|) to O(B), while Hessian-vector products implement the stationarity penalty and the two-loop scheme avoids fully solving each inner problem.

4 THEORETICAL ANALYSIS

The analysis establishes that the penalty reformulation enforces inner-objective stationarity and that the two-loop OFMU algorithm converges under convex and non-convex conditions. These results provide the theoretical foundation for the proposed method.

  • As ρ approaches infinity, accumulation points of minimizers of the penalized objective satisfy the inner stationarity condition ∇θΦ(θ) = 0.
  • Under convexity and Lipschitz-smoothness, the inner gradient-ascent loop achieves sublinear convergence with step size 0 < ηin ≤ 1/L.
  • The full penalty-based OFMU algorithm has convergence guarantees in both convex and non-convex regimes.
  • In the convex case, setting K, T = O(1/ϵ) ensures ϵ-optimality of the penalty objective.
  • In the non-convex case, OFMU converges to an ϵ-stationary point of F(θ), with stochastic-gradient variance represented by σ².

5 EXPERIMENTS

OFMU is evaluated on language and vision benchmarks to test unlearning efficacy, retained utility, baseline comparisons, and cross-domain generality. Across TOFU and CIFAR-10, it is reported to maintain a stronger balance between forgetting and retention than methods that over-optimize one objective.

  • 5.1 EXPERIMENTAL SETUP: The experiments cover TOFU forget01, forget05, and forget10 on two LLaMA architectures, plus CIFAR-10 class-wise and random forgetting settings.WMDP and CIFAR-100 results are deferred to the appendix.
  • 5.2 TOFU RESULTS: OFMU preserves competitive forgetting while retaining substantially more utility than aggressive baselines on TOFU, including MU = 0.65 on LLaMA-2 in forget05.GA attains higher raw FQ in forget05 but collapses to MU = 0.00; OFMU maintains MU close to the Retain baseline across scenarios.
  • 5.2 TOFU RESULTS: OFMU maintains high FTR across TOFU scenarios, whereas GA and GD degrade truthfulness and some NPO variants overfit to retain data.The reported FTR behavior indicates reliable rather than memorized or distorted responses after unlearning.
  • 5.3 CIFAR-10 RESULTS: OFMU achieves a balanced CIFAR-10 class-wise result of 81.51% UA, 93.51% RA, and 86.88% TA, with 59.76 MIA-Efficacy.IU reaches higher UA at 89.31% but requires repeated influence-function estimation and parameter adjustments, limiting practicality for large models.
  • 5.3 CIFAR-10 RESULTS: In random CIFAR-10 forgetting, OFMU achieves the best overall normalized performance while maintaining balance across UA, RA, TA, and MIA-Efficacy.The normalized score averages the four metrics within each scenario; higher values indicate better balance.

6 CONCLUSION AND FUTURE WORK

The paper concludes that OFMU provides stable, robust trade-offs across language and vision unlearning benchmarks with convergence guarantees in convex and non-convex regimes. It identifies continual unlearning, adaptive penalties, larger models, and additional modalities as future directions.

  • 6 CONCLUSION AND FUTURE WORK: OFMU combines a penalty-based bi-level framework and scalable two-loop algorithm to prioritize forgetting before utility preservation.The framework uses a similarity-aware penalty and is supported by convergence guarantees in convex and non-convex regimes.
  • 6 CONCLUSION AND FUTURE WORK: Empirical results report improved stability, robustness to hard-to-forget samples, and stronger resilience against membership inference attacks than existing approaches.
  • 6 CONCLUSION AND FUTURE WORK: Future work includes continual unlearning, adaptive penalty schedules, alternative gradient similarity measures, larger foundation models, and speech or multimodal learning.

7 APPENDIX

The appendix defines the optimization notation and records proof steps for the penalty formulation and convergence analysis. Its central theoretical claim is that increasing the penalty enforces stationarity of the inner objective at accumulation points.

  • 7.1 NOTATION SUMMARY •: The notation distinguishes retain and forget losses, the similarity-aware inner objective Φ(θ), the penalty objective F(θ), and two-loop learning rates and iteration counts.It also defines evaluation metrics including FQ, MU, FTR, UA, RA, TA, and MIA-Efficacy.
  • 7.2.1 PROOF OF LEMMA 1: As ρ increases, any accumulation point of minimizers must satisfy ∇θΦ(θ∗) = 0 because otherwise the penalty term diverges and contradicts bounded minimization.
  • 7.2.2 PROOF OF LEMMA 2: The convex proof derives an average-suboptimality bound by combining the update rule with convexity and smoothness inequalities.The supplied proof passages state the resulting average suboptimality without providing the complete bound.

7.3 CONVERGENCE GUARANTEES FOR PENALTY-BASED OFMU

OFMU is analyzed under convex and non-convex settings using an alternating inner maximization and outer minimization scheme with explicit smoothness, boundedness, and step-size assumptions. The analysis establishes sublinear convergence to an ε-optimal solution in the convex case and an ε-stationary point in the non-convex case.

  • 7.3.1 CONVERGENCE ANALYSIS: CONVEX CASE: The convex analysis alternates T gradient-ascent steps on Φ with an inexact gradient-descent update on F, whose error is controlled by inner maximization accuracy.The inner error is O(1/T), yielding a squared gradient error of O(1/T^2).
  • 7.3.1 CONVERGENCE ANALYSIS: CONVEX CASE: In the convex case, OFMU reaches an O(ε)-suboptimal solution with K = O(1/ε) outer iterations and T = O(1/ε) inner steps.The guarantee relies on convexity and smoothness, with the inner-loop error contributing O(K/T^2) cumulatively.
  • 7.3.2 CONVERGENCE ANALYSIS: NON-CONVEX CASE: The non-convex analysis bounds the final inner iterate through the expected squared gradient norm and combines this approximation error with outer-loop descent.The assumptions include L-smooth objectives, unbiased inner stochastic gradients with bounded variance, a lower-bounded smooth penalty objective, and bounds on gradient and Hessian norms.
  • 7.3.2 CONVERGENCE ANALYSIS: NON-CONVEX CASE: In the non-convex case, OFMU converges in expectation to an ε-stationary point of F when K = O(1/ε) and T = O(1/ε).The rate depends explicitly on K, T, the penalty parameter ρ, curvature bound H, and stochastic-gradient variance σ^2.

7.4 COMPUTATIONAL COMPLEXITY AND PRACTICAL EFFICIENCY

OFMU remains computationally practical because its inner loop uses a small number of updates and its outer loop computes Hessian information through Hessian–vector products rather than explicitly forming Hessians. Consequently, the method’s cost stays close to standard gradient-based optimization.

  • 7.4.1 PER-ITERATION COST OF THE INNER AND OUTER LOOPS: Each inner step costs Θ(Bd), while the outer step combines a retain-loss gradient with one HVP of the same order as a gradient.Here B is minibatch size and d is the number of trainable parameters.
  • 7.4.2 WHY HESSIAN–VECTOR PRODUCTS ARE COMPUTATIONALLY CHEAP: Hessian–vector products add negligible overhead because Pearlmutter’s method computes them with one additional backward pass and preserves the algorithm’s asymptotic scaling.The HVP cost is approximately the same as a gradient up to a small constant factor.
  • 7.4.3 WHY THE INNER LOOP IS COMPUTATIONALLY LIGHT: A small fixed inner budget, such as T = 5 or 10, is sufficient because the penalty term completes enforcement of inner-objective stationarity.The inner loop only needs to make progress toward reducing the stationarity violation rather than converge fully.
  • 7.4.4 TOTAL COMPLEXITY: Because T ≪ K and T is typically constant, OFMU’s total cost over K outer iterations remains close to standard gradient-based methods.Table 4 compares OFMU’s per-iteration complexity with standard unlearning baselines.

7.5 AUXILIARY RESULTS AND ABLATION STUDY

Auxiliary experiments across language and vision benchmarks show that OFMU consistently balances forgetting efficacy, retained utility, robustness, and computational practicality. Ablations and sensitivity analyses attribute these results to hierarchical optimization, similarity-aware gradient decorrelation, and controlled inner-loop updates.

  • 7.5.1 WMDP RESULTS: 72.8% unlearning efficacy on Biosecurity, 70.4% on Cybersecurity, and 74.6% utility accuracy on MMLU, with OFMU significantly outperforming all WMDP baselines.Gains over RMU are +1.5, +2.1, and +0.4 points respectively, with p < 0.05.
  • 7.5.2 CIFAR-100 RESULTS: On CIFAR-100 class-wise forgetting, OFMU achieves 74.1% UA, 71.9% RA, 69.3% TA, and 48.2 MIA-Efficacy while balancing forgetting and utility.For random forgetting, OFMU reaches 6.7% UA, 70.2% RA, and 67.8% TA, slightly exceeding retraining on UA.
  • 7.5.3 OVERALL PERFORMANCE SCORE CALCULATION: The normalized Overall Performance Score averages benchmark metrics to expose trade-offs among forgetting efficacy, utility preservation, generalization, and security across language and vision tasks.TOFU uses FQ, MU, and FTR, whereas CIFAR-10/100 uses UA, RA, TA, and MIA.
  • 7.5.4 HARD IN-SCOPE EVALUATION AND ROBUSTNESS: Across paraphrase, multi-hop, and cross-lingual TOFU probes, OFMU provides the most consistent balance, avoiding both RMU’s under-forgetting and GA’s utility collapse.OFMU is best or second-best on FQ while keeping MU and FTR competitive.
  • 7.5.5 SAMPLE-SELECTION SENSITIVITY: OFMU shows the lowest FQ variance and low MU variance across five random TOFU forget-set draws, although rankings can flip between other methods.This supports OFMU’s stability under sample-selection variation.
  • 7.5.6 MEASURING UNLEARNING DIFFICULTY: The Unlearning Difficulty Index combines update pressure, forget–retain gradient conflict, and loss margin, and OFMU shows the weakest coupling between difficulty and collateral utility loss.GA and GDiff show strong prediction of utility loss for high-UDI samples, whereas OFMU’s similarity-aware penalty and hierarchical updates regulate overcorrection.
  • 7.5.7 EMBEDDING ALIGNMENT WITH RETRAIN: OFMU improves embedding alignment with a retrained reference on hard samples while remaining realistic, whereas RMU preserves utility but under-forgets and misaligns on easy samples.The hard-sample robustness pattern is consistent with OFMU’s balanced performance elsewhere.
  • 7.5.8 COMPONENT-WISE ABLATION OF OFMU: Removing similarity-aware decorrelation lowers hard-sample embedding similarity from 0.73 to 0.71, while removing penalty reformulation reduces FQ from 0.38 to 0.33.Both components are therefore critical, although the full model’s gains are modest and consistent.

7.6 RELATED WORK

Related work spans input-based, data-based, and model-based unlearning methods. These approaches trade deployment simplicity, data-construction demands, suppression strength, and retained utility, motivating OFMU’s hierarchical treatment of forgetting and retention.

  • Input-based methods: Input-based methods modify prompts or generation controls without parameter updates, but adversarial queries, paraphrases, and prompt injection can expose residual memorization.Because the underlying parameters remain unchanged, these methods do not remove the internally encoded knowledge.
  • Data-based methods: Data-based methods fine-tune on curated refusal, negated, counterfactual, or preference examples, but require task-specific data and risk semantic drift or unrelated-domain degradation.Alternate Preference Optimization combines negative feedback on forget examples with positive in-domain alternatives.
  • Model-based methods: Model-based methods directly modify parameters and can suppress memorized knowledge more effectively, but gradient ascent, regularization, adapters, and neuron interventions face instability or conflicting forgetting–retention objectives.Scalarized fixed-weight formulations can cause catastrophic utility loss or incomplete forgetting in high-dimensional, non-convex LLM settings.
  • Model-based methods: OFMU advances model-based unlearning with penalty reformulation and similarity-aware gradient decorrelation, prioritizing forgetting in an inner problem and restoring utility in an outer loop.This differs from prior scalarization-based or heuristic bilevel approaches.

7.7 HESSIAN-VECTOR PRODUCT VIA AUTOMATIC DIFFERENTIATION

The Hessian-vector product enables OFMU’s penalty computation without explicitly constructing the Hessian. Pearlmutter’s method makes this second-order operation practical, while increasing ρk progressively enforces inner-objective stationarity.

  • Hessian-vector product computation: Explicit Hessian construction would require O(d^2) time and memory, which is prohibitive for large parameter dimension d.The Hessian-vector product instead uses the directional derivative of the gradient.
  • Hessian-vector product computation: Pearlmutter’s trick computes the Hessian-vector product at the cost of a single gradient evaluation, avoiding explicit Hessian formation.This makes the second-order penalty compatible with large-scale optimization.
  • Penalty enforcement: The penalty parameter ρk increases during training through a predefined or adaptive schedule, enforcing the stationarity constraint more strictly as optimization progresses.The schedule can also depend on the norm of ∇θΦ(θ).
Loading 2509.22483v2…