💬 NLP

Modeling LLM Unlearning as an Asymmetric Two-Task Learning Problem

This paper reframes LLM unlearning as an asymmetric two-task problem prioritizing retention, proposing a gradient synthesis framework that includes a novel SAGO method to theoretically and empirically achieve superior performance on the forgetting-retention trade-off by optimizing gradient geometry rather than loss re-balancing.

Original authors: Zeguan Xiao, Siqing Li, Yong Wang, Xuetao Wei, Jian Yang, Yun Chen, Guanhua Chen

Published 2026-04-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zeguan Xiao, Siqing Li, Yong Wang, Xuetao Wei, Jian Yang, Yun Chen, Guanhua Chen

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 brilliant, all-knowing librarian (the Large Language Model or LLM). This librarian has read almost everything on the internet. They are incredibly helpful, but they also remember some very dangerous secrets (like how to build a bomb) and some private gossip about specific people.

You want the librarian to forget these specific bad or private things, but you desperately want them to keep their general wisdom, grammar, and ability to help with everything else.

The Problem: The "Over-Correction" Trap

Traditionally, trying to make the librarian forget something is like trying to erase a specific sentence from a book by hitting the whole page with a giant eraser.

The Old Way (Gradient Ascent): You tell the librarian, "Forget this!" and they try so hard to un-learn it that they accidentally un-learn how to speak English, how to do math, or how to be polite. They become confused and useless.
The "Balanced" Way: Some researchers tried to tell the librarian, "Forget the bad stuff, but also remember the good stuff," treating both goals as equally important. But this often leads to a tug-of-war. The librarian gets stuck in the middle, doing neither job well.

The New Idea: The "Guardian First" Approach

This paper proposes a new way of thinking. Instead of treating "Forgetting" and "Remembering" as equal partners, they are asymmetric.

Remembering (Retention) is the Boss. The librarian's main job is to stay smart and helpful.
Forgetting is the Intern. The intern's job is to remove the bad stuff, but only if it doesn't annoy the Boss.

The authors call this an Asymmetric Two-Task Problem. The goal isn't to find a compromise; it's to protect the "Boss" (general knowledge) at all costs while quietly removing the "Intern's" target (bad knowledge).

The Solution: Two New Tools

To make this work, the authors created a framework that acts like a traffic controller for the librarian's brain updates. Every time the librarian learns something new, two signals arrive at once:

The "Forget" Signal: "Delete this!"
The "Remember" Signal: "Keep this!"

Sometimes these signals point in opposite directions (a traffic jam). The paper introduces two methods to fix this:

1. PCGrad (The "Sideways Step")

Imagine the "Remember" signal is a strong wind blowing North. The "Forget" signal is a wind blowing South. If you just add them, you go nowhere.

PCGrad says: "Okay, we can't go South because that hurts the North wind. But we can't go North either because we need to forget. So, let's take a sideways step (East or West)."
It projects the "Forget" signal onto a path that doesn't fight the "Remember" signal. It's a clever way to avoid a crash, but it's still a bit of a compromise.

2. SAGO (The "Strict Gatekeeper")

This is the paper's star invention. SAGO is even smarter and stricter.

Imagine the librarian's brain is made of millions of tiny switches.
SAGO looks at each switch individually.
- If the "Forget" switch and the "Remember" switch are trying to flip the same way? Great! Let the "Forget" signal go through.
- If the "Forget" switch is trying to flip the opposite way of the "Remember" switch? STOP! SAGO blocks the "Forget" signal for that specific switch and lets the "Remember" signal win.
The Metaphor: Think of SAGO as a bouncer at a club. The "Remember" crowd is VIPs. The "Forget" crowd wants to get in. SAGO checks every single person. If a "Forget" person is trying to push a VIP out of the way, they get kicked out immediately. If a "Forget" person is just walking in a different aisle where no VIPs are, they can pass.
The Result: The librarian never takes a single step that hurts their general knowledge. They only remove the bad stuff where it's safe to do so.

The Results: Why It Matters

The authors tested this on real-world benchmarks (like removing biosecurity risks or private info).

Before: When they tried to remove the bad info, the librarian's general intelligence dropped by half (e.g., from 100% smart to 45% smart).
With SAGO: They removed the bad info, but the librarian stayed 96% smart.

The Big Takeaway

The paper proves that the secret to fixing AI isn't just about balancing the math equations better. It's about changing the geometry of the learning process.

Instead of trying to balance two equal forces, you should build a system where protecting the model's core intelligence is the absolute rule, and forgetting is only allowed when it doesn't break that rule. It's like saying, "We will clean the house, but we will never break the furniture while doing it."

In short: SAGO is the ultimate "Do No Harm" filter for AI unlearning, ensuring the AI stays smart even after it's been scrubbed of its secrets.

1. Problem Statement

Large Language Models (LLMs) trained on internet-scale data often memorize sensitive personal information or hazardous knowledge (e.g., biosecurity threats). Machine Unlearning (MU) aims to remove this specific information while preserving the model's general capabilities.

However, a central challenge exists: the trade-off between unlearning and retention.

