💬 NLP

Don't Forget Your Embeddings: Robust Knowledge Erasure via Precise Editing of Embeddings

This paper introduces EMBER, a plug-and-play module that enhances robust knowledge erasure in language models by precisely removing concept-related features from token embeddings via sparse matrix factorization, thereby significantly improving resistance to relearning while minimizing coherence loss.

Original authors: Clara Haya Suslik, Or Shafran, Mor Geva

Published 2026-06-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Clara Haya Suslik, Or Shafran, Mor Geva

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a Large Language Model (LLM) like a giant, super-smart librarian who has read almost every book in the world. Sometimes, we need this librarian to "forget" specific things—maybe a character from a book that hasn't been published yet, a dangerous chemical recipe, or a celebrity's private information. This process is called Knowledge Erasure.

For a long time, scientists tried to make the librarian forget by rewriting the "instruction manuals" inside their brain (specifically, a part called the MLP layer). They would go in and erase the notes related to the forbidden topic.

The Problem: The Hidden Notebook
The paper argues that these previous attempts often fail. Why? Because the librarian has a second, hidden notebook: the Embedding Layer.

Think of the Embedding Layer as the librarian's index cards. Every word they know (like "Harry," "Potter," or "Wand") has a specific card. Even if you burn the instruction manual notes about Harry Potter, the index cards for those words still hold the secret. If someone asks, "Who is Harry?", the librarian can still pull up the card, see the connections, and accidentally "relearn" the forbidden knowledge very quickly.

The Solution: EMBER
The authors introduce a new tool called EMBER (Embedding ERasure). Instead of just burning the instruction manual, EMBER goes straight to the index cards.

Here is how EMBER works, using a simple analogy:

The Mix: Imagine the index card for the word "Harry" is actually a smoothie made of many different ingredients (features). Some ingredients are general (like "is a person's name"), and some are specific to the forbidden topic (like "wizard," "Hogwarts," "magic").
The Separation: EMBER uses a mathematical trick called Sparse Matrix Factorization. Think of this as a high-tech blender that can perfectly separate the "wizard" ingredients from the "person" ingredients on the card.
The Surgery: Once separated, EMBER carefully scoops out only the "wizard" ingredients from the "Harry" card. It leaves the "person" ingredients alone.
The Result: Now, when the librarian sees the word "Harry," the card no longer triggers the "wizard" connection. The librarian can still talk about Prince Harry (the real person) perfectly fine, but they are completely blank on Harry Potter.

Why This is a Big Deal
The paper tested this on two popular AI models (Gemma and Llama) with 18 different topics, ranging from "Harry Potter" to "World War II."

It's Harder to Cheat: Previous methods were like putting a "Do Not Read" sign on a book. The librarian could still peek and remember. EMBER is like removing the book from the shelf entirely. Even if someone tries to "retrain" the librarian with a few new notes, the librarian can't remember the forbidden topic because the foundation (the index card) is gone.
It Doesn't Break Everything: A common fear is that if you edit the librarian's brain, they might start speaking gibberish or forgetting how to talk about other things. The paper found that EMBER is incredibly precise. It only edits a tiny fraction of the words (less than 0.14% of the dictionary). If you erase "Harry Potter," the librarian can still talk about "Harry Styles" or "Prince Harry" without stumbling.
It Works with Old Methods: EMBER isn't a replacement for the old methods; it's a booster. When you use EMBER alongside the old instruction-manual edits, the results are much stronger. The "forgetting" becomes permanent and robust.

In Summary
The paper claims that to truly make an AI forget something, you can't just edit its high-level instructions; you must also clean its basic vocabulary cards. By using a precise mathematical "surgery" on these cards, EMBER ensures the AI forgets the specific concept without losing its ability to speak or remember other things, and it makes it nearly impossible for the AI to accidentally remember the forbidden topic again.

Technical Summary: Don't Forget Your Embeddings: Robust Knowledge Erasure via Precise Editing of Embeddings

Problem Statement

As language models (LMs) are increasingly deployed, the ability to permanently erase specific knowledge (e.g., private data, copyrighted content, or safety-sensitive topics) is critical. Existing persistent erasure methods primarily focus on updating the model's parameters, specifically within the Multi-Layer Perceptron (MLP) layers, often utilizing sparse autoencoders (SAEs) or weight steering. However, these methods frequently fail to achieve robust erasure. The target knowledge often remains recoverable through adversarial prompting or relearning (fine-tuning on concept-related data), and erasure efficacy is often form-dependent (failing to generalize from multiple-choice to open-ended generation).

The authors hypothesize that these limitations stem from existing techniques overlooking the token embedding layer. While MLPs are known to play a key role in knowledge recall, LMs also encode a nontrivial portion of their knowledge directly in their token embeddings. Ignoring this layer leaves a pathway for the model to "relearn" erased concepts and forces existing methods to rely on aggressive, broad parameter updates that degrade overall model utility.

