REVS: Unlearning Sensitive Information in Language Models via Rank Editing in the Vocabulary Space

Language models (LMs) risk inadvertently memorizing and divulging sensitive or personally identifiable information (PII) seen in training data, causing privacy concerns. Current approaches to address this issue involve costly dataset scrubbing, or model filtering through unlearning and model editing, which can be bypassed through extraction attacks. We propose REVS, a novel non-gradient-based method for unlearning sensitive information from LMs. REVS identifies and modifies a small subset of neurons relevant for constituent tokens that form sensitive information. To adequately evaluate our method on truly sensitive information, we curate three datasets: email and URL datasets naturally memorized by the models, and a synthetic social security number dataset that we tune the models to memorize. Compared to other methods, REVS demonstrates superior performance in unlearning sensitive information and robustness to extraction attacks, while retaining underlying model integrity.

Language models risk inadvertently memorizing and divulging sensitive or personally identifiable information (PII) seen in their training data, causing privacy concerns. Current approaches to address this issue involve costly dataset scrubbing or model filtering through unlearning and model editing, which can be bypassed through extraction attacks. We demonstrate that models like Llama-3-8B and GPT-J-6B naturally memorize sensitive information from The Pile training dataset, including email addresses and URLs.

REVS operates through two main phases: localization and editing. First, it selects the rarest tokens from the target sensitive information sequence, as accurately extracting the original sensitive data requires recovering the full token sequence. It then locates the relevant layers for those tokens and selects neurons within those layers that are both relevant to the given prompt and have high association with the target token information. REVS iteratively edits these selected neurons to reduce the rank of the target tokens below a specified threshold, effectively demoting the sensitive information while preserving the model's broader knowledge.

The key insight is that transformer MLP layers construct predictions by promoting specific tokens in the output vocabulary space. By identifying and modifying FF₂ columns (neurons) that contribute most to generating target tokens, REVS surgically removes the model's encoded tendency to generate sensitive data.

Our methodology REVS uses a sophisticated two-step localization process:

Layer Selection: We identify relevant layers by measuring how strongly each layer contributes to generating the target token. Layers where the target token rank exceeds a predetermined threshold are selected for editing.

Neuron Selection: Within selected layers, we identify neurons using two criteria: (1) Activation strength - how strongly a neuron is activated in response to the prompt, and (2) Token association - how strongly the neuron is associated with the target token when projected to vocabulary space. This hybrid approach targets neurons that are both contextually significant and semantically relevant, outperforming alternatives based solely on activations, gradients, or random selection.

We select the rarest tokens (typically 2 tokens per sequence) from the target sensitive information, as unlearning every token is unnecessary for preventing accurate recovery of the original data.

To rigorously evaluate unlearning on truly sensitive information, we curated three benchmark datasets:

Unlike prior work that evaluated editing methods on non-sensitive data, our benchmarks contain actual private information, enabling rigorous assessment of unlearning efficacy and robustness against extraction attacks.

We evaluate unlearning through three critical dimensions:

Our extraction attacks are stricter and more comprehensive than prior work, considering more candidate tokens across all model layers.

REVS demonstrates superior performance across all evaluation metrics compared to six strong baselines including MEMIT, Constrained Fine-tuning (FT-L), NPO-KL, RMU, Head-Projection, and Max-Entropy methods. Key findings include:

• Superior Unlearning Performance: REVS achieves the highest Unlearning Score across all datasets (89.58 on SSN, 62.37 on Emails, 44.25 on URLs)
• Robust Extraction Resistance: Achieves highest or second-highest Resistance Score across all datasets and models
• Preserved Model Integrity: Maintains strong specificity while achieving high effectiveness, with minimal impact on general capabilities (MMLU/GSM8K scores remain stable)
• Consistent Performance: Demonstrates robustness across different hyperparameters, candidate token set sizes, and number of targets

Notably, our experiments revealed inherent differences between synthetic and organically memorized information, with organically memorized data proving more challenging to unlearn, suggesting that information memorized during pre-training may be more deeply ingrained in model parameters.

As shown in the radar chart below and detailed results, REVS achieves the highest Resistance Score across all datasets and models, except for Emails on Llama-3-8B, where it ranks a close second. Notably, the results highlight a clear link between higher Effectiveness and improved Resistance Score. While strong Effectiveness often compromises specificity, REVS maintains both, achieving the best or second-best Resistance Score and demonstrating balanced robust unlearning.

The extraction resistance evaluation demonstrates that REVS not only successfully removes sensitive information from model outputs but also makes it significantly more difficult for adversarial attacks to recover the unlearned data. This robustness is crucial for real-world deployment where models may face sophisticated extraction attempts.