The Conflict: Standard unlearning methods (like Gradient Ascent on the forget set) often cause "over-forgetting," degrading the model's general utility.
The Limitation of Current Approaches: Existing methods that combine forgetting and retaining objectives (e.g., GradDiff, NPO+GD) often treat the two tasks symmetrically or simply sum their gradients. This leads to destructive interference where the "forget" gradient actively opposes the "retain" gradient, causing performance collapse.
The Core Insight: The authors argue that unlearning is not a balanced multi-task problem. Instead, it is an asymmetric two-task problem where Retention is the primary objective and Forgetting is an auxiliary objective constrained by a "do-no-harm" principle.

2. Methodology

The paper proposes a Retention-Prioritized Gradient Synthesis Framework (Algorithm 1) that decouples gradient extraction from conflict-aware combination. The framework operates iteratively:

Extract gradients for the forget set ( $g_f$ ) and retain set ( $g_r$ ).
Synthesize a final update direction ( $g_{final}$ ) that injects forgetting only where it does not conflict with retention.

The authors instantiate this framework with two specific methods:

A. Adapted PCGrad (Project Conflicting Gradients)

Mechanism: Based on the multi-task learning technique PCGrad, this method projects the forget gradient onto the normal plane of the retain gradient when they conflict (i.e., dot product $< 0$ ).
Innovation: Unlike standard PCGrad which flattens all parameters, this paper applies module-wise projection. It detects conflicts locally within specific model modules and projects only those specific gradients. This prevents unnecessary corrections in non-conflicting modules, preserving finer-grained retention.

B. SAGO (Sign-Align Gradient Optimization)

Core Idea: SAGO introduces a novel, element-wise (per-parameter) gating mechanism to ensure the update direction never opposes the retain signal.
Mechanism:
- It analyzes the signs of $g_f$ and $g_r$ for every parameter.
- Aligned Signs: If $g_f$ and $g_r$ have the same sign, the forgetting signal is allowed to pass (task-specific).
- Opposite Signs: If they have opposite signs (conflict), the forgetting signal is zeroed out, and only the retain gradient is used.
- Formula: $\tilde{g}_f = g_f \odot I(g_f \odot g_r \geq 0)$ and $\tilde{g}_r = g_r \odot I(g_f \odot g_r < 0)$ . The final gradient is a weighted sum of these gated components.
Theoretical Guarantee: SAGO ensures that the final update direction has non-negative cosine similarity with the retain gradient. Furthermore, it achieves strictly tighter alignment than PCGrad because it completely eliminates any coordinate where the update direction opposes the retain signal, whereas PCGrad's projection can still result in antagonistic components in high-dimensional spaces.

3. Key Contributions

Asymmetric Formulation: Reframing LLM unlearning as an asymmetric problem (Retention = Primary, Forgetting = Auxiliary) rather than a balanced multi-task learning problem.
Novel Gradient Synthesis:
- Adapting PCGrad with module-wise projection for unlearning.
- Proposing SAGO, a sign-constrained synthesis method that guarantees non-antagonistic updates.
Theoretical Analysis: Proving that both methods ensure non-negative cosine similarity with the retain gradient, with SAGO providing superior directional fidelity.
Empirical Validation: Demonstrating that reshaping gradient geometry is more effective than simply re-balancing loss weights.

4. Experimental Results

The methods were evaluated on two major benchmarks: WMDP (Biosecurity and Cybersecurity) and RWKU (Real-World Knowledge Unlearning).

Performance Metrics:
- Forgetting: Lower accuracy on forget sets (WMDP) or lower ROUGE-L (RWKU) is better.
- Retention: Higher MMLU scores (WMDP) or higher ROUGE-L on neighbor sets (RWKU) is better.
Key Findings:
- Pareto Frontier Shift: SAGO consistently pushes the Pareto frontier outward, achieving significantly higher retention for the same level of forgetting compared to baselines (Naive, GradDiff, PCGrad).
- Quantitative Gains (WMDP Bio):
  - Using SimNPO+GD as a base, naive unlearning recovered only 44.6% of the target model's MMLU performance.
  - Adding PCGrad improved this to 94.0%.
  - Adding SAGO further improved it to 96.0%, while maintaining comparable forgetting strength.
- RWKU Results: SAGO showed massive gains in neighbor retention (e.g., +13.3 points over PCGrad in GA+GD settings), indicating superior preservation of general knowledge.
- Gradient Geometry: Analysis of cosine similarities showed that SAGO minimizes the conflict between forget and retain gradients more effectively than PCGrad or GradDiff, maintaining the highest alignment with the retain gradient.

5. Significance

This paper fundamentally shifts the paradigm of LLM unlearning from "balancing losses" to "managing gradient geometry."

Safety Implications: By prioritizing retention, SAGO allows for the removal of dangerous knowledge (e.g., how to build biological weapons) without crippling the model's ability to perform general tasks. This is crucial for deploying safe LLMs in real-world scenarios.
Methodological Impact: The proposed framework is modular and compatible with various existing unlearning objectives (GA, NPO, SimNPO), making it a versatile drop-in solution for improving unlearning efficacy.
Theoretical Insight: The work demonstrates that treating unlearning as an asymmetric problem is theoretically sound and empirically superior to symmetric multi-task approaches, offering a new direction for future research in model editing and safety.