One-sentence Summary

Researchers from the University of Science and Technology of China and the National University of Singapore propose PALU (Prefix-Aware Localized Unlearning), a framework that maximizes local entropy by suppressing only sensitive prefixes and flattening $\text{top-}K$top-K logits, thereby reducing redundant optimization across the full vocabulary and parameter space to achieve superior forgetting efficacy with minimal collateral damage to general model performance.

Key Contributions

• PALU is introduced as a prefix-aware localized unlearning framework for Large Language Models, using a local entropy maximization objective over sensitive prefixes and top-K logits.
• By suppressing only the sensitive prefix and flattening only the top-K logits, PALU severs sensitive generation links while avoiding optimization over the full vocabulary and full parameter space, reducing collateral damage to general model performance.
• Comprehensive evaluations show that PALU achieves a stronger trade-off between forgetting efficacy and utility preservation than state-of-the-art gradient-based unlearning baselines.

Introduction

Large language models can memorize private or copyrighted content, making machine unlearning important for secure deployment and compliance with privacy regulations such as GDPR and CCPA. Existing methods mostly rely on negated cross-entropy objectives, which suppress target tokens globally across response positions and vocabularies, often causing over-correction, lost coherence, reduced utility, and high computational overhead. The authors leverage the observation that sensitive generation is usually triggered by a short prefix and dominated by a small set of high-probability logits. They introduce PALU, which applies localized entropy maximization over the sensitive prefix and top-K logits to disrupt the undesired generation path with minimal perturbation, achieving strong forgetting while better preserving general utility.

Method

The authors present PALU as a low-overhead machine unlearning framework for large language models. The design reduces intervention cost from two complementary perspectives: the token level and the vocabulary level. At the token level, it avoids applying the unlearning objective uniformly to every output token. Instead, it distinguishes sensitive initiating tokens, content agnostic tokens, and redundant sensitive tokens, assigning different training signals to each group. At the vocabulary level, it replaces naive negative cross entropy and full vocabulary entropy maximization with a localized entropy objective that targets only the dominant logits involved in decoding. As shown in the framework below, the token selection module determines where gradients are applied, while the vocabulary-level objective determines how each selected target should be transformed.

The first module constructs a token-level role partition. For an output sequence of length $T$T, the method identifies sensitive spans and defines a binary mask $m_t \in \{0,1\}$mt​∈{0,1}, where $m_t = 1$mt​=1 indicates that token $y_t$yt​ carries sensitive semantics and $m_t = 0$mt​=0 indicates a common token. Let

be the set of sensitive token indices. The authors further exploit temporal sparsity by assuming that only the first few tokens in each sensitive span are pivotal for triggering sensitive behavior. They therefore define an initiating target subset

by selecting the first $N$N sensitive tokens in each span.

Based on this partition, output tokens are assigned three optimization roles. Tokens in $\mathcal{T}_{\mathrm{init}}$Tinit​ receive the main unlearning signal because they are responsible for initiating the sensitive trajectory. Tokens outside $\mathcal{T}_{\mathrm{sens}}$Tsens​ are treated as common tokens and are protected with a KL divergence constraint to preserve general fluency and utility. The remaining sensitive tokens, those in $\mathcal{T}_{\mathrm{sens}} \setminus \mathcal{T}_{\mathrm{init}}$Tsens​∖Tinit​, are excluded from gradient computation under a minimal intervention principle.

The second module addresses the vocabulary-level optimization objective. Negative cross entropy can suppress the probability of a target token, but it does not guarantee that the distribution becomes truly uncertain. The suppressed probability mass may shift to a correlated synonym or another dominant token, producing an unstable and peaked distribution. Global entropy maximization offers a stronger theoretical target because it encourages the predictive distribution to become uniform over the full vocabulary, but optimizing entropy over all $|\mathcal{V}|$∣V∣ dimensions is computationally expensive and often unnecessary since long-tail low-probability tokens have little influence on generation.

To balance robustness and efficiency, the authors introduce local entropy maximization. For a target token $t \in \mathcal{T}_{\mathrm{init}}$t∈Tinit​, let $z_t \in \mathbb{R}^{|\mathcal{V}|}$zt​∈R∣V∣ denote the logit vector. They define a top-K set $\mathcal{V}_{\mathrm{top}}$Vtop​, identified using the frozen reference model $P_{\theta_{\mathrm{ref}}}$Pθref​​ and kept fixed throughout training. The localized objective is

