🤖 machine learning

OFMU: Optimization-Driven Framework for Machine Unlearning

This paper introduces OFMU, a penalty-based bi-level optimization framework that resolves the instability and utility degradation of existing machine unlearning methods by explicitly prioritizing data forgetting through a similarity-aware penalty while preserving model performance on retained data, achieving superior efficacy and convergence guarantees across vision and language benchmarks.

Original authors: Sadia Asif, Mohammad Mohammadi Amiri

Published 2026-06-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Sadia Asif, Mohammad Mohammadi Amiri

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart, well-read librarian (the AI model) who has memorized millions of books. One day, a customer asks the librarian to "unlearn" a specific, sensitive story they told them—perhaps because the story was a secret, copyrighted, or just no longer true. The customer wants the librarian to forget that story completely but still remember everything else perfectly so they can continue helping with other questions.

This is the problem of Machine Unlearning. The paper introduces a new method called OFMU to solve this tricky balancing act.

Here is how the paper explains the problem and their solution, using simple analogies:

The Problem: The "Tug-of-War"

Most current methods try to make the librarian forget the bad story by pulling on a rope in the opposite direction. However, they treat "forgetting the bad story" and "remembering the good stories" as two equal tasks happening at the same time.

The paper argues this is like a Tug-of-War where two teams are pulling on the same rope with equal strength.

If the "forgetting" team pulls too hard, the librarian forgets everything, including the good stories (the model breaks).
If the "remembering" team pulls too hard, the librarian never actually forgets the bad story (the unlearning fails).
Because the directions of "forgetting" and "remembering" often conflict, the librarian gets confused, leading to unstable results.

The Solution: OFMU (The "Strict Boss" Approach)

The authors propose OFMU, which changes the rules of the game. Instead of a Tug-of-War, they set up a hierarchical structure (a boss and an employee).

The Inner Loop (The Strict Boss): First, the system focuses only on the "forgetting" task. It acts like a strict boss who says, "Your only job right now is to erase this specific memory." To make sure this doesn't accidentally damage the other memories, it uses a special tool called a similarity-aware penalty.
- Analogy: Imagine the librarian is erasing a specific page from a book. The "penalty" is like a guard who watches the librarian's hand. If the librarian's hand starts moving in a way that would accidentally tear out a different page (a good memory), the guard stops them. This ensures the "forgetting" action doesn't mess up the "retaining" action.
The Outer Loop (The Helpful Employee): Once the "forgetting" is done and the bad memory is gone, the system switches to the "remembering" task. It gently tunes the librarian back up so they are sharp and helpful again, but only after the bad memory is confirmed to be gone.

How They Do It (The Two-Loop Algorithm)

The paper describes a "two-loop" process, which is like a rehearsal and performance routine:

Rehearsal (Inner Loop): The librarian practices forgetting the bad story for a few minutes. They don't need to master it perfectly every single time; they just need to get close to the point where the memory is gone.
Performance (Outer Loop): Once the practice is done, the librarian performs the main task: answering questions about the good stories. They use a mathematical "penalty" to ensure they didn't slip up and accidentally bring the bad story back.

The authors prove mathematically that this method is stable and converges (works reliably) even for very complex, large models.

The Results: Better Balance

The paper tested this on two types of "librarians":

Language Models (LLMs): Like LLaMA, used for text and conversation.
Vision Models: Used for recognizing images (like identifying cats vs. dogs).

What they found:

Old methods were like a sledgehammer: they either smashed the bad memory but broke the whole library (losing all utility), or they kept the library safe but failed to remove the bad memory.
OFMU was like a surgeon: it successfully removed the bad memory while keeping the rest of the library in perfect condition.
In tests, OFMU was better at "unlearning" difficult items without causing the model to collapse or become useless. It also made the model more secure against people trying to guess what data was inside it (a test called Membership Inference).

Summary

In short, the paper says: "Don't try to forget and remember at the same time with equal weight. First, force the model to forget the specific bad data using a smart, protective penalty system. Once that is done, fix the model's general performance. This 'forget-first' hierarchy creates a much more stable and effective way to clean AI models."

Technical Summary of OFMU: Optimization-Driven Framework for Machine Unlearning

Problem Statement
Large Language Models (LLMs) and other deep learning models trained on massive corpora often contain sensitive, copyrighted, or harmful information. Regulations like the GDPR mandate a "right to be forgotten," requiring models to remove the influence of specific data (the forget set, $D_f$ ) without retraining from scratch, while maintaining performance on the remaining data (the retain set, $D_r$ ).

Existing model-based unlearning methods typically formulate this as a scalarized multi-objective problem, combining forgetting and retention losses with fixed weights. The paper identifies three critical limitations in this approach:

Static Weighting: Fixed weights fail to adapt to the dynamic nature of unlearning, where early steps should prioritize forgetting and later steps should focus on utility restoration.
Instability: Scalarization leads to unstable training dynamics. If the forget objective dominates, model utility collapses; if the retain objective dominates, forgetting remains incomplete.
Gradient Entanglement: On "hard-to-unlearn" samples, gradients for forgetting and retention are strongly coupled. Aggressive updates to remove knowledge often cause disproportionate collateral damage to retained knowledge.

