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

2026-06-04

AI Summary

This work addresses the challenge of fine-grained knowledge unlearning in large language models by proposing an unsupervised method that identifies tokens most specific to the target knowledge to be forgotten. The approach innovatively leverages the conflict between forgetting and retention objectives as a token-level criterion for unlearning importance, enabling the model to learn genuine forgettable regions without external supervision. It employs an alternating optimization framework, ATWU, which jointly optimizes model parameters and hidden-state-based token weights. Evaluated on the TOFU and RWKU benchmarks, the method significantly outperforms existing sample-level, heuristic, and auxiliary-model-dependent approaches, achieving state-of-the-art trade-offs between forgetting efficacy and knowledge retention. Moreover, the learned token scores exhibit strong alignment with ground-truth forgettable segments.

Abstract

Machine unlearning aims to remove targeted knowledge from a trained model while preserving its general capabilities. For autoregressive language models, not all tokens in a forget sample are equally relevant to forgetting. Existing approaches either ignore this heterogeneity or rely on auxiliary models, heuristics, or external annotations to estimate each token's relevance for forgetting. We instead characterize it through the interaction with the retain objective: a token is forget-specific to the extent that minimizing the forget loss on that token does not conflict with retain optimality. We formalize this perspective as a joint optimization problem over the model parameters and the token weights and show that, under a natural separation condition, the resulting objective recovers the oracle forget-specific token support. Motivated by this formulation, we introduce Alternating Token-Weighted Unlearning (ATWU), a lightweight framework that jointly learns token forget-specificity and model parameters during unlearning using a simple linear scorer over the hidden states, without external token level supervision. Across TOFU and RWKU, ATWU achieves state of the art forget-retain trade-offs, outperforming sample-level methods, probability-based token weighting heuristics, and auxiliary-model-based approaches. Moreover, the learned scores align substantially better with ground truth forget-specific spans, indicating that ATWU identifies semantically meaningful token level forgetting signals. Overall, our results suggest that retain conflict provides an effective criterion for identifying what language models should forget, enabling unsupervised learning of token level forget-specificity directly from model representations with minimal computational overhead.

Problem

• machine unlearning
• token-level importance
• forget-retain trade-off
• autoregressive language models
• knowledge removal

Innovation

• machine unlearning
• token-level importance
• retain-forget conflict
• alternating optimization
• unsupervised token weighting