This objective reduces the variance among the dominant logits by pushing them toward a common target value $c$c. When the top-K logits become closer to one another, softmax normalization yields a more uniform distribution over decoding critical candidates. The anchor value $c$c also controls suppression strength: choosing a small $c$c can reduce the total probability mass assigned to the top-K candidates relative to the rest of the vocabulary.

The overall forget loss combines sparse token selection with local entropy maximization. It is written as

The first term optimizes only the sparse initiating tokens in the top-K logit space. The second term preserves general behavior by penalizing divergence from the reference model on common tokens. Redundant sensitive tokens receive no direct gradient because they are not included in either expectation.

From a computational perspective, this design changes the dependency of the backward pass from the full vocabulary to a small decoding subspace. Dense negative cross entropy incurs a cost proportional to $T|\mathcal{V}|$T∣V∣, while PALU induces sparse gradients on at most $T$T initiating token positions and only $K$K vocabulary dimensions. Thus the unlearning operation is bounded by $\mathcal{O}(TK)$O(TK), which is substantially lower than $\mathcal{O}(T|\mathcal{V}|)$O(T∣V∣) when $K \ll |\mathcal{V}|$K≪∣V∣.

Experiment

Experiments on TOFU and MUSE with Llama backbones against multiple unlearning baselines show that PALU provides a strong forgetting-utility balance, scales well with model capability, and closely approaches the ideal retain model in overall behavior. Ablations validate the dual-sparsity design by showing that limited vocabulary truncation and early prefix truncation are sufficient for effective unlearning, while global-mean flattening offers the best stability for preserving the model's distribution. Additional analyses show rapid convergence, low target memorization under likelihood-based separation metrics, and promising robustness across different forgetting settings.

In the forget 5% setting, PALU provides the strongest overall balance between forgetting and utility among the compared methods, and its advantage grows on the more capable Llama model. It avoids the severe utility collapse observed for GA and GD while remaining close to the retained reference performance profile. PALU achieves larger forgetting-quality gains over strong baselines, with the improvement becoming more pronounced on the stronger model. GA and GD degrade model utility substantially, while PALU preserves fluency and answer quality. The dual-locality approach supports erasure without sacrificing as much retained capability as competing methods.

The extended likelihood metrics indicate that several methods leave forget and holdout samples more separable than the retain model, while PALU produces consistently low separability on Llama-2-7B. On Llama-3.1-8B, PALU also suppresses target-related likelihood more strongly than TPO and can fall below the retain model, suggesting that forgotten content may receive lower target likelihoods than ordinary unseen data. PALU achieves strong target suppression across LOSS, ZLib, MinK, and MinK++ on Llama-2-7B. GA and GD retain higher likelihood separability, closer to the original model than PALU. On Llama-3.1-8B, PALU lowers MinK++ separability compared with TPO and can produce lower LOSS separability than Retain.

The comparison evaluates fixed prefix truncation ratios and PALU for Llama-2-7B on the 5% forget split of TOFU. A 25% truncation ratio causes a clear drop in forgetting quality, while the 50% and 100% ratios perform much better and preserve reasonable utility. PALU gives the best overall balance, slightly outperforming the closest fixed ratio while maintaining the highest utility among the compared settings. The 50% ratio is close to PALU, indicating that a moderate length-aware budget can approximate targeted removal of causally dominant tokens. Very aggressive truncation at 25% substantially weakens forgetting, showing that too little prefix can limit effective unlearning. Dynamic truncation appears promising, but consistent gains require better budget selection and finer hyperparameter tuning.

The experiments evaluate unlearning methods on forget 5% settings using both smaller and stronger Llama models, with additional extended likelihood analyses and a comparison against fixed prefix truncation strategies. PALU shows the strongest overall balance between forgetting and utility, avoiding the severe utility collapse seen with GA and GD while preserving fluency and answer quality. It also achieves strong suppression of target-related likelihoods, indicating robust forgetting without overly degrading retained capability.