Methodology: EMBER

The paper introduces EMBER (EMBedding ERasure), a plug-and-play module designed to precisely edit token embeddings to augment existing MLP-focused erasure methods. EMBER follows a "localize-then-edit" approach consisting of three stages:

Sparse Matrix Factorization (Sparse MF) for Feature Localization:
Instead of using SAEs, EMBER applies Sparse Matrix Factorization directly to the embedding matrix $E$ . It decomposes the embeddings of a specific set of tokens (those appearing in concept-related and neutral sentences) into a factor dictionary $Z$ (embedding features) and a coefficient matrix $Y$ .
- Sparsity Constraint: To ensure features correspond to localized, interpretable concepts, a "winner-takes-all" (WTA) constraint is applied to the coefficients $Y$ , keeping only the top 1% of magnitudes per factor. This forces the factorization to identify features shared by small, specific subsets of tokens.
Concept-Related Feature Identification:
The method scores the recovered features to distinguish concept-specific features from general ones. It computes a mass-ratio statistic comparing the contribution of a feature to concept-labeled tokens versus neutral tokens. Features with a high ratio are further filtered using an LLM-as-a-judge to verify that the associated tokens semantically match the target concept.
Precise Erasure via Subtraction:
Once concept-related features ( $F_C$ ) are identified, EMBER subtracts their contribution from the embeddings of the relevant tokens. For a token embedding $e_t$ , the update is:
$e_t \leftarrow e_t - \delta \sum_{i \in F_C} Y_{i,t} z_i$
where $\delta$ is a strength hyperparameter. This operation removes the concept-specific component while leaving the rest of the embedding intact. The method is applied only to tokens that have non-zero coefficients for the identified concept features.

EMBER is designed to be combined with existing MLP-based erasure methods (e.g., RMU, CRISP, SNMF) to create an ensemble that targets both the embedding layer and the MLP layers.

Key Contributions

Identification of the Embedding Bottleneck: The paper provides the first empirical study establishing the embedding layer as a critical bottleneck in robust concept erasure. It demonstrates that ignoring embeddings allows models to recover erased knowledge more easily.
EMBER Algorithm: A novel method using Sparse Matrix Factorization to disentangle and precisely remove concept-related features from token embeddings without requiring full model fine-tuning.
Plug-and-Play Augmentation: EMBER is shown to be compatible with existing erasure frameworks, significantly boosting their performance without requiring a complete redesign of the erasure pipeline.

Experimental Results

The authors evaluated EMBER on Gemma-2-2B-it and Llama-3.1-8B-Instruct across 18 diverse concepts (ranging from fictional works like Harry Potter to safety topics like Cannabis).

Efficacy and Specificity: Augmenting existing methods (RMU, CRISP, SNMF) with EMBER consistently improved erasure efficacy (lowering concept accuracy) and specificity (maintaining performance on similar domains and MMLU).
Robustness to Relearning: This is the primary gain. EMBER dramatically reduced the model's ability to relearn erased concepts.
- On Llama-3.1, relearning accuracy for prior methods (RMU, CRISP) was reduced from 70%–76% down to 35% when combined with EMBER (a ~50% reduction).
- On Gemma, relearning accuracy was reduced by 16%–24%.
- The full MF-based pipeline (SNMF + EMBER) achieved the lowest relearning accuracy (35% on Llama), significantly outperforming baselines.
Coherence and Localization: The cost to generation coherence was minimal. Analysis showed that EMBER modifies at most 0.14% of the vocabulary per concept. Coherence degradation was highly localized to "concept-exclusive" tokens (e.g., "Hogwarts" when erasing Harry Potter), while general tokens (e.g., "Harry" in the context of Prince Harry) remained largely unaffected.
Generalization: While hyperparameters were tuned on multiple-choice (MC) tasks, the improvements transferred effectively to open-ended (OE) question answering, suggesting the edits target conceptual knowledge rather than format-specific shortcuts.

Significance and Claims

The paper claims that precise, targeted editing of the embedding matrix is fundamental to achieving robust concept erasure. By shifting the focus beyond the MLP-centric view of prior work, EMBER addresses the "shallow unlearning" problem where knowledge is merely suppressed rather than removed.

The authors conclude that:

Embedding-level interventions are necessary to prevent relearning attacks.
Existing erasure methods can be significantly enhanced by augmenting them with embedding-level edits.
The trade-off between erasure efficacy and model utility (coherence) can be improved by localizing edits to specific concept-related features rather than applying broad weight updates.

The work establishes a new direction for model editing, highlighting that the embedding layer is a vital, yet previously underexplored, locus of knowledge storage that must be addressed for safety and compliance applications.