Learning What to Forget: Improving LLM Unlearning via Learned Token-Level Importance
For practitioners needing to remove specific knowledge from LLMs while preserving general capabilities, this work provides a lightweight, unsupervised token-level unlearning method that improves the forget-retain trade-off without external supervision.
The paper proposes Alternating Token-Weighted Unlearning (ATWU), a method that jointly learns token-level importance and model parameters for LLM unlearning by leveraging retain conflict. On TOFU and RWKU benchmarks, ATWU achieves state-of-the-art forget-retain trade-offs, outperforming sample-level and heuristic token-weighting methods.
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.