🤖 machine learning

Geometric and Behavioral Stratification in Transformer Residual Streams

This paper identifies the prediction direction as a privileged, content-defined anchor in transformer residual streams that geometrically and behaviorally stratifies high-dimensional computation into a narrow, structured prediction-proximal interface and a vast, scale-dependent complement, revealing how linear readout coexists with complex internal representations.

Original authors: Nelson Guda

Published 2026-08-14

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Nelson Guda

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 are trying to understand how a giant, super-smart robot brain works. This brain is built from a type of technology called a "transformer," which is the engine behind many of the AI chatbots and writing tools we use today. These robots don't think like humans; instead, they process information through a massive, high-dimensional "residual stream." Think of this stream as a giant, swirling river of numbers that flows through the robot's layers. At every step, the robot adds new information to this river, and at the very end, it has to decide what word to say next.

For a long time, scientists thought that if you made the robot bigger (by adding more parameters or making the river wider), the way it decided on words would just get bigger and more complex in the same way. They assumed the "decision-making" part of the brain would expand to fill all the new space. But this new research asks a different question: What if the robot is actually keeping its decision-making channel narrow and fixed, while using the extra space for something else entirely? The author is essentially trying to map the geography of this river to see where the "real" thinking happens and where the robot just stores its notes. They want to know if the robot's internal map is organized randomly or if it follows a hidden, special rule.

The Paper's Big Discovery: The "Prediction Anchor"

In this study, researchers Nelson Guda and colleagues looked at 18 different AI models, ranging from small ones (7 billion parameters) to massive ones (120 billion parameters). They wanted to see how these models organize their internal "river" of data. They discovered that the models don't just float around randomly; they organize everything around a very specific, moving target: the prediction direction.

Imagine the robot is trying to guess the next word in a sentence. It has a "target" in its mind—the specific direction in its brain that points to that word. The researchers found that the robot treats this target as a special anchor. Everything in its brain is measured relative to this anchor.

The "Narrow Door" vs. The "Huge Warehouse"

The most surprising finding is that the part of the brain that actually decides the next word is incredibly small and stays the same size, no matter how big the robot gets. The author calls this the prediction interface.

