OFMU: Optimization-Driven Machine Unlearning

• OFMU is an optimization-driven machine unlearning framework that formulates forgetting as a penalty-based bi-level problem prioritizing targeted data removal over utility restoration.
• It integrates a similarity-aware penalty to decorrelate conflicting gradient directions, ensuring that forgetting does not excessively degrade model performance.
• The framework employs a scalable two-loop algorithm with proven convergence guarantees, effectively balancing regulatory, privacy, and safety requirements in various benchmarks.

Searching arXiv for the OFMU paper and a small set of directly relevant machine unlearning context papers. OFMU, short for Optimization-Driven Framework for Machine Unlearning, is a machine unlearning framework designed for settings in which a trained model must remove the influence of targeted data while preserving performance on the remaining data. The framework is motivated by deployments of LLMs in sensitive applications that require the ability to unlearn specific knowledge, such as user requests, copyrighted materials, or outdated information, without retraining from scratch for regulatory compliance, user privacy, and safety. OFMU formulates unlearning as a penalty-based bi-level optimization problem that explicitly prioritizes forgetting and preserves retention through a hierarchical structure, rather than relying on scalarization by a weighted sum of forgetting and retention losses (Asif et al., 26 Sep 2025).

1. Conceptual setting and problem formulation

Machine unlearning, in the formulation used by OFMU, is the task of removing the influence of a forget set $\mathcal{D}_f$ from a trained model while preserving utility on the retain set $\mathcal{D}_r$ (Asif et al., 26 Sep 2025). The motivating difficulty is that common scalarization-based approaches combine forgetting and retention losses using a weighted sum, but this often results in unstable training dynamics and degraded model utility due to conflicting gradient directions (Asif et al., 26 Sep 2025).

The framework is built around an explicitly asymmetric priority structure. Forgetting is treated as paramount, while utility restoration proceeds only once sufficient forgetting is achieved. This hierarchical view contrasts with methods that optimize forgetting and retention as if they were commensurable objectives with a fixed trade-off coefficient (Asif et al., 26 Sep 2025).

OFMU expresses this asymmetry through a bi-level program over model parameters $\theta$:

$$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$

Here, $\mathcal{L}_r(\theta)$ denotes the empirical loss on the retain set, $\mathcal{L}_f(\theta)$ the empirical loss on the forget set, $\mathrm{Sim}(a,b)$ the cosine similarity between vectors $a$ and $b$, and $\beta > 0$ the regularization strength for the similarity-aware penalty (Asif et al., 26 Sep 2025).

The paper further defines

$\mathcal{D}_r$0

so that the inner problem becomes a maximization of forgetting subject to gradient decorrelation, while the outer problem restores retain-set utility (Asif et al., 26 Sep 2025).

2. Similarity-aware forgetting objective

A central feature of OFMU is the similarity-aware penalty embedded in the inner maximization. By including $\mathcal{D}_r$1, the framework explicitly discourages directions in parameter space where forgetting and retention conflict, making gradient updates more orthogonal and limiting collateral utility loss, particularly on “hard” samples (Asif et al., 26 Sep 2025).

The empirical loss terms are written as

$\mathcal{D}_r$2

and the gradient-alignment term uses cosine similarity:

$\mathcal{D}_r$3

These definitions formalize the claim that instability arises when forget and retain gradients are highly entangled, especially on hard-to-unlearn samples (Asif et al., 26 Sep 2025).

This structure also explains the framework’s departure from scalarization. Weighted-sum methods are described as non-adaptive, unable to shift focus as unlearning proceeds, and vulnerable to catastrophic forgetting or incomplete unlearning. OFMU instead imposes a hierarchy in which forgetting is enforced first and utility restoration is conditioned on the resulting state (Asif et al., 26 Sep 2025). A plausible implication is that the method is intended not merely to trade off two losses, but to encode the precedence relation that often exists in privacy, safety, or compliance-driven unlearning requests.

3. Penalty-based reformulation and optimization algorithm

