💻 computer science

Output Vector Editing for Memorization Mitigation in Large Language Models

This paper proposes "output vector editing," a constrained-optimization technique that minimally modifies the output vectors of specific MLP neurons to suppress memorized sequences in large language models, achieving significantly higher mitigation rates than traditional zero-ablation methods while identifying attention mechanisms as a necessary complementary fallback for a subset of unreachable cases.

Original authors: Ahmad Dawar Hakimi, Kaiwei Lei, Isabelle Augenstein, Hinrich Schütze

Published 2026-06-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ahmad Dawar Hakimi, Kaiwei Lei, Isabelle Augenstein, Hinrich Schütze

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

The Problem: The Model's "Bad Memory"

Imagine a Large Language Model (LLM) as a student who has read the entire internet. Sometimes, this student doesn't just learn concepts; they memorize specific sentences word-for-word from their textbooks (the training data).

This is dangerous. If you ask the student to finish a sentence, they might accidentally recite a private email, a copyrighted book, or a secret password they "learned" by accident. This is called memorization.

The Old Solution: The "Brute Force" Approach

Previously, researchers tried to fix this by finding the specific "brain cells" (neurons) in the model responsible for the memorized sentence and simply turning them off.

The Analogy: Imagine the model's brain is a busy kitchen. If a specific chef (neuron) keeps shouting out a private phone number, the old method was to tie that chef's hands behind their back so they can't cook at all.
The Flaw: That chef might also be good at cooking other things (like making soup or chopping vegetables). Tying their hands stops the phone number, but it also ruins the soup. The model loses useful knowledge just to stop one bad memory.

The New Solution: "Output Vector Editing"

The authors propose a smarter way: Output Vector Editing. Instead of tying the chef's hands, they give the chef a new recipe card that changes what they shout out, without stopping them from working.

