Maximizing Local Entropy Where It Matters: Prefix-Aware Localized LLM Unlearning

TL;DR

The paper addresses the challenge of forgetting sensitive content in large language models while preserving general utility. It introduces PALU, a dual-locality unlearning framework that surgically targets the sensitive prefix in time and flattens only the decoding-critical top-$K$ logits via localized entropy maximization, guided by a frozen reference model. This yields an efficient $\mathcal{O}(T K)$ optimization cost and superior forgetting-utility trade-offs compared to state-of-the-art baselines on TOFU and MUSE benchmarks. The work demonstrates strong empirical results, rapid convergence, and robust performance across models, while outlining limitations to text-only LLMs and proposing future work for multimodal extensions.

Abstract

Machine unlearning aims to forget sensitive knowledge from Large Language Models (LLMs) while maintaining general utility. However, existing approaches typically treat all tokens in a response indiscriminately and enforce uncertainty over the entire vocabulary. This global treatment results in unnecessary utility degradation and extends optimization to content-agnostic regions. To address these limitations, we propose PALU (Prefix-Aware Localized Unlearning), a framework driven by a local entropy maximization objective across both temporal and vocabulary dimensions. PALU reveals that (i) suppressing the sensitive prefix alone is sufficient to sever the causal generation link, and (ii) flattening only the top-$k$ logits is adequate to maximize uncertainty in the critical subspace. These findings allow PALU to avoid redundant optimization across the full vocabulary and parameter space while minimizing collateral damage to general model performance. Extensive experiments validate that PALU achieves superior forgetting efficacy and utility preservation compared to state-of-the-art baselines.

Figures (9)

• Figure 1: Illustration of the vocabulary-localized optimization. We specifically target sensitive tokens (red) while bypassing context-agnostic ones (gray). For each target position, the optimization is restricted to the top-$K$ vocabulary candidates (indicated by $\checkmark$), thereby pruning the computation on long-tail dimensions.
• Figure 2: Overview of PALU.(left) Token-level Unlearning. Comparison of how different methods (classic, current token-aware, and ours) distinguish between token roles to identify specific optimization targets. (right) Vocabulary-level Unlearning. Visualization of theoretical probability distributions induced by different objectives (PDU, -CE, and our Local Entropy Maximization) relative to the original LLM prediction.
• Figure 3: Analysis of PALU. We evaluate the impact of (left) the logit truncation size, (middle) the prefix length, and (right) the target threshold strategies, i.e., Uniform, Global Mean, Local Mean. Blue bars represent Forget Quality (left y-axis), and red bars represent Model Utility (right y-axis).
• Figure 4: Performance on TOFU forget 1% and 10% split for different unlearning methods on different models.
• Figure 5: Convergence Analysis. Model Utility (MU) and Forget Quality (FQ) versus unlearning epochs for PALU and NPO. The results are shown for the forget 5% split on the TOFU dataset over 10 epochs.