PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Naixin Zhai, Pengyang Shao, Binbin Zheng, Yonghui Yang, Fei Shen, Long Bai, Xun Yang

arXiv:2601.03190v4cs.CL

TL;DR

LLM unlearning seeks to remove sensitive knowledge without sacrificing general utility, but indiscriminate token and full-vocabulary treatment can cause unnecessary degradation. PALU uses dual-sided localized entropy maximization over sensitive prefixes and top-K logits. Across diverse benchmarks, it achieves a strong overall trade-off between forgetting sensitive knowledge and preserving general utility, while its validated scope remains text-based LLMs.

  • Problem

    Existing unlearning methods treat responses and vocabulary dimensions globally, causing unnecessary optimization beyond sensitive content and risking utility degradation.

  • Method

    PALU disrupts sensitive generation trajectories by applying localized entropy maximization to initiating sensitive prefixes and top-K logits.

  • Results

    PALU achieves the state-of-the-art overall trade-off, erasing sensitive knowledge while robustly preserving general LLM utility.

  • Takeaways & Limitations

    Effective unlearning can be achieved through dual locality rather than global suppression across all response tokens and vocabulary dimensions.

  • Takeaways & Limitations

    PALU is designed and validated exclusively on text-based LLMs and has not been extended to multimodal large language models.

Abstract

from arXiv · show

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 alleviate redundant optimization across the full vocabulary and parameter space while minimizing collateral damage to general model performance. Comprehensive evaluations validate that PALU achieves superior forgetting efficacy and utility preservation compared to state-of-the-art baselines. Our code is available at https://github.com/nxZhai/PALU.

1 Introduction

LLM unlearning must remove sensitive knowledge while preserving general utility, but global token and vocabulary optimization can cause unnecessary degradation and computation. PALU addresses this by disrupting sensitive generation prefixes and optimizing only top-K vocabulary candidates.

  • Motivation: Sensitive memorization raises security, ethical, and data-privacy concerns, motivating selective removal without retraining from scratch.Machine unlearning targets sensitive information while supporting compliant LLM deployment.
  • Limitations of Existing Methods: Existing methods can suppress content-agnostic words, disrupt linguistic coherence, degrade utility, and backpropagate through irrelevant positions and vocabulary dimensions.Their global treatment expands optimization beyond tokens and logits relevant to sensitive content.
  • PALU: PALU combines temporal sparsity, which targets pivotal sensitive prefixes, with vocabulary sparsity, which focuses on dominant high-probability candidates.These observations motivate localized optimization across both generation time and vocabulary dimensions.
  • PALU: PALU applies localized entropy maximization to top-K logits, pruning redundant computation and reducing collateral degradation.The method targets sensitive tokens while bypassing context-agnostic tokens and long-tail vocabulary dimensions.
  • PALU: PALU formulates unlearning as disrupting the sensitive generation trajectory with minimal necessary parameter intervention.The framework is based on intervention efficiency rather than global suppression.
  • Results: Empirical results position PALU as improving the trade-off between forgetting quality and utility preservation.The supplied introduction characterizes this trade-off as a new standard.

2 Related Work

Prior LLM unlearning methods use loss-based suppression or token-level intervention, while entropy maximization offers a principled way to increase uncertainty. PALU’s related-work motivation follows from the observation that uncertainty need not be optimized across the full output space.

  • LLM Unlearning: GradientAscent and GradientDiff maximize loss on forget samples, while bounded methods such as NPO address instability through reference-model anchoring.SimNPO removes reference-model bias, and AltPO is another related variant.
  • Token-Level Unlearning: Token-level methods reduce unnecessary perturbations by intervening on subsets of tokens, but typically retain full-vocabulary gradient computation.This leaves vocabulary-space redundancy unresolved.
  • Entropy Maximization: Entropy maximization is used to discourage over-confident predictions and encourage exploration and robustness.The principle selects a less-committal distribution subject to constraints.
  • Localized Entropy: Full-output-space entropy is not always necessary because final predictions are often dominated by a small set of high-probability candidates.This motivates localized entropy maximization over decoding-critical candidates rather than long-tail dimensions.

3 Preliminary

Standard gradient-based unlearning balances forgetting and retention through competing objectives, but the forget loss can introduce unnecessary parameter perturbations. This motivates examining whether negated cross-entropy is an ideal unlearning objective.

  • Objective Formulation: Most gradient-based LLM unlearning methods optimize conflicting forget and retain objectives.The overall objective combines forget loss Lf with retain loss Lr.
  • Retain Objective: The retain loss Lr is typically standard cross-entropy on the retain set Dr, preserving next-token prediction performance.It acts as an indirect constraint against degradation of general capabilities.
  • Forget Objective: The forget loss Lf suppresses the likelihood of generating target output y conditioned on input x.This disrupts the learned mapping from x to y on the forget set.
  • Trade-off: Optimizing Lf can unnecessarily perturb model parameters and degrade general capabilities, while Lr mitigates that degradation.This tension motivates minimizing intervention cost during unlearning.
  • Intervention Efficiency: The paper frames unlearning cost as achieving effective forgetting with minimal necessary parameter perturbations and questions whether negated cross-entropy is ideal.This reframes the preliminary objective as an intervention-efficiency problem.

