💻 computer science

DualOptim+: Bridging Shared and Decoupled Optimizer States for Better Machine Unlearning in Large Language Models

The paper proposes DualOptim+, a novel optimization framework that bridges shared and decoupled optimizer states to achieve superior trade-offs in machine unlearning for large language models, alongside a memory-efficient quantized variant called DualOptim+ 8bit.

Original authors: Xuyang Zhong, Qizhang Li, Yiwen Guo, Chen Liu

Published 2026-05-22

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xuyang Zhong, Qizhang Li, Yiwen Guo, Chen Liu

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

The Big Problem: The "Eraser" Dilemma

Imagine you have a giant, incredibly smart encyclopedia (a Large Language Model) that has read almost everything on the internet. Sometimes, you need to remove specific pages from it—maybe because they contain private information, outdated facts, or harmful instructions. This process is called Machine Unlearning.

The problem is tricky: You want to erase the bad or private pages without accidentally deleting the rest of the encyclopedia's knowledge. If you try to scrub out the bad info too hard, the model might start forgetting how to write a poem, solve a math problem, or even say "hello."

The Old Ways: Two Flawed Approaches

Before this paper, researchers tried two main ways to fix this:

The "Mix-and-Mash" Approach (Joint Updating): Imagine trying to clean a muddy shoe while simultaneously polishing a diamond, using the same cloth for both. You mix the "cleaning" and "polishing" instructions together. The result? The shoe stays dirty, or the diamond gets scratched. The model forgets too much useful info.
The "Switch-Back" Approach (Alternate Updating): Imagine cleaning the shoe for one minute, then polishing the diamond for the next, switching back and forth. This is better, but if you use the same cloth for both tasks, the dirt from the shoe eventually gets on the diamond, and the polish from the diamond makes the shoe slippery. The "memory" of the cloth gets confused.

The New Solution: DualOptim+

The authors propose a new framework called DualOptim+. Think of it as giving the model two specialized cleaning kits that work together perfectly.

1. The "Shared Base" (The Common Ground)

Imagine the model has a Shared Base State. This is like a common foundation or a "general knowledge" bucket.

When the model learns to forget something and when it learns to keep something, they both update this bucket.
Why? Because forgetting and remembering often share some basic brain structures. This bucket captures what the model needs to keep doing regardless of the specific task.

2. The "Delta States" (The Specialized Tools)

Then, there are Delta States (or "Delta" for short). Think of these as specialized, separate toolkits for the specific job.

Delta-F: A toolkit just for "Forgetting." It holds the difference between what needs to be erased and the common foundation.
Delta-R: A toolkit just for "Retaining." It holds the difference between what needs to be saved and the common foundation.

How It Works Together

Instead of using one confused cloth or two completely separate cloths that never talk to each other, DualOptim+ does this:

It updates the Shared Base with both the forgetting and remembering instructions.
It updates the Delta tools with only the unique parts of the instructions (the "residuals").
When the model actually changes its brain (updates its parameters), it combines the Shared Base + the Right Delta Tool.

The Magic: If the "forget" and "remember" instructions are fighting each other (conflicting), the Delta tools step in to handle the conflict. If they agree, the Shared Base handles the work efficiently. It's like a smart manager who knows when to use a shared team resource and when to assign a specialist.

The "8-bit" Trick: Saving Space

Usually, having these extra toolkits (Delta states) takes up a lot of computer memory (RAM), which is expensive.

The authors created a DualOptim+ 8-bit version.
Analogy: Imagine your toolkits are usually made of heavy, solid gold (32-bit precision). The 8-bit version compresses them into lightweight, high-tech foam. They weigh much less (saving memory) but still do the exact same job with almost no loss in performance.

What Did They Prove?

The paper ran many tests to see if this new method works:

Fake Scenarios: They tested it on made-up data (like fictional authors) to see if the model could forget specific names without losing its ability to write.
Real Scenarios: They tested it on real-world data, like removing specific people's private info from a model.
Safety: They tested if it could remove "harmful" instructions (like how to make a bomb) without making the model refuse to answer any questions (a common side effect called "over-refusal").
Multi-Tasking: They even tested it on teaching the model three different skills at once (coding, science, and math) to see if the "Shared Base" helps balance different tasks.

The Result: In almost every test, DualOptim+ was better at erasing the bad stuff while keeping the good stuff intact compared to the old methods. It found the "sweet spot" between forgetting and remembering.

Summary

DualOptim+ is a smarter way to train AI models to "unlearn" things. Instead of using a confused mix or two completely separate systems, it uses a hybrid system: a shared foundation for common sense and specialized tools for specific tasks. This ensures the AI forgets exactly what it's supposed to, without losing its general smarts. Plus, they figured out how to make it run on cheaper computers by compressing the memory.

Technical Summary of DualOptim+

Problem Statement
Machine Unlearning (MU) in Large Language Models (LLMs) aims to erase the influence of specific training data (the forget set) while preserving the model's general utility on the remaining data (the retain set). Existing approaches face a significant challenge in balancing these conflicting objectives.

