Don't Forget Your Embeddings: Robust Knowledge Erasure via Precise Editing of Embeddings
For developers needing robust and persistent removal of specific knowledge from LLMs, this work addresses a critical safety and compliance bottleneck by showing that embedding-level intervention is necessary and effective.
Existing knowledge erasure methods for language models overlook the embedding layer, leading to vulnerability to adversarial recovery. EMBER, a plug-n-play module using sparse matrix factorization, precisely erases concept-related features from token embeddings, reducing relearning accuracy by up to 50% (to 35% on Llama vs. 70-76% for prior methods) with minimal coherence loss.
As language models are increasingly deployed in real-world applications, the ability to erase specific knowledge from them becomes critical for safety and compliance. Prominent methods seek persistent removal by updating the model's parameters, yet the target knowledge often can be recovered through adversarial prompting or relearning. In this work, we hypothesize this limitation stems in part from existing methods overlooking the embedding layer. To address this, we introduce EMBedding ERasure (EMBER), a plug-n-play erasure module that leverages Sparse Matrix Factorization for precise erasure of concept-related features from token embeddings. Through comprehensive evaluations across diverse concepts on Gemma-2-2B-it and Llama-3.1-8B-Instruct, we find that augmenting existing methods with EMBER consistently improves erasure efficacy and specificity across task formats, with minimal coherence loss. Moreover, it dramatically improves robustness to relearning, reducing regained accuracy by up to 50%, limiting it to 35% on Llama compared to 70%-76% for prior methods. Further analysis shows that the coherence cost is localized, affecting only a small set of concept-exclusive tokens. Our work establishes that precise embedding-level intervention is necessary for robust concept erasure, and demonstrates that existing methods can benefit from such augmentation.