🤖 machine learning

Geometric Evolution Maps: Extracting Stable Concept Probes from Transformer Residual Streams

This paper introduces Geometric Evolution Maps (GEMs), a method that tracks the directional trajectory of concepts in transformer residual streams to identify stable "handoff layers" where representations settle, thereby extracting more reliable concept probes than traditional fixed-layer or peak-score approaches across diverse architectures.

Original authors: James Henry

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: James Henry

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 Transformer AI model as a massive, multi-story factory where raw ideas enter at the bottom and get processed layer by layer until they become a finished product at the top.

For a long time, researchers trying to understand what the AI is thinking (like whether it understands "anger," "truth," or "danger") had a simple rule: Just look at the very top floor. They assumed that by the time an idea reached the final layer, it was fully formed and stable.

This paper argues that rule is wrong. It's like trying to understand a recipe by only looking at the final plated dish, ignoring the fact that the ingredients were chopped, mixed, heated, and stirred in completely different ways on every single floor before they got there.

Here is the simple breakdown of what the paper discovered and proposed:

1. The Problem: Ideas Spin Like a Top

The authors found that when an AI "thinks" about a concept (like "sarcasm" or "threat"), the way it represents that idea in its internal memory spins around wildly as it moves up the factory floors.

The Analogy: Imagine a group of people trying to form a human pyramid. At the bottom, they are just grabbing each other's hands and spinning in circles, trying to find their balance. They are all over the place.
The Data: The paper measured this "spin." In most cases, the direction of the idea at the start of the assembly process was almost completely different (like pointing North vs. South) from where it ended up.
The Mistake: If you try to "probe" (detect) the idea while it's still spinning in the middle of the factory, you might catch it pointing in the wrong direction. You might think the AI is thinking about "danger" when it's actually just organizing the pieces to eventually think about danger.

2. The Solution: The "Handoff" Layer

The authors introduced a new method called Geometric Evolution Maps (GEMs). Instead of guessing which floor to look at, GEMs track the idea as it moves up the factory to find the exact moment it stops spinning and locks into place.

The Analogy: Think of a relay race. The runners (layers) pass the baton (the concept) back and forth. For a while, the baton is wobbling and being passed awkwardly. But then, there is a specific moment—the Handoff—where the runner finally catches the baton, steadies their grip, and starts running in a straight line.
The Discovery: The paper found that this "steady grip" happens at a specific floor, which they call the Handoff Layer. Once the idea passes this layer, it stops spinning and stays stable all the way to the top.
The Result: If you want to know what the AI is thinking, you shouldn't look at the very top (where it might have drifted) or the middle (where it's still spinning). You should look exactly at the Handoff Layer, where the idea has settled into its final, stable shape.

3. Testing the Theory

The researchers tested this on 23 different AI models (ranging from tiny ones to very large ones) and 17 different types of concepts (like "sarcasm," "urgency," "deception," etc.).

The Comparison: They compared their new "Handoff" method against the old "Peak" method (looking at the floor where the idea seemed loudest, even if it was still spinning).
The Outcome: The Handoff method was better 68% of the time. In many cases, it was significantly more accurate.
The Catch: It works best in larger models. In very small models, the "spin" is sometimes so chaotic or the factory is so short that the idea doesn't have time to settle properly before it hits the top.

4. Special Rules for Different Factory Types

The paper noticed that different types of AI factories behave differently:

Standard Factories (MHA): These almost always have a clear "Handoff" layer where the idea settles. The new method works great here.
Grouped Factories (GQA): These have a more complex internal structure. They sometimes settle earlier or behave differently, so the "Handoff" method is less dominant but still useful.
The "Near-Top" Rule: Sometimes, the idea settles so late that it's almost at the very top floor. If you try to check the idea there, you might accidentally mix it up with the final "packaging" step (getting ready to speak). The authors created a special rule to handle these rare cases, ensuring they don't get confused by the packaging process.

