💬 NLP

Forget What Matters, Keep the Rest: Selective Unlearning of Informative Tokens

This paper proposes Entropy-guided Token Weighting (ETW), a token-level unlearning regularizer that leverages predictive entropy to distinguish between informative and structural tokens, thereby enabling more effective unlearning of adversarial behaviors while better preserving the overall utility of large language models compared to existing methods.

Original authors: Seunghee Koh, Sunghyun Baek, Youngdong Kim, Junmo Kim

Published 2026-04-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Seunghee Koh, Sunghyun Baek, Youngdong Kim, Junmo Kim

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 AI model) who has read every book in the world. One day, you ask the librarian to "forget" a specific set of books—maybe because they contain sensitive information or because the author wants to retract them.

The challenge is tricky: You want the librarian to completely erase the knowledge from those specific books, but you don't want them to forget how to speak, how to write sentences, or how to answer questions about other books.

The Problem: The "Brute Force" Approach

Previously, when trying to make an AI "unlearn" something, researchers used a blunt instrument. They told the AI: "Every time you see a word from the book you need to forget, punish yourself!"

Think of this like trying to remove a stain from a white shirt by scrubbing the entire shirt with bleach.

The Stain: The specific fact you want to forget (e.g., "Carmen Montenegro was born in 1977").
The Shirt: The sentence structure (e.g., "The", "is", "of").

If you scrub the whole shirt, you remove the stain, but you also bleach the fabric until it's weak and holes appear. In AI terms, this "brute force" method damages the model's ability to speak naturally, making it sound robotic or nonsensical.

The Insight: Not All Words Are Created Equal

The authors of this paper noticed something clever. In any sentence, some words carry the heavy lifting of meaning, while others are just scaffolding.

Structural Words (The Scaffolding): Words like "the," "a," "is," or "and." These are predictable. If I say, "The cat sat on the...", you know the next word is almost certainly "mat" or "floor." There is very little surprise here.
Informative Words (The Heavy Lifting): Words like "Carmen," "Montenegro," or "Optometrist." These are unpredictable. If I say, "The author's name is...", the next word could be anything. There is a lot of surprise (or "uncertainty") here.

The authors realized: To forget a fact, you only need to target the "surprising" words, not the predictable ones.

The Solution: The "Entropy" Compass

The paper introduces a new method called Entropy-guided Token Weighting (ETW).

Imagine the AI is a detective trying to guess the next word in a sentence.

If the detective is 100% sure the next word is "the," the "surprise meter" (Entropy) is low.
If the detective is guessing wildly between "Carmen," "John," "Sarah," or "Bob," the "surprise meter" is high.

ETW works like this:

Scan the sentence: Look at every word the AI is about to generate.
Check the Surprise Meter:
- Low Surprise (Structural words): "Oh, it's just 'the' again." -> Ignore it. Don't punish the AI for knowing this. Keep the sentence flowing naturally.
- High Surprise (Informative words): "Wow, the AI is really unsure if the name is Carmen or someone else!" -> Focus here. This is where the specific knowledge lives. Apply a strong "forgetting" penalty to these specific words.

The Analogy: The Gold Miner

Think of the AI's memory as a river full of rocks.