Direct bi-level optimization is described as computationally intractable for modern models. OFMU therefore converts the problem into a single-level penalty-based optimization:

$\mathcal{D}_r$4

where $\mathcal{D}_r$5 is an increasing penalty parameter (Asif et al., 26 Sep 2025). The penalty enforces the stationarity condition $\mathcal{D}_r$6, indicating a maximum for the forgetting objective, and in the limit $\mathcal{D}_r$7, any solution $\mathcal{D}_r$8 satisfies the original bi-level constraint (Asif et al., 26 Sep 2025).

To make this practical, the framework uses a scalable two-loop algorithm. The inner loop performs forgetting and decorrelation by running $\mathcal{D}_r$9 steps of gradient ascent on $\theta$0:

$\theta$1

with starting point $\theta$2 (Asif et al., 26 Sep 2025).

The outer loop then carries out utility preservation and penalty enforcement through a gradient descent step on $\theta$3:

$\theta$4

where

$\theta$5

The Hessian-vector term is computed efficiently using automatic differentiation, the penalty schedule $\theta$6 is gradually increased, and all steps are computed using mini-batches for scalability (Asif et al., 26 Sep 2025).

The method is therefore “optimization-driven” in a literal sense: the algorithmic design follows directly from the bi-level constrained formulation rather than from heuristic modifications of a single objective. This suggests that OFMU is intended as a principled alternative to ad hoc unlearning updates.

4. Theoretical guarantees

OFMU includes explicit theoretical guarantees for both convex and non-convex settings (Asif et al., 26 Sep 2025). The paper states a penalty stationarity guarantee: as $\theta$7, every accumulation point of minimizers of $\theta$8 satisfies $\theta$9, meaning that the forgetting constraint is enforced (Asif et al., 26 Sep 2025).

For the inner maximization, the paper gives a convergence statement for convex, smooth $$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$0. After $$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$1 steps of gradient ascent,

$$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$2

with the proof leveraging standard convex optimization analysis (Asif et al., 26 Sep 2025).

For the full two-loop algorithm, the paper reports the following rates. In the convex case, after $$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$3 outer iterations with $$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$4 inner steps,

$$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$5

Accordingly, with $$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$6, one obtains $$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$7-optimality (Asif et al., 26 Sep 2025). In the non-convex case,

$$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$8

where $$\begin{align} & \min_{\theta \in \mathbb{R}^d} \quad \mathcal{L}_r(\theta) \ & \text{subject to} \quad \theta \in \arg\max_{\theta'} \left[ \mathcal{L}_f(\theta') - \beta \cdot \mathrm{Sim} \big(\nabla_\theta \mathcal{L}_f(\theta'), \nabla_\theta \mathcal{L}_r(\theta')\big) \right]. \end{align}$$9 is the stochastic gradient variance (Asif et al., 26 Sep 2025).

These results are presented as evidence that the method avoids trivial solutions and yields efficient, stable, and practical convergence even for deep, non-convex models (Asif et al., 26 Sep 2025). A plausible implication is that the penalty formulation is not merely a surrogate objective, but the mechanism through which the original hierarchy between forgetting and retention is enforced in a scalable manner.

5. Empirical evaluation and benchmarks

OFMU is evaluated on both language and vision benchmarks, specifically TOFU and WMDP for language, and CIFAR-10/100 for vision (Asif et al., 26 Sep 2025). The reported metrics differ by domain. For language, the paper uses FQ (Forget Quality), MU (Model Utility), and FTR (Forget Truth Ratio). For vision, it uses UA (Unlearning Accuracy), RA (Retain Accuracy), TA (Test accuracy), and MIA-Efficacy (Membership Inference robustness) (Asif et al., 26 Sep 2025).

On TOFU, using synthetic QA with LLaMA-2/3, OFMU is reported to achieve among the highest overall composite scores across FQ, MU, and FTR in all forgetting scenarios (Asif et al., 26 Sep 2025). Competing methods such as GA, GradDiff, and NPO are described as tending either to achieve high FQ while destroying utility or to preserve utility while under-forgetting (Asif et al., 26 Sep 2025). On WMDP, the paper reports that OFMU attains the largest gains in accuracy in safety-critical domains and consistently outperforms competitive baselines (Asif et al., 26 Sep 2025).