Methodology: OFMU
The authors propose OFMU (Optimization-Driven Framework for Machine Unlearning), a penalty-based bi-level optimization framework that explicitly separates the objectives into a hierarchical structure. The core philosophy is that forgetting is a non-negotiable primary objective, while utility preservation is a secondary goal conditional on successful forgetting.

Bi-Level Formulation:
- Inner Maximization (Forgetting): The inner problem seeks parameters $\theta^*$ that maximize the loss on the forget set while minimizing the similarity between the forget and retain gradients. This is defined by the objective $\Phi(\theta) = L_f(\theta) - \beta \cdot \text{Sim}(\nabla L_f, \nabla L_r)$ , where $\beta$ controls the gradient decorrelation penalty.
- Outer Minimization (Utility): The outer problem minimizes the loss on the retain set $L_r(\theta)$ , subject to the constraint that the parameters are stationary points of the inner objective ( $\nabla \Phi(\theta) = 0$ ).
Penalty-Based Reformulation:
Directly solving the bi-level problem is computationally prohibitive due to the need for repeated inner maximization and higher-order derivatives. OFMU reformulates this into a single-level unconstrained problem using a penalty term:
$F(\theta) = L_r(\theta) + \rho \|\nabla \Phi(\theta)\|^2$
Here, $\rho$ is a penalty parameter that enforces the stationarity of the inner objective. As $\rho \to \infty$ , the minimizer of $F(\theta)$ satisfies the original bi-level constraints.
Two-Loop Algorithm:
To ensure scalability, the authors develop a two-loop algorithm:
- Inner Loop: Performs $T$ steps of gradient ascent on $\Phi(\theta)$ to maximize forgetting and decorrelate gradients.
- Outer Loop: Performs a gradient descent step on $F(\theta)$ to restore utility. This step involves computing a Hessian-vector product (via automatic differentiation) to handle the penalty term $\|\nabla \Phi(\theta)\|^2$ .
- Penalty Schedule: The penalty parameter $\rho$ is increased over iterations to gradually strengthen the enforcement of the inner stationarity condition.

Key Contributions

Novel Framework: Introduction of OFMU, a bi-level framework that prioritizes forgetting over utility preservation, capturing the inherent asymmetry of the unlearning task.
Scalable Algorithm: Development of a two-loop algorithm with provable convergence guarantees that avoids the computational bottlenecks of traditional bi-level optimization (e.g., full inner convergence).
Similarity-Aware Penalty: Design of a penalty term that dynamically decorrelates forget and retain gradients, mitigating destructive interference during updates.
Theoretical Analysis: Rigorous convergence analysis showing that the algorithm converges to an $\epsilon$ -stationary point in both convex and non-convex settings, with specific bounds on suboptimality.

Experimental Results
The authors evaluated OFMU on language and vision benchmarks:

Language (TOFU Benchmark): Using LLaMA-2-7B and LLaMA-3.2-1B, OFMU was tested on removing 1%, 5%, and 10% of the dataset.
- Performance: OFMU achieved a superior balance between Forgetting Quality (FQ), Model Utility (MU), and Forget Truth Ratio (FTR).
- Comparison: Unlike Gradient Ascent (GA) and Gradient Difference (GradDiff), which achieved high FQ but caused utility to collapse (MU $\approx$ 0), OFMU maintained high utility (e.g., MU $\approx$ 0.65 on LLaMA-2 for 5% forgetting) while enforcing strong forgetting. It outperformed baselines like NPO and RMU in overall normalized scores.
Vision (CIFAR-10): Tested under class-wise and random forgetting settings.
- Class-wise: OFMU achieved 81.51% Unlearning Accuracy (UA) and 93.51% Retain Accuracy (RA), offering a better trade-off than Influence Unlearning (IU) which had higher UA but was computationally expensive.
- Random: In the challenging random forgetting scenario (10% of data), OFMU achieved 7.71% UA, slightly outperforming retraining (6.79%), while maintaining robustness against Membership Inference Attacks (MIA).

Significance and Claims
The paper claims that OFMU addresses the fundamental instability of scalarized unlearning methods by recognizing the hierarchical nature of the problem. By enforcing forgetting as a primary constraint before optimizing for utility, OFMU achieves state-of-the-art trade-offs across both language and vision tasks. The authors assert that their approach provides:

Stability: Consistent performance even on hard-to-unlearn samples where other methods fail.
Robustness: Stronger resilience against membership inference attacks compared to existing baselines.
Theoretical Grounding: A principled approach with convergence guarantees in non-convex regimes, moving beyond heuristic weighting schemes.

The paper concludes that while OFMU represents a significant step forward, future work could explore continual unlearning scenarios, adaptive penalty schedules, and application to larger foundation models and multimodal data.