The Analogy: Imagine a massive, multi-story warehouse (the model's brain). Inside this warehouse, there is a tiny, narrow door (the prediction interface) that leads to the outside world (the next word).
The Finding: Whether the warehouse is the size of a house or the size of a city, that door stays exactly the same size. The researchers found that this "door" only uses about 4 to 10 dimensions (a tiny slice of the brain's total space).
The Rest of the Brain: The rest of the warehouse—the "prediction-distal complement"—is huge. As the models get bigger, this warehouse gets bigger and bigger, but the door does not. The extra space isn't used to make the door wider; it's used to build a massive storage area around the door.

Why This Matters: The "Signal vs. Noise" Game

Why would a robot keep such a tiny door? The paper suggests it's about keeping the signal clear. If the door were wide and messy, all the noise from the huge warehouse would get mixed up with the decision. By keeping the door narrow and the warehouse huge, the robot can do all its complex calculations in the warehouse, but it only lets the clean, clear "signal" pass through the tiny door to make its choice.

The researchers found that the "warehouse" (the complement) is actually doing a lot of heavy lifting. Even though it's far away from the decision door, if you mess with the direction of the data in the warehouse, the robot's output goes crazy immediately. However, if you just make the data in the warehouse smaller (without changing its direction), the robot is fine. This suggests the robot cares about where the data points, not how big it is.

The "Anti-Clumping" Mystery

Here is a weird twist the researchers found. When they looked at how the robot groups similar prompts (like questions about math vs. questions about history), they expected the "warehouse" to be messy and unorganized. Instead, they found something strange: the warehouse actually does the opposite of what you'd expect.

The Door (Prediction Interface): This part is very good at telling different groups of prompts apart. It keeps math questions far away from history questions.
The Warehouse (Complement): This part is "anti-discriminatory." It actually pushes prompts from the same group further apart and pulls prompts from different groups closer together.

Think of it like a party. The "door" is the bouncer who checks IDs and separates the VIPs from the regulars. The "warehouse" is the dance floor. On the dance floor, the VIPs might be dancing in a chaotic circle, while the regulars are huddled together. The warehouse seems to be doing a special kind of balancing act to make sure that even if the "door" sees different things, the final result (the next word) still makes sense.

The "Task Frame" vs. The "Storyline"

The researchers also tested what happens when they "scrambled" different parts of the brain. They found a clear hierarchy in how the robot behaves:

The "Stance" Layer (Closest to the door): If you mess with the part of the brain closest to the decision door, the robot immediately forgets what it's supposed to be doing. It might stop writing a story and start giving grammar advice, or switch from being a helpful assistant to a grumpy critic. It changes the frame of the conversation instantly.
The "Trajectory" Layer (A bit further out): If you mess with the next layer, the robot keeps the same frame (it stays a helpful assistant), but it changes the content of the story. It might start talking about a cat instead of a dog, but it's still telling a story.
The "Dynamic" Layer (The moving part): The researchers also found that the huge warehouse has a "static" part (like a permanent scaffold) and a "dynamic" part (like a moving envelope that updates every second). If you freeze the dynamic part and don't let it update, the robot crashes immediately. It needs that constant update to keep generating text.

What This Means for the Future

The paper suggests that as AI models get bigger and smarter, they aren't getting "wider" in their decision-making. Instead, they are getting better at packing more complex calculations into the massive warehouse surrounding that tiny, fixed door.

This changes how we might evaluate these models. If we only look at the final answer (what comes out of the door), we might miss 99% of what the model is actually doing inside the warehouse. The model could be doing incredibly complex reasoning in that huge, hidden space, but because the door is so narrow, we can't see it. The author suggests that to truly understand AI, we need to look inside the warehouse, not just at the door.

In short, the paper reveals that these giant AI brains are organized like a fortress with a tiny, unchanging gate. The magic happens in the vast, hidden courtyard behind the gate, where the real work is done, and the gate itself just ensures the final message comes out clean and clear.

Technical Summary: Geometric and Behavioral Stratification in Transformer Residual Streams

Problem Statement

Recent research has established that trained transformer models develop "privileged bases"—specific coordinate axes with statistical properties distinct from the rest of the residual stream. However, the functional role and geometric nature of these privileged directions remain unclear. While prior work identified these axes as anomalous coordinate sets, it did not explain what specific direction they select or how they organize the high-dimensional computation of the model.

This paper addresses the gap between the high-dimensional nature of transformer computation (thousands of dimensions) and the low-dimensional nature of the linear readout (a single projection onto vocabulary logits). The author investigates whether the prediction direction—the unembedding vector of the token the model currently predicts—serves as a content-defined privileged anchor that organizes the residual stream's geometry and behavior.

Methodology

The author introduces PDSF (Prediction-anchored Decomposition of Subspace Features), a coordinate system that partitions the residual stream state ( $h$ ) at every layer into four strictly orthogonal components, ordered by their proximity to the model's current prediction direction ( $\hat{u}$ ):

P (Predictive): The rank-1 projection of $h$ onto the normalized unembedding vector $\hat{u}$ of the current argmax token.
D (Discriminative): The top- $k$ PCA directions of the residual after removing $P$ .
S (Situational): The top- $k$ PCA directions of the residual after removing both $P$ and $D$ .
F (Framework/Complement): The orthogonal remainder, representing the vast majority of the residual stream dimensions.

The combined region $P + D + S$ is termed the prediction interface, while $F$ is the prediction-distal complement.

Experimental Setup:

Testbed: 18 transformer models across 6 architecture families (Llama, Gemma, Mistral, Mixtral, Qwen, GPT-OSS), spanning 7B to 120B parameters, including both base and instruction-tuned variants.
Prompt Sets: Three structured sets designed to probe different variations:
- SpecA: Constrained prompts with fixed answers but varied surface forms (224 prompts).
- SpecB: Open-ended narrative continuations (80 prompts).
- Diverse: Prompts across 8 linguistic regimes (narrative, code, math, etc.) (84 prompts).
Interventions: The author employs targeted interventions, including single-pass perturbations, persistent rotational scrambling (Haar measure rotations applied at ~70% depth), and magnitude attenuation, to test causal roles.

Key Contributions and Results

1. A Scale-Invariant Prediction Interface

The effective dimensionality of the prediction-proximal subspaces ( $D$ and $S$ ) is low and scale-invariant.

On constrained prompts, the effective rank of $D$ is consistently single-digit (mean $\approx 4$ ) across all models, regardless of hidden dimension ( $d_{model}$ ranging from 2,880 to 8,192).
The dimensionality of the interface is determined by the semantic complexity of the prompt, not the model's width.
Instruction tuning sharpens the internal organization of the interface (increasing discrimination and reducing paraphrase sensitivity) but does not widen it.
Conclusion: Scaling increases the size of the prediction-distal complement ( $F$ ), not the prediction interface.

2. Geometric Stratification: Orthogonality and Manifold Complexity

Orthogonality: The prediction direction is nearly orthogonal (mean angle $\approx 84^\circ$ – $86^\circ$ ) to the principal variance axes of the residual stream. High-variance directions are displaced away from the readout axis.
Manifold Complexity Gradient: There is a steep gradient in manifold complexity along the prediction-proximity axis.
- Prediction-proximal ( $D$ ): Highly folded, structured geometry (high information density).
- Prediction-distal ( $F$ ): Flat, uniform geometry.
- This gradient is significantly steeper when measured relative to the prediction direction compared to standard variance-ordered PCA controls, especially on linguistically heterogeneous prompts.

3. Anti-Discrimination in the Complement

Contrary to the expectation that group discrimination fades with distance from the prediction direction, the complement ( $F$ ) exhibits anti-discrimination.

Prompts from the same semantic group are systematically farther apart in $F$ -space than prompts from different groups.
This suggests $F$ carries prompt-specific configuration that compensates for variations in the prediction-proximal states to achieve similar readout outcomes.

4. Direction Over Magnitude (Thin-Shell Hypothesis)

The paper validates the "thin-shell" hypothesis: residual-stream variation is driven by direction, not magnitude.

Disrupting the direction of the complement ( $F$ ) via rotation causes immediate and severe output divergence.
Reducing the magnitude of $F$ (attenuation) while preserving direction has a negligible effect.
This confirms that the complement acts as a directional substrate essential for computation, despite its weak alignment with the readout.

5. Behavioral Stratification

Persistent rotational interventions reveal a behavioral hierarchy aligned with the prediction-proximity axis:

Subspace $D$ (Immediate Framing): Disruption causes immediate divergence (mean first-divergence token $\approx 1.2$ ) and frequent task-frame shifts (e.g., changing from answering to meta-commentary).
Subspaces $S$ and $F$ -topK (Trajectory): Disruption causes delayed divergence (mean first-divergence token $\approx 5.6$ for $S$ , $\approx 6.5$ for $F$ -topK) and content substitution within a preserved task frame.
This indicates that $D$ controls the initial framing of the response, while $S$ and the high-variance portion of $F$ control the subsequent trajectory.

6. Static/Dynamic Partition of the Complement

The complement $F$ is further organized by temporal dynamics:

It separates into a static scaffold (temporal mean) and a dynamic envelope (token-to-token updates).
The dynamic envelope is causally load-bearing: injecting "stale" (one-step delayed) $F$ causes universal output divergence by the first token.
Rotating the directions of the dynamic envelope causes catastrophic failure (92.3% failure rate), distinct from the redirection caused by rotating $D$ or $S$ .

Significance and Claims

The paper claims to establish the prediction direction as a privileged, content-defined anchor that organizes residual-stream geometry in a way distinct from previously described coordinate-based axes.

Architectural Constraint: The findings suggest that the linear readout imposes a "signal-to-noise" constraint, forcing the model to isolate prediction-relevant structure into a narrow, scale-invariant interface while expanding the computational substrate ( $F$ ) orthogonally to it.
Evaluation Blindness: Because the majority of the residual stream (the complement) is weakly aligned with the readout, output-based evaluation metrics (logits, next-token accuracy) likely under-measure the internal computation and structural changes occurring in the model as it scales.
Interpretability: The geometric stratification offers a framework for feature-level interpretability, suggesting that task framing resides in $D$ while trajectory and state maintenance reside in $S$ and $F$ .
Geometric vs. Variance: The paper argues that geometric organization is not merely a shadow of variance structure; anchoring to the prediction direction reveals organizational principles (like the anti-discrimination of $F$ and the behavioral hierarchy) that variance-based analyses alone fail to capture.

The author concludes that transformer residual streams are best understood not as diffuse high-dimensional activations, but as stratified geometries featuring a fixed prediction interface embedded in a dynamically maintained, causally load-bearing substrate.