Unlearning What Matters: Token-Level Attribution for Precise Language Model Unlearning

1. Problem Statement

Large Language Models (LLMs) often memorize sensitive, copyrighted, or harmful data, necessitating machine unlearning to remove specific knowledge without retraining from scratch.

• Current Limitation: Existing unlearning methods operate at the sequence level, applying uniform gradient updates across all tokens in a sequence.
• The Core Issue: Knowledge is not encoded uniformly; only a sparse subset of tokens (e.g., specific entities, dates, facts) carries the targeted information, while others serve structural or contextual roles.
• Consequences of Uniform Updates:
  1. Gradient Noise: Updating non-critical tokens introduces noise, degrading the signal-to-noise ratio (SNR) of the unlearning process.
  2. Utility Degradation: Indiscriminate updates damage the model's performance on unrelated knowledge (collateral damage).
  3. Suboptimal Forgetting: The model fails to effectively remove the targeted knowledge because the gradient signal is diluted by irrelevant tokens.

2. Methodology: TokenUnlearn

The authors propose TokenUnlearn, a framework that shifts unlearning from the sequence level to the token level by identifying and selectively targeting "knowledge-critical" tokens.

A. Token Attribution Mechanism

To identify which tokens encode the knowledge to be forgotten, the framework computes a composite importance score ($\phi_i$) for each token using two signals:

1. Knowledge-Aware Attribution (Masking):
   • Mechanism: The model generates predictions on the original input ($z_{orig}$) and on an input where knowledge-relevant portions (specifically nouns in the prompt) are masked ($z_{mask}$).
   • Metric: The absolute log-probability shift ($\Delta^{unlearn}_i$) is calculated. A large shift indicates the token's prediction is heavily dependent on the targeted knowledge.
2. Entropy-Aware Complementary Signal:
   • Mechanism: Calculates the predictive entropy ($H_i$) of the token.
   • Rationale: High entropy may indicate decision points where the model is uncertain about knowledge-dependent continuations (e.g., choosing between multiple plausible facts).
   • Combination: The final score is a weighted sum: $\phi_i = \alpha \cdot \bar{\Delta}^{unlearn}_i + (1-\alpha) \cdot \bar{H}_i$.

B. Token-Selective Optimization Strategies

Once importance scores are generated, the framework applies unlearning objectives only to high-importance tokens using two strategies:

1. Hard Selection:
   • Selects the top-$r$ fraction of tokens (e.g., top 20%) based on $\phi_i$.
   • Applies unlearning loss exclusively to these tokens; other tokens receive zero gradient updates.
2. Soft Weighting:
   • Modulates the gradient contribution of every token based on a normalized importance score (using a softmax with temperature $\tau$).
   • Allows for smooth interpolation between uniform updates and fully selective updates.

C. Integration with Existing Methods

The framework is model-agnostic and integrates with standard gradient-based unlearning objectives, including:

• Gradient Ascent (GA)
• Weighted Gradient Ascent (WGA)
• Negative Preference Optimization (NPO)
• Representation Misdirection Unlearning (RMU)
• Retention: All methods are combined with KL-divergence regularization on the "retain set" to preserve general capabilities.

3. Theoretical Analysis

The paper provides a theoretical justification for why token-level selection improves unlearning:

• Signal-to-Noise Ratio (SNR): The authors model unlearning as movement along a low-dimensional subspace ($U$) of parameters.
• Noise Reduction: Uniform updates accumulate noise from all $T$ tokens. Token-level selection concentrates the gradient weight on the subset $K$ of knowledge-critical tokens.
• Result: The SNR of the gradient estimator improves proportionally to the inverse of the selection ratio ($T/|K|$). By suppressing updates on non-critical tokens, the method reduces gradient noise and preserves unrelated knowledge more effectively.

4. Experimental Results

The authors evaluated TokenUnlearn on two benchmarks across three model architectures (Phi-1.5, Llama-2-7B, Qwen-3-8B).

Benchmarks

1. TOFU (Task of Fictitious Unlearning): Fine-grained knowledge unlearning (removing facts about fictitious authors).
2. WMDP (Weapons of Mass Destruction Proxy): Safety alignment (removing hazardous bio/cybersecurity knowledge).

Key Findings

• Superior Forgetting: Token-level variants (T-* and S-*) consistently outperformed sequence-level baselines in forgetting effectiveness (lower accuracy on forget sets).
  • Example: On TOFU with Qwen-3-8B, T-WGA improved forgetting effectiveness significantly compared to standard WGA.
• Better Utility Preservation: Token-level methods maintained higher model utility (retain scores) on non-targeted data.
  • Example: On Qwen-3-8B, the retain score improved by up to 19.0% compared to sequence-level baselines while achieving better forgetting.
• Generalizability: The approach worked consistently across different model sizes and architectures.
• Ablation Studies:
  • Selection Ratio ($r$): A ratio of 20% (top 20% of tokens) provided the optimal trade-off between forgetting and retention.
  • Entropy Weight: Combining entropy signals (approx. 30% weight) provided complementary benefits, capturing uncertain decision points that masking alone might miss.

5. Key Contributions

1. Token-Level Attribution Framework: Introduced a principled method to identify knowledge-critical tokens using counterfactual masking and entropy signals without requiring auxiliary models.
2. Selective Unlearning Strategies: Proposed "Hard Selection" and "Soft Weighting" strategies that extend existing unlearning algorithms to operate at the token level.
3. Theoretical Insight: Proved that token-selective updates improve the gradient SNR, explaining why uniform updates fail to optimize the dual objectives of forgetting and retention.
4. Empirical Validation: Demonstrated consistent improvements across three models and two distinct benchmarks (knowledge and safety), achieving up to 32.6% improvement in forgetting and 19.0% improvement in utility preservation.

6. Significance and Impact

• Privacy & Regulation: Enables more precise "Right to be Forgotten" compliance (GDPR/CCPA) by removing specific data traces without degrading the model's general utility.
• Safety: Offers a more effective mechanism for removing hazardous capabilities (e.g., biosecurity threats) from pre-trained models.
• Efficiency: Challenges the assumption that unlearning requires uniform sequence updates, suggesting that knowledge in LLMs is sparse and can be targeted precisely.
• Limitations: The method introduces computational overhead due to the additional forward pass for masking. It is an approximate method and does not guarantee 100% removal of knowledge traces (residual traces may exist).

In conclusion, TokenUnlearn represents a paradigm shift in machine unlearning, moving from coarse-grained sequence updates to fine-grained token attribution, thereby solving the fundamental tension between effective forgetting and utility preservation.