💬 NLP

Learning What to Forget: Improving LLM Unlearning via Learned Token-Level Importance

This paper introduces Alternating Token-Weighted Unlearning (ATWU), a lightweight framework that improves machine unlearning in large language models by jointly learning token-level importance through the conflict between forget and retain objectives, thereby achieving state-of-the-art performance without requiring external supervision or auxiliary models.

Original authors: Gizem Yüce, Giorgos Nikolaou, Nicolas Flammarion

Published 2026-06-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Gizem Yüce, Giorgos Nikolaou, Nicolas Flammarion

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" That Smudges the Whole Page

Imagine you have a giant, incredibly smart library (a Large Language Model) that has read almost everything on the internet. One day, you realize it memorized a specific secret diary entry that belongs to someone else, and you need to remove it immediately.

The problem is that the library is so huge that you can't just tear out the specific page where the diary was written. If you try to "unlearn" that specific story by retraining the whole library without it, you might accidentally erase other useful things too, like how to write a polite email or solve a math problem.

Current methods of "unlearning" are like using a giant, blunt eraser. They try to wipe out the entire sentence or paragraph containing the secret. But often, they wipe out the grammar, the punctuation, and the common words (like "the" or "and") along with the secret. This leaves the library's writing style broken and clumsy.

The Core Idea: A Smart, Surgical Scalpel

The authors of this paper propose a new way to think about forgetting. They realized that not every word in a secret sentence is equally important to the secret.

The Secret Part: The specific name, date, or fact you want to remove (e.g., "Hina Ameen").
The Structural Part: The boring, necessary words that hold the sentence together (e.g., "The author is...").

If you try to make the model forget the word "The," you aren't removing the secret; you're just breaking the grammar. The model needs to keep knowing how to use "The."

The Solution: ATWU (Alternating Token-Weighted Unlearning)

The paper introduces a method called ATWU. Think of this as giving the library a smart, surgical scalpel instead of a blunt eraser.

Here is how it works, step-by-step:

The Conflict Test: The authors realized that the best way to figure out which words to erase is to see how much "pain" it causes the library to forget them.
- If the library tries to forget the word "Hina" (the secret), it doesn't hurt its ability to write good sentences.
- If the library tries to forget the word "The," it starts writing gibberish.
- The Insight: If forgetting a word doesn't hurt the library's general skills, that word is a good candidate for removal. If it hurts the skills, keep it.
The "Scorekeeper" (The Linear Scorer):
- Instead of hiring a human to read every sentence and circle the words to delete (which is slow and expensive), the authors built a tiny, simple "scorekeeper" inside the model.
- This scorekeeper looks at the "brain" of the model (the hidden states) and assigns a score to every single word.
- High Score: "This word is the secret; we must erase it."
- Low Score: "This word is just grammar; leave it alone."
The Dance (Alternating Optimization):
- The model and the scorekeeper take turns learning.
- First, the scorekeeper decides which words are important to forget.
- Then, the model tries to forget those specific words while keeping the rest of its skills intact.
- Then, the scorekeeper looks at how the model changed and updates its scores to be even more accurate.
- They keep dancing back and forth until the model learns exactly what to forget without breaking anything else.

Why This is Better (The Results)

The paper tested this method on two different "libraries" (datasets) and found that ATWU is much better than previous methods:

Precision: It successfully removed the specific secrets (like the name "Hina Ameen") while keeping the model's ability to write normal, high-quality text.
No Extra Help Needed: Unlike other methods that need a second, separate AI to help figure out what to delete, ATWU figures it out all by itself using the model's own internal signals.
Better Balance: It achieves a better trade-off: it forgets the bad stuff more effectively while keeping the good stuff better than any other method tested.

The "Magic" of the Metaphor

Imagine you are trying to remove a specific stain from a white shirt.

Old Methods: You scrub the whole shirt with bleach. The stain is gone, but the shirt is now yellow and ruined.
ATWU: You use a laser that only targets the color of the stain. It burns away the stain perfectly, leaving the white fabric of the shirt completely untouched.

Summary of Claims

The paper claims that by treating "forgetting" as a joint problem where the model learns what to forget and how to forget it simultaneously, we can achieve "machine unlearning" that is:

Unsupervised: It doesn't need humans to label which words to delete.
Efficient: It uses a very simple, lightweight mechanism (a linear scorer) to do the work.
Effective: It creates a much cleaner separation between the "secret" information and the "general knowledge" the model needs to keep.

The authors conclude that this approach proves we don't need external tools or expensive annotations to teach a model what to forget; the model's own internal conflict between "remembering the secret" and "keeping its skills" provides all the clues it needs.

Technical Summary: Learning What to Forget: Improving LLM Unlearning via Learned Token-Level Importance