5. What This Means (and What It Doesn't)

What it proves: The paper proves that AI concepts are not static; they are dynamic processes that move and rotate before they stabilize. To understand them, you must find the moment they stop moving.
What it doesn't claim: The paper does not claim this makes the AI safer, or that we can now perfectly control AI behavior. It simply provides a better "microscope" to see where the AI's thoughts are actually stable.
One Failure: The paper admits one specific small model (gpt2) broke the rules. In this tiny factory, the "Handoff" happened so close to the top that the method got confused by the packaging process. This is a known limitation, not a flaw in the general theory.

Summary

Think of the AI's brain as a river. The old way of studying it was to look at the ocean where the river ends, assuming the water was calm. This paper says, "No, the water is churning and spinning all the way down."

The Geometric Evolution Map is like a sensor that floats down the river, waiting for the exact moment the water stops swirling and flows straight. Once it finds that calm spot (the Handoff Layer), it takes a snapshot. This snapshot is a much clearer, more accurate picture of what the AI is actually thinking than any snapshot taken at the top or bottom.

Technical Summary: Geometric Evolution Maps (GEMs)

Problem Statement

Linear probes extracted from transformer residual streams are only as reliable as the specific layer from which they are extracted. Current literature relies on three naive heuristics for layer selection: probing the final layer, probing the layer with the peak separation score, or probing a fixed middle-to-late layer.

The paper argues that these approaches fail to account for a fundamental structural reality: concept representations in transformer residual streams are not static. They emerge gradually, undergo substantial directional rotation during their assembly phase (within the Concept Allocation Zone or CAZ), and only settle into a stable direction at a characteristic "handoff layer" after the primary CAZ. Probing at the separation peak (which occurs inside the CAZ) or at the CAZ entry extracts a direction that has not yet stabilized, leading to unreliable concept probes.

Methodology: Geometric Evolution Maps (GEMs)

The paper introduces Geometric Evolution Maps (GEMs), a method to track the full directional trajectory of a concept through residual stream activations to identify the precise layer where the concept direction stabilizes.

Core Definitions

Concept Direction ( $u^{(l)}$ ): At each layer $l$ , the dominant concept direction is defined as the $L_2$ -normalized difference between the mean residual stream activations of positive and negative contrastive pairs (Difference-in-Means).
Angular Velocity ( $\omega^{(l)}$ ): The rate of directional change is calculated as $\omega^{(l)} = 1 - |u^{(l)} \cdot u^{(l-1)}|$ . High velocity indicates active rotation; low velocity indicates stability.
Concept Allocation Zone (CAZ): The layer range where a concept undergoes its primary separation event, detected via Fisher-normalized separation scores.
Handoff Layer ( $L_H$ ): Operationally defined as the first layer immediately following the CAZ end where the angular velocity drops below a threshold ( $\epsilon = 0.05$ ). This is the layer where the concept direction "locks in."
Settled Probe: The probe extracted is the unit vector $u^{(L_H)}$ at the handoff layer.

Adaptive Width Rule

When evaluating probes via ablation, the paper introduces an adaptive width rule ( $w$ ). The default is to ablate 3 consecutive layers starting at $L_H$ . However, for near-final-layer handoffs ( $L_H/N > 0.85$ ), ablating 3 layers risks extending into the unembedding-preparation zone, diluting the signal. For these cases, the rule sets $w=1$ to extract the direction exactly at the convergence point.

Key Contributions

Formalization of GEM: A tractable method to compute the handoff layer and extract the settled concept direction.
Empirical Validation of Rotation: Demonstrated that concept directions rotate substantially during assembly. Across 391 concept $\times$ model pairs (23 models, 17 concepts), the mean Entry-Exit Cosine (EEC) within CAZs is 0.233, with 93.9% of pairs showing EEC < 0.5. This confirms that probing at CAZ entry or the separation peak captures a direction orthogonal to the final settled state.
Superiority of Handoff Probes: GEM-extracted probes are at least as precise as peak-layer probes in 68.5% (268/391) of trials and strictly outperform in 66.2% (259/391). A Wilcoxon signed-rank test on per-model proportions confirms significance ( $p=0.010$ ; $p=0.0036$ excluding gpt2).
Architecture-Specific Findings:
- MHA vs. GQA: MHA models favor the handoff layer in 78.3% of trials, whereas GQA models favor it in only 47.1%.
- Scale: GEM structure exists even in 70M models, but probe quality improves significantly in the >500M range.
Direction-Specificity Control: Ablation of the specific concept direction suppresses separation at a median rate 377 $\times$ higher than random-direction ablation (45.6% vs. 0.24%), confirming the effect is carried by the concept direction itself, not just depth.
Depth-Matched Control: A control experiment comparing GEM against a non-settling layer at the same relative depth confirmed that the settling criterion contributes an additional +32.5 percentage points of separation suppression, ruling out depth-of-extraction as the sole driver.
Structured Failure Modes:
- gpt2: Identified as a genuine failure mode where ablation at the handoff layer (due to proximity to unembedding in a 12-layer model) pathologically increases separation.
- OPT-6.7b: Initially appeared as a failure in a smaller corpus (41% improvement) but achieved 76% improvement at $N=250$ , indicating the initial finding was corpus-size-dependent.

Results and Observations

Handoff Depth: Most concepts settle in the final 15–30% of model depth. Safety-relevant and epistemic concepts (e.g., threat_severity, certainty) hand off deepest (mean relative depth $\ge$ 0.86), while syntactic concepts (e.g., specificity, plurality) hand off earlier.
Relay Structure: Multi-node ablation experiments on pythia-6.9b reveal that concepts often exist in a causal relay. The deepest GEM node is the most causally concentrated intervention point, but shallow nodes are causally upstream (cross-disruption observed in 9/15 multi-node concepts).
Adaptive Width Efficacy: The near-final rule ( $w=1$ ) improves probe quality in 75.9% of triggered cases (60/79), with a mean gain of +7.44pp for those cases.

Significance and Claims

The paper claims that GEM provides the extraction methodology underlying the CAZ Validation and PRH papers in this research program. Its primary significance lies in:

Correcting Probe Extraction: Establishing that the "settled" direction at the handoff layer is the geometrically precise representation of a concept, superior to the transient directions found at separation peaks.
Mapping Assembly Dynamics: Providing a method to visualize and quantify the "rotation" of concepts, revealing that concepts are not static but undergo a dynamic assembly process before stabilizing.
Causal Characterization: Demonstrating that GEM identifies the most accessible points in a causal relay for intervention, though it does not claim to capture the entirety of concept encoding (which remains distributed).

The paper remains modest regarding its scope: it does not claim to explain the full mechanism of concept encoding (e.g., the specific contributions of attention vs. MLP sublayers) nor does it claim to capture all concept information, only the "peaks" of the distribution where the signal is large enough to be isolated and tested. It explicitly notes that behavioral ablation (measuring downstream output changes) is designed but not yet run at scale.