4 The Proposed Method: PALU

PALU reduces unlearning intervention along both token and vocabulary dimensions: it targets sparse initiating sensitive tokens and maximizes entropy only among top-K logits. This design preserves common tokens through KL constraints, excludes redundant sensitive tokens, and lowers forget-objective complexity from O(T|V|) to O(TK).

  • PALU reduces intervention overhead by localizing unlearning across token positions and vocabulary dimensions.The framework combines semantic-aware token selection with a localized entropy objective.
  • Token-level Unlearning: PALU selects the first N tokens from each sensitive span as initiating targets that disrupt the sensitive generation trajectory.Remaining sensitive tokens are treated as redundant and excluded from the computation graph.
  • Token-level Unlearning: Common content-agnostic tokens receive KL-divergence preservation constraints instead of the unlearning objective.This protects general model capabilities, loss, and fluency while avoiding uniform intervention across the output sequence.
  • Vocabulary-level Unlearning: PALU maximizes entropy within the decoding-critical top-K logits rather than across the full vocabulary.The top-K set is identified using a frozen reference model and remains fixed during unlearning.
  • Vocabulary-level Unlearning: The localized objective flattens dominant logits toward target c, producing a locally higher-entropy distribution while suppressing top-K probability mass when c is sufficiently small.It minimizes variance among dominant logits and anchors their aggregate probability relative to the rest of the vocabulary.
  • Efficiency: O(TK) replaces O(T|V|) for the forget objective when K ≪|V|.Sparse gradients affect only N initiating tokens and top-K vocabulary dimensions, while common-token KL terms and redundant sensitive-token positions remain localized.

5 Experiments

PALU is evaluated on TOFU and MUSE using utility, forgetting, memorization, privacy, and likelihood-based metrics. Across TOFU settings, it improves forgetting while preserving utility, with ablations identifying sparse prefix and top-K interventions as effective and efficient.

  • 5.1 Experiment Setup: PALU is evaluated on TOFU and MUSE with complementary utility, forgetting, memorization, privacy, fluency, and likelihood-based metrics.TOFU includes 1%, 5%, and 10% forget splits; MUSE evaluates real-world News and Books memorization and privacy risks.
  • 5.2 Overall Performance: 13.4% and 28.0% relative FQ improvements over TPO are reported on Llama-2-7B and Llama-3.1, respectively, on TOFU forget 5%.The improvement widens on the more capable Llama-3.1 model.
  • 5.2 Overall Performance: MU reaches 0.6238 on Llama-2-7B, exceeding TPO’s 0.5862 and nearly matching Retain’s 0.6266.The reported comparison accompanies PALU’s strongest overall forgetting-utility trade-off.
  • 5.2 Overall Performance: Across forget 1% and 10% settings, PALU’s performance profile remains closest to Retain, while NPO and DPO deteriorate in the demanding 10% setting.The comparison spans Llama-2-7B and Llama-3.1-8B.
  • 5.3 Ablation Study: Validating Dual-Sparsity: Performance stabilizes at prefix length N = 3, while extending optimization beyond three tokens adds virtually no forgetting gains and increases computational burden linearly.Optimizing only one token is reported as too abrupt and suboptimal for utility; the prefix result supports temporal sparsity.
  • 5.5 Training Efficiency and Convergence: PALU saturates FQ by Epoch 5 and recovers MU within 2 epochs, effectively halving training duration relative to NPO.The convergence analysis is conducted on TOFU forget 5% over 10 epochs, and the efficiency claim is paired with gradient sparsity.

6 Conclusion

PALU uses dual-sided localized entropy maximization to target sensitive generation paths while avoiding unnecessary global intervention. Across diverse benchmarks, it reports an overall trade-off between erasing sensitive knowledge and preserving general utility.

  • 6 Conclusion: PALU targets the sensitive prefix temporally and the top-K logits in vocabulary space rather than suppressing all tokens globally.This dual-locality mechanism is designed to sever sensitive generation paths while reducing redundant updates.
  • 6 Conclusion: PALU achieves the state-of-the-art overall trade-off across diverse benchmarks, erasing sensitive knowledge while preserving general model utility.

Limitations

The paper’s theoretical analysis contrasts negated cross-entropy with localized entropy maximization, while the framework is currently limited to text-based LLMs. The multimodal setting remains outside the validated scope.

  • Limitations: PALU is designed and validated exclusively on text-based LLMs, not multimodal models integrating visual or audio modalities.Adapting sensitive-prefix detection and logit-confusion measurement to multimodal sequences requires further investigation.
  • Limitations: Negated cross-entropy preserves logit ratios among non-target tokens, making it vulnerable to synonym substitution, whereas entropy maximization induces uncertainty in the semantic space.

A.2 Limitations of Negated CE