Problem Statement
Machine unlearning aims to remove targeted knowledge (e.g., personal data, copyrighted text, harmful content) from trained Large Language Models (LLMs) without retraining from scratch. A critical challenge in autoregressive LLMs is that not all tokens within a "forget sample" (the data to be removed) are equally relevant to the targeted information. Many tokens are structural (e.g., function words, punctuation, syntactic patterns) and are necessary for general language capabilities. Existing token-weighted unlearning approaches attempt to selectively apply the forget loss to specific tokens but rely on external annotations, auxiliary models, probability-based heuristics, or task-specific rules to estimate token relevance. These methods often fail to distinguish effectively between forget-specific content and structural tokens, leading to either insufficient forgetting or unnecessary degradation of general utility.

Methodology: Alternating Token-Weighted Unlearning (ATWU)
The authors propose that token-level forget-specificity should be characterized by the interaction between the forget and retain objectives themselves, rather than external signals. They formalize this as a joint optimization problem over model parameters ( $\theta$ ) and token weights ( $z$ ).

Theoretical Formulation: The paper defines a "retain conflict" metric, $\kappa(A)$ , which measures the cost of suppressing a set of tokens $A$ while maintaining optimality on the retain set. Under a "retain-conflict separation" assumption, the authors prove that the oracle set of forget-specific tokens can be recovered as the minimizer of a joint objective that balances the retain loss and a weighted forget loss. This theoretical result suggests that token weights can be learned directly from the unlearning objective without external supervision.
Algorithm (ATWU): To make this tractable, the authors introduce Alternating Token-Weighted Unlearning (ATWU).
- Scorer: Instead of learning independent weights for every token, ATWU uses a lightweight linear scorer, $g_w$ , parameterized by a vector $w$ in the model's hidden-state space. The score for a token is $\hat{z}_{x,t} = \sigma(\langle w, h_\theta(x_t) \rangle)$ , where $h_\theta$ is the hidden representation.
- Joint Optimization: The framework optimizes a relaxed, continuous objective that includes the retain loss, a saturated negative cross-entropy forget loss (SATGA+), an entropy regularizer to push scores toward binary values, and a budget constraint to control the number of tokens selected.
- Alternating Updates: The model parameters $\theta$ and the scorer parameters $w$ are updated alternately. When updating $\theta$ , scores are detached; when updating $w$ , $\theta$ is frozen. This stabilizes the training by decoupling the feedback loop between the changing model geometry and the token selection.

Key Contributions

Retain-Conflict Characterization: The paper provides a theoretical characterization of token-level forget-specificity based on the tension between forget and retain objectives, proving that a joint optimization can recover the oracle forget-specific support under natural separation conditions.
Unsupervised Learning of Token Weights: ATWU is a lightweight framework that learns token forget-specificity jointly with model parameters using only the model's internal representations. It requires no external token-level supervision, auxiliary models, or expensive annotations.
Adaptive Loss Integration: The method introduces SATGA+, a score-modulated variant of saturated negative cross-entropy where the learned token score is injected into the saturation exponent. This allows for smoother, uncertainty-aware updates during the early phases of optimization.

Results
The authors evaluate ATWU on two benchmarks: TOFU (synthetic fictitious author biographies) and RWKU (real-world public figure knowledge).

Performance: ATWU achieves state-of-the-art Unlearning Quality (UQ), a metric balancing forget quality (FQ) and retain degradation (RD). On TOFU, it achieves a UQ of 91.7 (exceeding the runner-up by 5.8 pp), and on RWKU, 58.7 (exceeding the runner-up by 6.8 pp).
Utility Preservation: Unlike aggressive unlearning methods that collapse generation quality, ATWU preserves general utility (measured by MMLU, repetitiveness, and pairwise win rates) close to the original checkpoint.
Token Score Quality: The learned token scores align significantly better with ground-truth forget-specific spans (measured by AUROC) compared to probability-based heuristics, auxiliary-model approaches, and other token-weighting baselines.
Versatility: The framework is loss-agnostic; when combined with other forget losses (DPO, NPO, SIMNPO), ATWU consistently improves the forget–retain trade-off over their unweighted counterparts.

Significance and Claims
The paper claims that the "retain conflict" provides an effective, intrinsic criterion for identifying what language models should forget. By shifting token-weighted unlearning from a pipeline of external heuristics to an end-to-end joint optimization problem, ATWU demonstrates that:

The unlearning objective itself contains sufficient latent signal to recover forget-specific tokens without external supervision.
The bottleneck in selective unlearning is not merely finding the right tokens, but the principled co-design of the scoring mechanism and the forget objective.
Scalable, privacy-preserving, and highly targeted unlearning is achievable without relying on external annotations, which often introduce privacy risks and scalability issues.

The authors note limitations, including the reliance on a hyperparameter for the token-selection budget and the assumption of a strict separation between forget-specific and structural tokens, which may only hold approximately in natural language. Additionally, evaluations are currently limited to specific benchmarks and single runs due to computational costs.