🤖 AI

Probe-Geometry Alignment: Erasing the Cross-Sequence Memorization Signature Below Chance

This article introduces Probe-Geometry Alignment (PGA), a surgical intervention that aligns model activities such that feature discrimination for cross-sequence memorization in large language models is reduced below chance level while their functional capabilities are preserved.

Original authors: Anamika Paul Rupa, Anietie Andy

Published 2026-05-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Anamika Paul Rupa, Anietie Andy

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 you own a library of books (a Large Language Model) that has memorized a specific secret story. You ask the librarian to "unlearn" this story, meaning they should no longer tell it to anyone.

Most current methods for "unlearning" are like telling the librarian: "If someone asks for this story, simply say 'I don't know' or invent a different ending." The librarian obeys and stops telling the story. However, the paper argues that the story is still written in the librarian's brain; they have merely learned to hide it. If you ask the right tricky questions, the librarian might accidentally reveal that they still know it.

This paper introduces a method to determine whether the story has truly vanished from the librarian's brain, as well as a new method to actually delete it without causing the librarian to forget how to do their job.

The Problem: The "Ghost" in the Machine

The authors discovered that a model, even after it stops telling a memorized secret, still internally knows it. They refer to this as a "Cross-Sequence Signature."

The Analogy:
Imagine the librarian has a hidden "Yes/No" switch in their brain that lights up as soon as they think about the secret story.

Old Unlearning: You train the librarian to keep their mouth shut. They stop telling the story.
The Reality: The hidden "Yes/No" switch still lights up brightly when you ask about the story. The knowledge is still there, just suppressed.

The authors built a special test (a "probe") to check if this switch lights up. They found that this "ghost" of memory exists in models of all sizes, from tiny toy models to massive ones like Mistral-7B.

The Discovery: Memory and Speech Are Separate

One of the paper's biggest insights is that remembering and speaking occur in different parts of the brain.

The Analogy:
Imagine the model as a radio station.

Storage: The secret is stored in the "recording studio" (the deep layers of the model).
Broadcasting: The "On-Air" switch (the Attention Heads) decides whether the recording is played.

The authors showed that you can destroy the "On-Air" switch so the secret is never broadcast (the model stops saying it). However, the recording in the studio remains perfectly clear and intact. You can even point to the recording and say, "That is the secret!" even though the radio is silent.

The Solution: "Probe-Geometry Alignment" (PGA)

Since old methods only destroyed the "On-Air" switch, the authors invented a new surgical tool called Probe-Geometry Alignment (PGA).

The Analogy:
Instead of just destroying the microphone, PGA goes into the recording studio and aligns the sound waves.

