💬 NLP

Model Unlearning Objectives Vary for Distinct Language Functions

This paper argues that language model unlearning should be tailored to specific functions rather than treated as a monolithic task, demonstrating through experiments on 7-8B models that distinct objectives—specifically a cosine-based meta-learned approach for dangerous knowledge and a multi-layer probe-based method for toxicity—yield superior results for their respective goals.

Original authors: Berk Atil, Vipul Gupta, Rebecca J. Passonneau

Published 2026-05-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Berk Atil, Vipul Gupta, Rebecca J. Passonneau

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 Large Language Models (LLMs) as incredibly talented, but slightly reckless, students who have read almost everything on the internet. They are great at writing stories and answering questions, but they've also picked up some bad habits: sometimes they know dangerous secrets (like how to build a bomb), and sometimes they use mean or toxic language.

The paper argues that trying to "unlearn" these bad habits is not a one-size-fits-all job. Just as a teacher uses different methods to stop a student from cheating on a math test versus stopping them from being rude to classmates, AI researchers need different tools for different problems.

Here is a breakdown of their approach using simple analogies:

1. The Core Problem: One Size Does Not Fit All

The authors say that "unlearning" (teaching the AI to forget specific things) shouldn't be treated as a single, generic task.

Dangerous Knowledge is like a student memorizing a specific, dangerous fact (e.g., "How to make poison"). This is stored in the model like a specific file in a library.
Toxicity is like a student's bad attitude or tendency to be rude. This isn't just one fact; it's a vibe or a pattern that is spread out across the whole student's personality.

The paper claims that because these two problems are stored differently inside the AI's "brain," you need two different strategies to fix them.

2. Strategy A: Forgetting Dangerous Facts (The "Cosine" Approach)

When the AI knows dangerous facts, the authors tried a new method based on direction rather than distance.

The Old Way (L2 Loss): Imagine trying to move a toy car away from a cliff. The old method measured the exact distance between the car and the cliff. If the car moved 1 inch, it counted as progress. But if the car was huge, moving 1 inch might not be enough.
The New Way (Cosine Loss): The authors suggest we shouldn't care about the exact distance. Instead, we should care about the direction the car is facing. If the car is facing the cliff, we turn it 180 degrees so it faces the opposite way. Even if it's still close to the edge, it's now pointing away from danger.
The Result: They also used a "smart coach" (meta-learning) that automatically figures out how hard to push the car to turn it without making it forget how to drive (general knowledge). This worked very well for removing dangerous facts.

3. Strategy B: Forgetting Toxicity (The "Multi-Layer" Approach)

When the AI is being toxic, the "turn the car around" method failed. Why? Because toxicity isn't stored in one specific spot; it's like a stain that has soaked into the fabric of the AI's brain across many different layers.

The Problem: If you only scrub one spot on a stained shirt, the stain is still there.
The Solution: The authors built a "multi-layer scrubber." They looked at the AI's brain at different depths (early, middle, and late layers) and found that the "toxicity signal" looks different at each level.
The Method: They trained special detectors (probes) at various layers to find exactly where the toxicity is hiding. Then, they pushed the AI to change its behavior at all those layers simultaneously, effectively scrubbing the stain from the whole shirt, not just one spot.

4. The Scorecard: S-Unlearning

How do you know if the AI is better? You can't just say, "It's less toxic," because maybe it also stopped being helpful.
The authors created a new score called S-Unlearning.

Imagine a balance scale. On one side is "How well did we remove the bad stuff?" and on the other is "How well did we keep the good stuff?"
The S-Unlearning score is like the area of a rectangle formed by these two sides. You want a big rectangle (high removal of bad stuff AND high retention of good stuff). If you remove the bad stuff but break the AI so it can't talk anymore, your rectangle shrinks to zero.

5. What They Found

They tested this on four different open-source AI models (like Llama, Mistral, and Qwen).

For Dangerous Facts: Their new "direction-based" method worked better than previous methods at making the AI forget dangerous info without losing its general smarts.
For Toxicity: Their "multi-layer scrubber" was much more effective than trying to use the same method used for dangerous facts. They found that toxicity is indeed spread out, and you have to target multiple layers to fix it.
The Big Takeaway: You can't use a single "magic wand" to fix all AI problems. You have to understand how the AI stores the problem (as a specific fact or as a scattered behavior) and choose the right tool to fix it.

In short: To fix a specific fact, change the direction. To fix a bad attitude, scrub the whole system layer by layer.

Technical Summary: Model Unlearning Objectives Vary for Distinct Language Functions

Problem Statement

Large Language Models (LLMs) acquire undesirable properties during pretraining, including dangerous factual knowledge (e.g., biosecurity risks) and toxic language generation. While post-training methods use distinct objectives to shape different behaviors (e.g., instruction following vs. safety), existing unlearning approaches often treat the removal of these distinct properties as a single, generic problem.