Joint Updating: Most methods minimize the sum of forgetting and retaining losses simultaneously. This often leads to significant degradation in model utility due to excessive forgetting.
Alternate Updating: Alternating between forgetting and retaining objectives improves performance but suffers from instability and sensitivity to hyperparameters.
Decoupled Optimizers (DualOptim): A previous approach, DualOptim, decouples optimizer states (e.g., momentum terms) for each objective to prevent gradient entanglement. While effective in computer vision, DualOptim yields only marginal improvements in LLMs. The authors observe that in LLM unlearning, the gradients for forgetting and retaining are not always in direct conflict; they often share common representations, especially in later training phases. Fully decoupled states fail to leverage these shared features, while fully shared states fail to handle conflicts.

Methodology: DualOptim+
The paper proposes DualOptim+, a plug-and-play optimization framework designed to bridge the gap between fully shared and fully decoupled optimizer states. The core innovation is the decomposition of the optimizer state into two components:

Shared Base State ( $B$ ): This state captures the common representations shared by both the forgetting and retaining objectives. It is updated jointly using gradients from both objectives ( $\nabla_\theta L_f$ and $\nabla_\theta L_r$ ).
Decoupled Delta States ( $\Delta_f, \Delta_r$ ): These states preserve the objective-specific residuals. They are updated based on the difference between the specific objective's gradient and the bias-corrected base state.
- $\Delta_f \leftarrow \beta \Delta_f + (1-\beta)(\nabla_\theta L_f - \hat{B})$
- $\Delta_r \leftarrow \beta \Delta_r + (1-\beta)(\nabla_\theta L_r - \hat{B})$

Parameter Update:
The final momentum term used for parameter updates is reconstructed by combining the bias-corrected base state with the respective bias-corrected delta state ( $\hat{B} + \hat{\Delta}_f$ or $\hat{B} + \hat{\Delta}_r$ ).

Key Theoretical Insight:
Theoretical and numerical analyses demonstrate that DualOptim+ functions as an adaptive intermediate.

When gradients are positively correlated (shared features dominate), the delta states converge to zero, and the method behaves like an Alternate scheme (utilizing the shared base).
When gradients are negatively correlated (conflicting objectives), the base state converges to zero, and the method behaves like a DualOptim scheme (utilizing decoupled states).
This adaptivity allows the optimizer to dynamically adjust based on the directional conflict between forgetting and retaining gradients.

DualOptim+ 8bit
To address the memory overhead associated with storing additional optimizer states (3x the memory of a standard Adam optimizer), the authors introduce DualOptim+ 8bit. This variant quantizes the optimizer states to 8-bit precision, significantly reducing memory consumption (to roughly 3/4 of a standard Adam) without compromising performance.

Key Contributions

Novel Framework: The proposal of DualOptim+, which introduces a shared base state and decoupled delta states to effectively bridge shared and decoupled optimizer states.
Plug-and-Play Design: A framework compatible with any optimizer that stores states (e.g., AdamW, Muon) and applicable to multi-objective optimization.
Empirical Validation: Extensive experiments across fictitious datasets (TOFU), real-world scenarios (Llama-3 unlearning), safety alignment, and multi-task learning tasks.
Efficiency: The introduction of the 8-bit quantized variant that mitigates memory overhead.

Experimental Results
The paper evaluates DualOptim+ against Joint, Alternate, and DualOptim baselines across various LLMs (Phi-1.5, Llama-2, Llama-3) and tasks:

Fictitious Scenarios (TOFU): DualOptim+ consistently achieves the best Overall Performance (OVR), defined as a trade-off between Forget Efficacy (FE) and Model Utility (MU). It outperforms DualOptim, particularly in targeted unlearning, by better leveraging shared features.
Real-World Scenarios: In unlearning real-world individuals from Llama-3, DualOptim+ maintains superior utility on downstream tasks (ARC-c, MMLU, etc.) compared to baselines, whereas DualOptim sometimes underperforms the simpler Alternate approach.
Safety Alignment: In a task to unlearn unsafe knowledge while retaining instruction-following capabilities, DualOptim+ achieves the highest safety-utility trade-off and the lowest over-refusal rate on the XSTest benchmark.
Multi-Task Learning: Unlike in unlearning where gradient conflicts are high, in multi-task learning where conflicts are less pronounced, DualOptim+ continues to deliver substantial performance boosts, whereas DualOptim shows negative gains.
Memory Efficiency: The 8-bit variants achieve near-parity in memory usage with standard baselines while maintaining the performance benefits of the full-precision method.

Significance and Claims
The authors claim that DualOptim+ provides a more stable and effective trade-off between knowledge erasure and utility preservation than existing methods. They position the method not just as a unlearning tool, but as a generalizable optimization framework applicable to broader scenarios involving multiple conflicting learning objectives, such as LLM alignment and multi-objective learning. The paper emphasizes that the adaptive nature of bridging shared and decoupled states is crucial for handling the dynamic gradient relationships found in LLM unlearning.