Find the Signal: First, they use their special test to find the exact direction in the brain where the secret hides.
Surgical Alignment: Next, they perform a tiny, precise adjustment in every layer of the model. They do not delete the whole brain; they simply shift the specific "direction" where the secret lives so that it no longer looks like a secret. It is like turning a clear, high-resolution photo into static noise only in the specific area where the secret was located, while the rest of the photo (the model's general knowledge) remains perfectly sharp.

The Results:

The Ghost is Gone: After applying PGA, the special test no longer lights up. In fact, the test performs worse than random guessing, meaning the model has truly forgotten the internal structure of the secret.
No Side Effects: Crucially, this operation did not prevent the librarian from doing anything else. Their ability to answer general questions, write stories, or solve logic puzzles remained exactly the same.

Key Takeaways in Simple Language

Silence is Not Forgetting: Just because a model stops saying a secret does not mean it has forgotten it. The memory is still hiding inside.
We Can See the Hiding Place: The authors developed a way to detect these hidden memories across models of different sizes.
We Can Delete Them: They developed a method (PGA) that surgically removes these hidden memories.
It Is Safe: This deletion is so precise that it does not damage the model's general intelligence. It is like removing a specific stain from a white shirt without the shirt shrinking or changing color.

The paper concludes that to truly "unlearn" something from an AI, you must delete the internal representation, not just silence the output. Their new method, PGA, does exactly that.

Technical Summary: Aligning Probe Geometry for the Erasure of Memorization Signatures

1. Problem Statement

Current frameworks for machine unlearning in Large Language Models (LLMs) suffer from a critical measurement gap: they evaluate behavioral forgetting (whether the model stops generating target content) but fail to assess representational forgetting (whether the model no longer internally encodes this content). Recent studies suggest that behavioral suppression often leaves internal traces that can be recovered via adversarial probes. This article argues that suppressing generation for privacy or security is insufficient, as underlying memorization signatures persist in the model's residual stream and can potentially be exploited by attackers.

The core challenge is twofold:

Detection: Existing probes often fail to distinguish between genuine memorization signatures across sequences and superficial string artifacts or model layer shifts.
Erasure: Surgical interventions that successfully suppress behavioral retrieval (e.g., head-level unlearning) often fail to erase the underlying representational signature, which remains linearly separable and detectable.

2. Methodology

2.1 Cross-Sequence Leave-One-Out (LOO) Probing

To rigorously detect memorization, the authors introduce a Cross-Sequence LOO probing protocol. Unlike single-sequence probes that may overfit to specific token identities, this protocol tests for a generalizable memorization signature:

Procedure: A linear probe is trained on activation pairs (memorized vs. control) from a set of sequences, with one memorized sequence held out. The probe is then tested on the held-out sequence.
Significance: If the probe classifies the held-out sequence as memorized, it has learned a feature of the memorization process itself rather than just specific training examples.
Controls: The protocol includes a "Pure-Distinguishability" baseline (trained on non-memorized sequences with adjusted lexical structure) and a "Shuffled-Label" null value to isolate the memorization-specific gap.

2.2 Causal Separability Analysis

The article investigates whether the direction identified by the probe is causally responsible for retrieval. Using causal tracing and Forward Hooks on Pythia-70M, the authors project the fitted probe direction out of the residual stream.

Result: This intervention collapses the local memorization signature (the gap drops from +0.44 to -0.19), while behavioral retrieval (log-probability of the secret) and general capabilities remain largely intact. This proves that the probe-able signature and the retrieval-producing mechanism occupy causally separable directions.

2.3 Probe-Geometry Alignment (PGA)

To address the failure of existing unlearning methods (which suppress outputs but leave representations intact), the authors propose Probe-Geometry Alignment (PGA).

Mechanism: PGA is a surgical erasure technique that aligns activations specifically along the live readout direction of the probe at every depth.
Goal: It minimizes the scalar projection of the difference between memorized and clean activations onto the probe's weight vector ( $\hat{w}_d$ ) at each depth $d$ , while preserving cross-entropy on clean data to maintain capabilities.
Constraint: Unlike isotropic alignment methods (e.g., AAE) that act on all $d_{model}$ dimensions, PGA acts on a single scalar magnitude per depth (the projection along the probe's readout direction).
Adversarial Extension: To defeat attackers who re-fit probes on PGA-treated activations, the authors introduce Adversarial PGA, which iteratively expands the projection basis with orthogonal components of re-fitted probes.

3. Main Contributions

Cross-Sequence LOO Protocol: A validated method for detecting genuine memorization signatures that generalize across held-out sequences and distinguish them from model layer shifts or representational drift.
Causal Separability: Empirical evidence that the linear direction readable by a probe differs from the direction required for behavioral retrieval. Projecting out the probe direction collapses the signature without significantly impairing retrieval.
Distinct Memorization Regimes: Demonstration that naturally pre-trained memorization and fine-tuning-injected secrets leave representationally distinct traces. A probe trained on natural memorization does not classify injected secrets as memorized.
Probe-Geometry Alignment (PGA): A constructive surgical erasure method that drives cross-sequence probe accuracy below random chance (<0.50) across four model scales (from 0.8M toy model to 7.24B Mistral-7B) while preserving zero-shot capabilities.

4. Key Results

4.1 Existence of the Signature

The cross-sequence memorization signature is real and consistent across scales:

Pythia-70M: Memorization-specific gap of +0.32 (Peak +0.54 at L6).
GPT-2 Medium: Gap of +0.19 (Peak +0.45 at L21).
Mistral-7B: Gap of +0.30 (Peak +0.47 at L11).
Cluster Specificity: The signature is strong for formal registers of English and licenses, but nearly zero for code and pseudo-Latin.

4.2 Failure of Standardized Unlearning Methods

Standard unlearning methods (Gradient Ascent, NPO, RMU, MEMIT, and the authors' own MLDU) successfully suppress behavioral retrieval (e.g., $P(secret) < 10^{-4}$ ) but fail to collapse the probe. In toy models and at scale, these methods leave the linear probe's accuracy at ceiling (1.000) or near ceiling, confirming that behavioral suppression does not equate to representational erasure.

4.3 Success of PGA

PGA successfully collapses the signature below random chance (<0.50) across all tested scales:

Toy Model (0.8M): The probe drops to 0.17 at depth 4.
Pythia-70M: The probe drops to 0.11 ± 0.04 at L6.
Mistral-7B: The probe drops to 0.42 in middle layers.
GPT-2 Medium: Using MD-PGA (Eigenbasis variant), the probe drops to 0.061 at L21.
Robustness: The collapse holds against six adversarial probe variants (different seeds, regularizations, and MLPs). Adversarial PGA defeats re-fitted attacker probes at all memorization-critical depths.

4.4 Preservation of Capabilities

PGA preserves model capabilities with negligible degradation:

Across five zero-shot benchmarks (HellaSwag, PIQA, BoolQ, ARC-Easy, WinoGrande), the mean $\Delta$ -accuracy is +0.2pp.
The maximum single-task regression is 2.9pp (on BoolQ for the adversarial variant).
Perplexity (PPL) remains stable (e.g., 1.40 $\to$ 1.42 for the toy model).

5. Significance and Claims

The article claims that suppressing what a model says is not the same as erasing what it represents. The cross-sequence signature is a real, causally separable, regime-specific property of pre-trained representations.

Auditability: The cross-sequence probe enables representational verification post-unlearning, allowing third parties to audit whether a model truly forgets or merely suppresses outputs.
Constructive Erasure: PGA offers a concrete method to surgically remove these signatures (one scalar per depth) without retraining from scratch or sacrificing general capabilities.
Mechanistic Insight: The work proposes a decomposition of memorization in Transformers where information is encoded in the residual stream (storage) and routed to the output via attention heads (expression). Unlearning methods that only disrupt routing fail to erase the storage.

The authors position this work not as a complete solution for all unlearning scenarios, but as a first step toward empirically auditable privacy in post-hoc unlearning, reframing the field from a binary behavioral pass/fail to a constructive audit pipeline.