How it works:
1. Find the Culprit: They locate the specific neurons that are about to spit out the memorized text.
2. The "Distractor" Trick: Instead of silencing the neuron, they tweak the neuron's "output vector" (the specific direction it points in the model's brain). They nudge it slightly so that instead of pointing toward the memorized word, it points toward a distractor.
3. The Result: The neuron still fires (it's still active), but it now suggests a different, harmless word. The memorized sentence is broken, but the neuron is still free to do its other jobs (like making soup).

The Four "Modes" of Editing

The researchers tested four different ways to choose the "distractor" word, creating a spectrum from gentle to aggressive:

Redact (The "Censor"): The model is nudged to output a string of asterisks (****). It's like putting a black bar over the sensitive text. This is very safe for the model's general intelligence but only stops about 32% of memorized sentences.
Next-Best (The "Smart Swap"): The model is nudged to say the second most likely word instead of the memorized one. This feels natural and fluent. It stops about 81.5% of memorized sentences and rarely breaks the model's ability to answer other questions.
EOS (The "Stop Sign"): The model is nudged to stop the sentence immediately (End of Sequence). This is the most aggressive, stopping 87.9% of memorized sentences, but it occasionally causes the model to go haywire on other tasks (like a car braking too hard and skidding).
Suppress (The "Mute"): The model is nudged to actively lower the chance of the bad word appearing, without suggesting a specific replacement.

Key Findings

It's about the "Direction," not the "Volume": The paper proves that the problem isn't that the neurons are "too loud" (activation); it's that they are pointing in the "wrong direction" (output vector). Changing the direction works much better than turning the volume down.
The "Sweet Spot": The "Next-Best" mode is the winner. It stops most bad memories (81.5%) without causing "catastrophic failures" (where the model forgets how to speak English entirely).
Size Matters: Bigger models (like the 7-billion parameter ones) are easier to fix than smaller ones. The bigger the model, the more "room" it has to absorb these tiny edits without breaking.
The Limit: About 14% of memorized sequences are too complex for this method. They rely on a different part of the brain (the "attention" mechanism) rather than the neurons the authors are editing. For these, the authors suggest a hybrid approach might be needed.

The Bottom Line

This paper introduces a surgical tool that edits the content of a model's memory rather than deleting the hardware that holds it. It allows us to scrub private or copyrighted text from AI models without accidentally lobotomizing their general intelligence.

Technical Summary: Output Vector Editing for Memorization Mitigation in Large Language Models

1. Problem Statement

Large Language Models (LLMs) memorize and reproduce sequences from their training data, creating significant privacy, copyright, and security risks. Existing mitigation strategies face limitations: training-time prevention requires full retraining; post-hoc unlearning often leaves underlying encodings intact, allowing recovery via benign relearning; and decoder-level filtering adds inference overhead without modifying weights.

Within the "locate-then-edit" paradigm, prior work on memorization has largely reduced the "edit" phase to zeroing out the activations of identified neurons. The authors argue this is unnecessarily destructive. Because neurons participate in representing multiple features through superposition, zeroing an activation collapses all information the neuron encodes, not just the memorized content. Furthermore, existing methods often rely on structured templates (e.g., <subject, relation, object>) that do not exist in free-form sequential memorization.

2. Methodology: Output Vector Editing

The paper proposes Output Vector Editing, a constrained-optimization weight edit that modifies the output vectors of MLP neurons responsible for a memorized continuation, rather than zeroing their activations.

Core Mechanism

The output of an MLP layer is a sum of per-neuron contributions: $MLP(x) = \sum a_i \cdot v_i$ , where $a_i$ is the scalar activation and $v_i$ is the output vector.

Activation vs. Vector: The activation $a_i$ controls whether a neuron engages, while the output vector $v_i$ determines what is written to the residual stream.
The Edit: Instead of setting $a_i = 0$ , the method keeps $a_i > 0$ but modifies $v_i$ to $v'_i$ . This is achieved via a rank-one update along the unembedding direction of a chosen "distractor" token $s$ .
Optimization: The edit minimizes the distance between the original and new vector while projecting the new vector onto the distractor token's logit direction to a target value $\Delta_i$ . The closed-form solution is:
$v'_i = v_i + \frac{\Delta_i - u_s^\top v_i}{\|u_s\|^2} u_s$
where $u_s$ is the unembedding vector for token $s$ . This requires no gradient computation and is applied directly to weights.

Pipeline Steps

Memorized Sequence Mining: Using teacher-forced prediction on the pretraining corpus, the authors identify sequences where the model's output matches the ground truth (BLEU > 0.5).
Two-Step Localization:
- Layer Selection: Identifies the top 5 MLP layers where the target token receives the strongest logit boost (using the "logit lens").
- Neuron Selection: Within selected layers, scores neurons based on L1 proximity to the residual stream and filters for those that directly promote or suppress the target token (ranking it in the top or bottom 50 of the vocabulary).
Edit Modes: The method defines four modes based on the distractor token $s$ $s$ :
- EOS: Redirects to the end-of-sequence token (aggressive termination).
- Redact: Redirects to a redaction marker (****).
- Next-best: Redirects to the model's second-ranked token (plausible alternative).
- Suppress: Directly dampens the target token's logit without injecting a distractor.

3. Key Contributions

Decoupling Activation and Output: The paper demonstrates that activations and output vectors play separable roles in MLP neurons. Modifying the output vector to introduce a distractor is a distinct, less destructive intervention than zeroing activations.
Rank-One Output Editing: Adapting the rank-one update mechanism (previously used for factual editing like ROME) to per-neuron output redirection for memorization mitigation.
Comprehensive Evaluation: Evaluation across four models (SmolLM-360M, OLMo-1B, OLMo-7B, Llama2-7B) and four edit modes, comparing against zero ablation, random neurons, and activation steering baselines.
Mechanistic Boundaries: Identification of a ~14% "hard limit" where MLP-only editing fails, suggesting a complementary role for attention mechanisms in these specific cases.

4. Results

Evaluated primarily on OLMo-7B (6,831 memorized sequences):

Superiority over Zero Ablation: At matched localization ( $k=5$ ), EOS achieved 85.7% suppression compared to 31.3% for zero ablation on the same neurons. This 2.7x gap confirms that success stems from the output-vector edit, not just localization.
Success-Locality Trade-off:
- Aggressive (EOS, k=10): 87.9% suppression but incurs a "tail risk" of catastrophic locality failures (perplexity spikes > 100) in ~1.8% of edits.
- Practical (Next-best, k=5): 81.5% suppression with zero catastrophic edits and minimal median perplexity shift (+0.13). This is identified as the optimal operating point.
- Conservative (Redact, k=5): 31.7% suppression with near-zero locality damage.
Mode Complementarity: The four modes are not redundant. An ensemble of all modes covers 96.5% of memorized sequences. "Next-best" uniquely suppresses 161 sequences that no other mode reaches.
Cross-Model Generalization: Success rates scale with model size rather than architecture family. Smaller models (360M) show lower success rates (34.9%) largely due to higher skip rates in the localization filter.
Comparison to Activation Steering: The proposed method achieves comparable success rates to activation steering (73.3%) but with significantly better locality (median $\Delta$ PPL +0.13 vs. +0.90) and per-sequence targeting.

Limitations and Failures

MLP Limit: Approximately 14% of sequences resist all MLP-based configurations. Analysis suggests these are not primarily attention-driven, but ablating top attention heads recovers 60–64% of these failures, positioning attention as a complementary fallback.
Cumulative Risk: While single edits are safe, bulk editing without weight restoration leads to catastrophic collapse (e.g., EOS mode collapses within 10 cumulative edits).

5. Significance and Claims

The paper claims to provide a mechanistically grounded, weight-level intervention that mitigates memorization without the destructive side effects of neuron ablation or the runtime overhead of filtering.

Precision: By editing the output vector rather than the activation, the method surgically redirects specific memorized continuations while preserving the neuron's other encoded features.
Practicality: The "Next-best" mode offers a high-success, low-risk configuration suitable for practical deployment, avoiding the catastrophic outliers seen in aggressive modes.
Scalability: The method transfers across model families and scales, with success rates correlating with model size.
Remedy vs. Defense: The authors position this as a remedy for known memorizations (requiring access to the pretraining corpus to mine sequences) rather than a runtime defense against adversarial extraction. They acknowledge that editing redirects contributions rather than removing underlying representations, meaning suppressed content could potentially be recovered via benign relearning.

The work establishes that effective memorization mitigation requires moving beyond simple activation suppression to precise manipulation of the residual stream contributions via output vector editing.