Negated cross-entropy mainly lowers the target token’s logit while preserving relative relationships among alternatives, allowing a synonym to recover the underlying concept. Localized entropy maximization instead flattens the critical top-K subspace without globally flattening the vocabulary.

  • A.2 Limitations of Negated CE: Negated cross-entropy minimizes the likelihood of the target token, primarily decreasing its logit during optimization.
  • A.2 Limitations of Negated CE: For non-target tokens, negated cross-entropy does not explicitly alter logit distances, so their probability ratios remain invariant under a uniform normalization update.
  • A.2 Limitations of Negated CE: As the target probability approaches zero, its probability mass is redistributed proportionally according to the original probabilities.
  • A.2 Limitations of Negated CE: The synonym can therefore recover the underlying concept after the specific target word is forgotten, leaving sensitive information incompletely erased.
  • A.2 Limitations of Negated CE: Localized entropy maximization minimizes MSE between top-K logits and a uniform target, flattening the decoding-critical candidate set that includes the target and its synonym.
  • A.2 Limitations of Negated CE: Restricting entropy maximization to the decoding-critical subspace avoids suppressing tail tokens that support linguistic fluency, unlike global vocabulary flattening.

B Training Procedure

PALU’s training procedure computes localized unlearning objectives over selected initial sensitive tokens and top-K logits, while evaluating forgetting, utility, privacy, memorization, and fluency. The experimental setup compares PALU with multiple unlearning baselines and uses several complementary metrics.

  • B Training Procedure: Algorithm 1 computes PALU’s forget loss by enforcing local entropy maximization on top-K candidate logits within initial tokens.
  • B Training Procedure: PALU is evaluated alongside GradientAscent, GradDiff, DPO, SimNPO, and PDU as representative unlearning approaches.
  • B Training Procedure: Additional metrics assess fluency, verbatim memorization, knowledge memorization, privacy leakage, and truth ratio.
  • B Training Procedure: Forget Quality measures distributional similarity between the unlearned model and a retain model using a Kolmogorov-Smirnov test.
  • B Training Procedure: Model Utility is the harmonic mean of normalized probability, ROUGE score, and truth ratio across retain and non-forget datasets.
  • B Training Procedure: The procedure identifies decoding-critical subspaces using reference logits and applies preservation constraints to common tokens.

E Additional Implementation Details

PALU’s implementation and positional analyses use controlled tuning and token-selection studies to test whether localized interventions preserve the forgetting–utility balance. Results across model scales support targeting initial prefix tokens rather than random or full-sequence positions.

  • Training and tuning: All methods use 10 training epochs, batch size 32, paged AdamW, and a one-epoch linear warmup on 8 NVIDIA A800 GPUs.Learning rates and λ are tuned over predefined grids; PALU additionally searches top-K and Initial-N sparsity settings.
  • Token-position analysis: Table 6 compares token positional strategies on Llama-2-7B using the 5% TOFU forget split.The table is the controlled evaluation of positional choices described in this section.
  • Token-position analysis: Random-token suppression yields suboptimal forgetting, while full-sequence suppression improves forgetting but slightly reduces utility.These findings identify token position and intervention scope as important design factors.
  • Token-position analysis: Targeting only the first few prefix tokens achieves the best balance between forgetting efficacy and utility preservation.The result supports the prefix-locality assumption.
  • Scale analysis: At the 13B scale, localizing interventions rather than enforcing uncertainty across entire responses and vocabularies yields effective forgetting.The reported conclusions are consistent with those observed at the 7B scale.

F.3 Performance on MUSE

On MUSE, PALU reduces memorization measures on the News subset while retaining general knowledge, illustrating the intended benefit of prefix-aware suppression. The reported results emphasize privacy-risk mitigation without the model collapse observed for GA.

  • News subset: VerbMem on News falls from 0.5789 initially to 0.2700 with PALU, outperforming NPO and approaching the Retain model.This result is reported as evidence of improved privacy-risk mitigation on the MUSE benchmark.
  • News subset: KnowMem on Df decreases from 0.6443 for the Original model to 0.4442 with PALU on News.The passage describes this as deeper unlearning than the Original model.
  • General-knowledge retention: PALU maintains robust general knowledge on Dr, unlike GA, which suffers complete model collapse.This comparison accompanies the News-subset memorization results.
  • Mechanism: Suppressing initial dominant tokens is reported to be sufficient to sever generation of sensitive sequences.The result connects PALU’s prefix-aware mechanism with its observed MUSE behavior.

G.2 Analysis of Sensitive Information Lengths

TOFU analyses find that sensitive information occupies a sparse portion of responses and is usually concentrated in short entities. Generation examples then contrast irrelevant, garbled, and targeted suppression behaviors across unlearning methods.

  • Entity-length distribution: Most sensitive entities contain only 2 to 6 tokens, while entities exceeding 10 tokens are extremely rare.This right-skewed distribution supports using a small initiating prefix budget.
  • Efficiency rationale: PALU’s mask-based strategy focuses optimization on sparse sensitive positions and their prefixes, reducing redundant computation on non-sensitive text.The passage links sparse target-token distributions to the efficiency rationale for localized optimization.
  • Generation examples: NPO often produces irrelevant generic or background text, whereas TPO’s stronger sensitive-token suppression can produce garbled outputs.The examples compare these behaviors with PALU and retained ground-truth information.
Loading 2601.03190v4…