The authors argue that unlearning should be goal-dependent, analogous to the diversity of post-training objectives. Specifically, they posit that dangerous knowledge (factual/procedural information) and toxicity (socially harmful behavior) are represented differently within the model's internal mechanics. Consequently, a single unlearning objective may be ineffective for both. Current methods often fail to robustly remove these capabilities or inadvertently degrade general model utility.

Methodology

The paper proposes distinct unlearning strategies tailored to the specific mechanistic nature of the target property, evaluated across four open-source 7–8B models (Llama3.1-8B, Mistral-7B, Qwen2.5-7B, Olmo3-7B).

1. Dangerous Knowledge Unlearning

Based on mechanistic evidence that factual knowledge is stored in relatively stable locations (subject, object, and last-token positions) across layers, the authors modify the Representation Misdirection for Unlearning (RMU) framework:

Cosine-Based Loss: They replace the standard L2 distance loss with a cosine-based loss. The hypothesis is that representational direction is more critical than exact Euclidean position. This optimizes the alignment between the updated representation and the target direction without penalizing harmless magnitude changes.
Meta-Learned Coefficient ( $\alpha$ ): Instead of fixing the trade-off coefficient between the forget and retain objectives, the authors treat $\alpha$ as a learnable parameter. Using REINFORCE, the model adaptively updates $\alpha$ during fine-tuning to balance forgetting and retention, removing the need for extensive manual tuning.

2. Toxicity Unlearning

Mechanistic studies suggest toxicity is represented differently, often elicited in later MLP layers and not necessarily removed by alignment methods like DPO. A single-layer or single-direction approach proved ineffective.

Multi-Layer Objective: The authors propose a forget loss that targets toxicity-relevant signals across multiple layers.
Layer-Specific Probes: They train linear logistic regression probes on the last-token hidden representations of a frozen reference model at various layers to identify layer-specific toxicity directions ( $LR_i$ ).
Orthogonalization: The forget loss encourages the updated model's representations of toxic inputs to become orthogonal to the toxicity direction at each targeted layer. The loss squares the dot product to ensure non-negativity and symmetric gradients, giving larger gradients to representations strongly aligned with toxicity.
Layer Selection: Experiments indicate that effective unlearning requires coordinated updates across dispersed depths (early, middle, and late layers), as toxicity directions rotate significantly across the network depth.

3. Evaluation Metric: S-Unlearning

To summarize the trade-off between unlearning performance ( $U$ ) and retained utility ( $R$ ), the authors introduce S-unlearning:

Utility and unlearning scores are chance-corrected.
The final score is the product $\bar{U} \cdot \bar{R}$ .
This metric ranges from 0 to 1, maximizing when unlearning is strong and utility is high, and dropping to zero if utility falls to chance levels. It is geometrically equivalent to the area of the rectangle defined by the point $(\bar{U}, \bar{R})$ .

Key Results

The proposed methods were evaluated on four models using WMDP for dangerous knowledge and RealToxicityPrompts for toxicity.

Dangerous Knowledge: The cosine-based, meta-learned RMU variant achieved strong results, outperforming standard RMU and Adaptive RMU on most models (Mistral, Qwen, Olmo) while preserving general capabilities (MMLU). It achieved the best S-unlearning trade-off on three of the four models.
Toxicity: The multi-layer objective significantly outperformed baselines. While standard RMU and Adaptive RMU only reduced toxicity by 2–3%, the proposed method achieved reductions of 6–24%.
- The method yielded the highest S-unlearning scores for toxicity across all models.
- Layer Analysis: Probes revealed that toxicity directions are largely layer-specific (cosine similarity < 0.5 between most cross-layer pairs), validating the need for multi-layer intervention.
- Ablation: Updating three or more layers (specifically a mix of middle and late layers) yielded the best safety-utility trade-offs. Single-layer interventions often failed to reduce toxicity significantly.

Significance and Claims

The paper claims that unlearning should be studied as a family of problems rather than a single generic intervention.

Mechanistic Awareness: The results support the view that different language functions (factual recall vs. toxic generation) rely on distinct internal mechanisms. Therefore, unlearning objectives must be customized to the specific function being removed.
Geometry of Unlearning: For factual knowledge, preserving directional structure via cosine loss is superior to L2 distance. For toxicity, the distributed and rotating nature of the representation requires multi-layer, probe-guided orthogonalization.
Evaluation: The S-unlearning metric provides a unified way to capture the complex trade-offs inherent in unlearning, facilitating better comparison between methods.

The authors conclude that effective unlearning requires a deeper understanding of what is being forgotten, moving beyond one-size-fits-all approaches to methods that respect the specific geometric and structural properties of the target capability within the LLM.