On CIFAR-10 and CIFAR-100, including class-wise and random forgetting settings, OFMU is reported to maintain high unlearning performance, competitive overall accuracy, and superior resistance to membership inference attacks (Asif et al., 26 Sep 2025). Under random forgetting, the paper states that OFMU outperforms retrain in some scenarios (Asif et al., 26 Sep 2025). Across these benchmarks, the framework is described as achieving the top normalized overall score among all considered metrics (Asif et al., 26 Sep 2025).

The empirical section also reports robustness findings. In harsh scenarios involving paraphrases, multi-hop reasoning, and cross-lingual inputs, OFMU maintains a superior balance in FQ, MU, and FTR compared to baselines (Asif et al., 26 Sep 2025). It exhibits the lowest variance across different random forget sets, indicating stability, and achieves the weakest correlation between unlearning difficulty and collateral utility loss among the baselines, which the paper attributes to the similarity-aware inner penalty (Asif et al., 26 Sep 2025).

6. Contributions, interpretation, and limitations of scope

The paper presents OFMU as the first scalable, theoretically rigorous unlearning framework that explicitly prioritizes forgetting via inner-loop maximization and restores model utility in a structured outer-loop minimization (Asif et al., 26 Sep 2025). Its principal contributions are the penalty-based bi-level formulation, the similarity-aware gradient decorrelation term, a scalable two-loop algorithm with provable convergence guarantees under both convex and non-convex regimes, and empirical demonstrations across language and vision benchmarks (Asif et al., 26 Sep 2025).

The framework addresses a common misconception in machine unlearning: that forgetting and retention can always be handled adequately by a single weighted loss. OFMU’s formulation is built on the opposite premise, namely that forgetting and retention can have conflicting gradients and that this conflict is especially consequential in hard-to-unlearn cases (Asif et al., 26 Sep 2025). The similarity-aware penalty is therefore not an auxiliary regularizer in the ordinary sense; it is part of the operational definition of how forgetting should proceed without excessive collateral damage.

At the same time, the paper’s claims are bounded by the experimental settings it reports. The evidence is drawn from TOFU, WMDP, and CIFAR-10/100, with evaluations centered on FQ, MU, FTR, UA, RA, TA, and MIA-Efficacy (Asif et al., 26 Sep 2025). This suggests that OFMU should be understood as a framework with demonstrated breadth across language and vision benchmarks rather than as a universal guarantee for all unlearning regimes. The paper’s emphasis on convergence and stability indicates that its main intervention is methodological: it reorganizes the optimization problem so that the asymmetry between forgetting efficacy and retained utility is directly reflected in the training dynamics.

7. Position within machine unlearning research

Within machine unlearning research, OFMU belongs to a line of work concerned with removing targeted information from already trained models while avoiding full retraining. The framework is explicitly motivated by applications involving regulatory compliance, user privacy, and safety, and by the need to remove user requests, copyrighted materials, or outdated information from LLMs (Asif et al., 26 Sep 2025).

Its defining distinction is the rejection of scalarization as the primary organizing principle. Instead of a weighted-sum objective, OFMU adopts a hierarchical optimization structure in which forgetting is enforced through inner maximization and utility is restored through outer minimization (Asif et al., 26 Sep 2025). The paper argues that this design yields better trade-offs between forgetting efficacy and model utility compared to prior methods, and its empirical results are presented as consistent with that claim across both vision and language tasks (Asif et al., 26 Sep 2025).

For technically oriented readers, the importance of OFMU lies less in any single benchmark score than in the optimization perspective it advances: machine unlearning is treated as a constrained, gradient-conflicted, and inherently hierarchical problem. The resulting framework combines penalty methods, bilevel reasoning, gradient similarity control, and scalable stochastic optimization into a single formulation intended to be both theoretically analyzable and practically applicable (Asif et al., 26 Sep 2025).