Old Method: You dump a giant net over the whole river to catch the gold (the facts you want to forget). You catch the gold, but you also catch thousands of useless pebbles (structural words) and damage the riverbed (the model's utility).
ETW Method: You use a metal detector that only beeps when it finds gold. You only dig up the specific spots where the gold is hidden. The riverbed stays intact, and the water flows smoothly, but the gold is gone.

Why This Matters

By using this "Surprise Meter" (Entropy), the paper shows that:

Better Forgetting: The AI forgets the specific facts much more effectively because it targets the exact words that hold that information.
Better Utility: The AI keeps its ability to speak, write, and reason because it stops punishing the boring, structural words that hold the sentence together.
No Extra Tools Needed: Unlike previous methods that needed external dictionaries or grammar checkers to guess which words were important, this method lets the AI look at its own "uncertainty" to decide what to forget.

In short: Instead of scrubbing the whole shirt to remove a stain, ETW finds the exact spot of the stain and scrubs only that, leaving the rest of the shirt pristine.

1. Problem Statement

Machine unlearning in Large Language Models (LLMs) aims to remove specific knowledge (e.g., sensitive data or copyrighted content) while preserving the model's general utility. A critical challenge in this domain is selective unlearning: determining which tokens within a training sample should be penalized to forget the target knowledge without degrading the model's overall language capabilities.

The Limitation of Current Methods: Existing token-level unlearning approaches typically rely on ground-truth confidence (how confident the model is in the correct token) or external linguistic parsers (e.g., identifying nouns/entities via spaCy).
- Confidence-based methods fail because a token can have high confidence but still be semantically critical (informative), or low confidence but be structurally repetitive.
- Linguistic-based methods fail because they ignore context; for example, a noun might be a structural repetition of the prompt rather than the core answer.
The Consequence: Uniform or poorly targeted penalization leads to unnecessary degradation of model utility (the model becomes "dumber" generally) or ineffective forgetting (the model retains the forbidden knowledge).

2. Methodology: Entropy-guided Token Weighting (ETW)

The authors propose Entropy-guided Token Weighting (ETW), a token-level regularizer that uses the entropy of the model's predictive distribution as a proxy for token informativeness.

Core Intuition

Structural Tokens: Function words (e.g., "the", "a") or repetitive prompt mentions are highly predictable. The model assigns probability mass almost entirely to one token, resulting in low entropy.
Informative Tokens: Content words carrying specific knowledge (e.g., names, dates, specific facts) admit multiple plausible alternatives. The model distributes probability mass across several candidates, resulting in high entropy.

Mathematical Formulation

Given an autoregressive model $M_\theta$ and a completion sequence $y = (y_1, ..., y_n)$ :

Token-wise Entropy: For the $i$ -th token, the entropy $H(y_i)$ is calculated over the vocabulary $\mathcal{V}$ based on the predictive distribution $p(v | y_{<i}, x; \theta)$ :
$H(y_i) = -\sum_{v \in \mathcal{V}} p(v | y_{<i}, x; \theta) \log p(v | y_{<i}, x; \theta)$
Weighting Scheme: The forgetting loss is reweighted by normalizing these entropy values across the sequence length $n$ $n$ :
$\omega_i^{ETW} = \frac{n \cdot H(y_i)}{\sum_{j=1}^n H(y_j)}$
- High Entropy $\rightarrow$ High Weight: Informative tokens receive stronger penalties during the unlearning process (Gradient Ascent).
- Low Entropy $\rightarrow$ Low Weight: Structural tokens are preserved, minimizing utility degradation.
Implementation: The entropy is computed using a stop-gradient copy of the model parameters ( $\hat{\theta}$ ) to ensure the weighting calculation does not interfere with the gradient updates of the unlearning objective. A temperature parameter $T$ is applied to the softmax logits to control the sharpness of the distribution.

3. Key Contributions

Novel Regularizer: Introduction of ETW, which leverages the full predictive distribution (entropy) rather than just ground-truth confidence or external linguistic rules to identify informative tokens.
Theoretical Insight: Demonstration that entropy is a superior signal for distinguishing informative tokens from structural tokens, even when ground-truth confidence levels are similar.
Empirical Superiority: ETW achieves a better trade-off between Forget Quality (how well the model forgets) and Model Utility (how well the model retains general capabilities) compared to state-of-the-art baselines.

4. Experimental Results

The authors evaluated ETW on the TOFU (fictitious authors) and WMDP (hazardous knowledge) benchmarks using LLaMA-3.2 (1B, 3B) and LLaMA-2 (7B) models.

Token Discrimination (AUC):
- ETW achieved the highest ROC-AUC (0.72+) in distinguishing informative vs. structural tokens, significantly outperforming confidence-based methods (e.g., Imp, SatImp) and linguistic methods (SCN, SCE).
- Confidence-based methods often failed to distinguish tokens with similar confidence but different semantic roles.
Unlearning Performance (TOFU):
- Forget Quality: ETW achieved the lowest $-\log(FQ)$ scores (indicating the unlearned model's distribution is statistically indistinguishable from a retrained model) across 10%, 5%, and 1% forget splits.
- Model Utility: ETW preserved model utility ( $\Delta MU$ ) significantly better than Gradient Ascent (GA) and other token-weighted methods. For instance, in the 10% forget split for the 1B model, ETW had an aggregated score of 6.40, outperforming the next best (NPO at 23.10).
- Probability Alignment: The probability distribution of informative tokens in the ETW-unlearned model was closest to the retrained model (lowest $|\Delta RT|$ ), indicating that ETW successfully suppresses specific knowledge without distorting the model's general probabilistic behavior.
WMDP Benchmark: ETW showed robust performance in reducing accuracy on hazardous knowledge (Bio/Cyber) while maintaining competitive performance on general knowledge (MMLU), often outperforming NPO and SCE in the trade-off.

5. Significance and Impact

Context-Aware Unlearning: ETW moves beyond static rules (like "penalize all nouns") to dynamic, context-aware penalization. It understands that a noun like "author" in a prompt is structural, while "Carmen" in the answer is informative.
Efficiency: By focusing unlearning pressure only on high-entropy (informative) tokens, the method prevents the "collateral damage" to structural tokens that causes general language degradation.
Scalability: The method proved effective across different model sizes (1B to 7B) and varying forget ratios, suggesting it is a robust strategy for practical LLM unlearning.
Paradigm Shift: The paper challenges the reliance on ground-truth confidence for unlearning, proposing that the uncertainty (entropy) of the model's prediction is a more reliable indicator of semantic importance.

In conclusion, ETW offers a principled, entropy-driven approach to machine unlearning that effectively targets the "what matters" (informative content) while preserving the "rest" (structural integrity), setting a new standard for selective knowledge removal in LLMs.