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

Abstract: Machine unlearning aims to remove targeted knowledge from a trained model while preserving its general capabilities. For autoregressive LLMs, 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 LLMs should forget, enabling unsupervised learning of token level forget-specificity directly from model representations with minimal computational overhead.

• The paper introduces ATWU, a novel method that learns token-level importance to selectively erase sensitive information without degrading overall model utility.
• It formulates unlearning as a bilevel optimization problem, using a linear scorer to assign token weights based on the balance between forgetting and retention objectives.
• Experimental results show that ATWU achieves up to 91.7% unlearning quality, outperforming existing methods while maintaining robust language performance.

Token-Level Relevance in Machine Unlearning: Alternating Token-Weighted Unlearning for LLMs

Introduction

Machine unlearning for large-scale autoregressive LLMs remains a technically complex challenge at the intersection of privacy, legal compliance, and model utility. The paper "Learning What to Forget: Improving LLM Unlearning via Learned Token-Level Importance" (2606.06320) addresses the problem of determining which tokens in a forget set are actually responsible for the information that should be unlearned. Rather than relying on external annotation, heuristics, or auxiliary models to assign token importance, the authors formalize token forget-specificity based on the compatibility between forgetting and retaining objectives. They introduce Alternating Token-Weighted Unlearning (ATWU), a theoretically-justified and empirically robust framework for end-to-end learning of token importance in LLM unlearning.

Problem Formalization and Theory

Standard unlearning protocols often degrade model utility by naively applying the forget loss to all tokens in the forget sample, failing to recognize that most tokens may encode only non-specific linguistic regularities rather than the sensitive knowledge to be erased. Token-weighted unlearning has been introduced in recent work to address this, but prior assignments of token weights are extrinsic, based on human or LLM-generated annotation, auxiliary contrastive models, or confidence heuristics. The present work replaces such mechanisms with an endogenous, model-driven approach based on an explicit notion of retain conflict.

Concretely, the central insight is that a token is forget-specific if suppressing its likelihood under the forget loss does not incur substantial degradation with respect to the retain objective. Under a monotonicity and separation assumption on the tokenwise losses, the joint minimization of model parameters and a token-weight vector (subject to a budget constraint) provably recovers the desired subset of forget-specific tokens. The work establishes global optimality and exact recovery guarantees for this bilevel optimization when the retain and forget loss are suitably separated.

Alternating Token-Weighted Unlearning (ATWU)

To instantiate this principle computationally, the authors relax the combinatorial binary optimization of token-weight vectors into a regularized, continuous optimization with entropic and budget penalties driving token weights towards sparsity and prescribed support. Crucially, for scalability and generalization, these token weights are not learned independently, but are parameterized by a single linear scorer applied to each token’s hidden state. This design choice is motivated by the structural probing literature, which shows that intermediate representations in transformers are linearly separable by a wide range of linguistic and semantic properties [alain2017understanding, tenney-etal-2019-bert].

The ATWU procedure alternates between updating model parameters using the scorer-fixed token-weighted loss and updating the linear token importance scorer to optimize end-task performance with the model frozen. The forget loss used in ATWU is a bounded, saturation-weighted negative cross-entropy that robustly penalizes forget-specific tokens without catastrophic degradation of model utility. The framework supports alternative loss functions including DPO, NPO, and SimNPO by inserting the scorer as a weighting function into any autoregressive forget-side objective.

Experimental Results

ATWU is evaluated on the TOFU [maini2024tofu] and RWKU [jin2024rwku] LLM unlearning benchmarks with strong baselines spanning sample-level, probability-based, and externally-annotated token-weighted objectives. Quantitative and qualitative results demonstrate that ATWU achieves the best or second-best unlearning quality—measured as net reduction in sensitive knowledge recall under robust paraphrase attacks and minimal degradation in retain set and general model utility—across a range of forget/retain splits and model scales.

Numerical highlights include:

• On TOFU (Llama-3.1-8B-Instruct, forget10 split), ATWU attains unlearning quality (UQ) of 91.7%, exceeding previous state-of-the-art by 5.8 percentage points and matching or surpassing all competitors on MMLU and repetitiveness metrics.
• ATWU outperforms auxiliary-model, probability-, and annotation-based baselines in both score assignment (per-sample AUROC up to 75±9 against ground-truth annotation) and unlearning efficacy.
• ATWU is effective across forget losses (DPO, NPO, SimNPO), always improving UQ over the unweighted variant—evidence that the method improves the core selective forgetting signal, not simply the implementation of a particular objective.

  Figure 1: ATWU’s scorer sharply focuses on the answer-bearing tokens (forget-specific) compared to probability-based and heuristic baselines, which spread mass over both structural and forget-specific tokens.

Token Scorer Analysis

The linear scorer parameterized over LLM hidden representations, trained in the absence of token-level annotation, aligns strongly with ground-truth forget annotations. Supervised training with gold labels is only marginally stronger than the unsupervised ATWU variant, and the gap is closed further when the scorer is injected multiplicatively in the saturation exponent. Qualitative analysis shows that ATWU's token importance is both sparse and highly aligned with the semantically relevant answer spans; transfer experiments indicate that the learned scorer generalizes to out-of-distribution queries and real-world entities.

Ablations and Robustness

Ablation studies confirm:

• Performance improvements depend crucially on budget regularization and entropy constraints in the scorer’s learning objective.
• Token-weighted forget losses with poor underlying gradient scaling (e.g., non-saturated negative cross-entropy) degrade both forgetting and retention, even with oracle annotation. By contrast, saturation-weighted losses recover most of the supervised upper bound using unsupervised ATWU scoring.
• Frozen, pre-trained scorers are significantly less effective; strong results require scorer and model co-adaptation during unlearning.
• Scheduled scorer refresh (alternating optimization) is superior to joint parameter updates or fixed scorer approaches.

Generalization and Limitations

While ATWU significantly revises the paradigm of token selection in unlearning objectives, its theoretical guarantees hold under structural token separation that may only be approximately satisfied in practical data. Empirical efficacy is robust across benchmarks, but gains are largest when the forget set is sufficiently informative for the scorer to generalize. The procedure requires no external annotation, supporting privacy-by-design in sensitive applications.

Conclusion

The ATWU framework recasts token-weighted unlearning as a joint optimization of model and token importance, providing both a theoretical blueprint and practical mechanism for extracting selective unlearning signals from model representations. The results challenge the assumption that improved token identification alone yields robust selective unlearning, highlighting the importance of objective design and the synergy between scorer and model optimization. By sidestepping the need for external supervision and auxiliary models, ATWU sets a new operational standard for utility-preserving, privacy-compliant LLM unlearning and motivates further study in joint identification and intervention in large-scale generative models.

Outlook

The framework’s generality immediately suggests several avenues for future research: (1) non-linear or context-aware token importance scoring, (2) adaptive budget selection and dynamic support constraints, and (3) integration within broader privacy, audit, and policy compliance toolchains for foundation models. As demands for targeted erasure rise in deployed LLMs, token-level, end-to-end optimized approaches such as ATWU will likely underpin both practical and regulatory advances in machine unlearning.