---
title: 'OFMU: Optimization-Driven Machine Unlearning'
url: https://www.emergentmind.com/topics/ofmu
type: topic
---

# OFMU: Optimization-Driven Machine Unlearning

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 large language models 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 [2509.22483].

## 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$ [2509.22483]. 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 [2509.22483].

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 [2509.22483].

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 [2509.22483].

The paper further defines
$$
\Phi(\theta) := \mathcal{L}_f(\theta) - \beta \cdot \mathrm{Sim}\Big(\nabla_\theta \mathcal{L}_f(\theta), \nabla_\theta \mathcal{L}_r(\theta)\Big),
$$
so that the inner problem becomes a maximization of forgetting subject to gradient decorrelation, while the outer problem restores retain-set utility [2509.22483].

## 2. Similarity-aware forgetting objective

A central feature of OFMU is the **similarity-aware penalty** embedded in the inner maximization. By including $\mathrm{Sim}\big(\nabla_\theta \mathcal{L}_f, \nabla_\theta \mathcal{L}_r\big)$, 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 [2509.22483].

The empirical loss terms are written as
$$
\mathcal{L}_r(\theta) = \frac{1}{|\mathcal{D}_r|} \sum_{(x, y) \in \mathcal{D}_r} \ell(f_\theta(x), y), \quad
\mathcal{L}_f(\theta) = \frac{1}{|\mathcal{D}_f|} \sum_{(x, y) \in \mathcal{D}_f} \ell(f_\theta(x), y),
$$
and the gradient-alignment term uses cosine similarity:
$$
\mathrm{Sim}(\nabla_\theta \mathcal{L}_f, \nabla_\theta \mathcal{L}_r) =
\frac{\langle \nabla_\theta \mathcal{L}_f, \nabla_\theta \mathcal{L}_r \rangle}{\|\nabla_\theta \mathcal{L}_f\| \| \nabla_\theta \mathcal{L}_r \|}.
$$
These definitions formalize the claim that instability arises when forget and retain gradients are highly entangled, especially on hard-to-unlearn samples [2509.22483].

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 [2509.22483]. 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:
$$
F(\theta) = \mathcal{L}_r(\theta) + \rho \cdot \|\nabla_\theta\Phi(\theta)\|^2,
$$
where $\rho > 0$ is an increasing penalty parameter [2509.22483]. The penalty enforces the stationarity condition $\nabla_\theta\Phi(\theta) = 0$, indicating a maximum for the forgetting objective, and in the limit $\rho \to \infty$, any solution $\theta^*$ satisfies the original bi-level constraint [2509.22483].

To make this practical, the framework uses a scalable **two-loop algorithm**. The inner loop performs forgetting and decorrelation by running $T$ steps of gradient ascent on $\Phi(\theta)$:
$$
\theta'^{(t+1)} = \theta'^{(t)} + \eta_{\text{in}} \nabla_\theta \Phi(\theta'^{(t)}),
$$
with starting point $\theta'^{(0)} = \theta^{(k)}$ [2509.22483].

The outer loop then carries out utility preservation and penalty enforcement through a gradient descent step on $F(\theta)$:
$$
\theta^{(k+1)} = \theta_{\text{in}}^{(k)} - \eta_{\text{out}} \nabla_\theta F(\theta_{\text{in}}^{(k)}),
$$
where
$$
\nabla_\theta F(\theta) = \nabla_\theta \mathcal{L}_r(\theta) + 2\rho \, \nabla^2_\theta \Phi(\theta)\, \nabla_\theta \Phi(\theta).
$$
The Hessian-vector term is computed efficiently using automatic differentiation, the penalty schedule $\rho_k$ is gradually increased, and all steps are computed using mini-batches for scalability [2509.22483].

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 [2509.22483]. The paper states a **penalty stationarity guarantee**: as $\rho \to \infty$, every accumulation point of minimizers of $F(\theta)$ satisfies $\nabla_\theta \Phi(\theta) = 0$, meaning that the forgetting constraint is enforced [2509.22483].

For the inner maximization, the paper gives a convergence statement for convex, smooth $\Phi(\theta)$. After $T$ steps of gradient ascent,
$$
\Phi(\theta^*_{\text{in}}) - \Phi(\theta'^{(T)}) \leq \frac{\|\theta^*_{\text{in}} - \theta'^{(0)}\|^2}{2T\eta_{\text{in}}},
$$
with the proof leveraging standard convex optimization analysis [2509.22483].

For the full two-loop algorithm, the paper reports the following rates. In the convex case, after $K$ outer iterations with $T$ inner steps,
$$
F(\theta^{(K)}) - F^* \leq \mathcal{O}\left(\frac{1}{K}\right) + \mathcal{O}\left(\frac{K}{T^2}\right).
$$
Accordingly, with $K,T = \mathcal{O}(1/\epsilon)$, one obtains $\epsilon$-optimality [2509.22483]. In the non-convex case,
$$
\min_{k = 0, \dots, K-1} \mathbb{E}\|\nabla F(\theta^{(k)})\|^2
\leq \mathcal{O}\left(\frac{1}{K}\right) + \mathcal{O}\left(\frac{1}{T}\right) + \mathcal{O}(\sigma^2),
$$
where $\sigma^2$ is the stochastic gradient variance [2509.22483].

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 [2509.22483]. 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 [2509.22483]. 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) [2509.22483].

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 [2509.22483]. 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 [2509.22483]. On **WMDP**, the paper reports that OFMU attains the largest gains in accuracy in safety-critical domains and consistently outperforms competitive baselines [2509.22483].

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 [2509.22483]. Under random forgetting, the paper states that OFMU outperforms retrain in some scenarios [2509.22483]. Across these benchmarks, the framework is described as achieving the top normalized overall score among all considered metrics [2509.22483].

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 [2509.22483]. 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 [2509.22483].

## 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 [2509.22483]. 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 [2509.22483].

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 [2509.22483]. 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 [2509.22483]. 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 large language models [2509.22483].

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 [2509.22483]. 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 [2509.22483].

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 [2509.22483].

Source: https://www.emergentmind.com/topics/ofmu