English

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

Computation and Language 2026-05-07 v2

Abstract

Machine unlearning has emerged as a critical capability for addressing privacy, safety, and regulatory concerns in large language models (LLMs). Existing methods operate at the sequence level, applying uniform updates across all tokens despite only a subset encoding the knowledge targeted for removal. This introduces gradient noise, degrades utility, and leads to suboptimal forgetting. We propose TokenUnlearn, a token-level attribution framework that identifies and selectively targets critical tokens. Our approach combines knowledge-aware signals via masking, and entropy-aware signals to yield importance scores for precise token selection. We develop two complementary strategies: hard selection, applying unlearning only to high-importance tokens, and soft weighting, modulating gradient contributions based on importance scores. Both extend existing methods to token-level variants. Theoretical analysis shows token-level selection improves gradient signal-to-noise ratio. Experiments on TOFU and WMDP benchmarks across three model architectures demonstrate consistent improvements over sequence-level baselines in both forgetting effectiveness and utility preservation.

Keywords

Cite

@article{arxiv.2605.00364,
  title  = {Unlearning What Matters: Token-Level Attribution for Precise Language Model Unlearning},
  author = {Jiawei Wu and Doudou Zhou},
  journal= {arXiv preprint arXiv:2605.00364},
  year   = {2026}
}

Comments

17 pages, 2 figures

R2 v1 2026-07-01T12:44:44.152Z