Escaping Flatland: Weight-Space Regularization and Hyperbolic Geometry in Mechanistic Interpretability
Abstract
The frontier of AI alignment and safety currently depends on Mechanistic Interpretability—the attempt to reverse-engineer the dense, superimposed activations of Large Language Models (LLMs) into discrete, human-understandable concepts. The foundational engine of this field is the Euclidean Sparse Autoencoder (SAE). Yet, as researchers attempt to scale SAEs to tens of millions of features, they hit a structural mathematical wall: massive dictionary sizes yield diminishing returns in feature purity, plagued by dead latents, feature-splitting, and explosive compute overhead.
This bottleneck is a geometric pathology. Current SAE architectures operate under the implicit assumption that semantic latent space is flat (Euclidean). However, human semantic concepts are intrinsically hierarchical (ontological trees). Because the spatial volume of discrete trees grows exponentially with depth, embedding them into a Euclidean manifold—where volume grows strictly polynomially—inevitably forces disparate concepts into geometric interference, violating sphere-packing limits.
By shifting the interpretability framework from linear algebra to the differential geometry of Riemannian manifolds, we propose the Hyperbolic Sparse Autoencoder (HypSAE). Reconciling theoretical topology with the unyielding realities of GPU optimization—specifically the RMSNorm spherical trap, Algebraic Cancellation, and the Category Error of applying geometry to transient tokens—we execute a complete topological decoupling. By utilizing a strictly Euclidean forward pass and applying the Poincaré manifold exclusively as a Weight-Space Regularizer, we leverage learned radial scalars, Momentum Contrast, and Hyperbolic Entailment Cones to natively group concepts hierarchically. This provides a rigorous mathematical pathway to monosemanticity while preserving perfect causal interventional compatibility.
1. The Dimensionality Crisis in Mechanistic Interpretability
In modern transformer architectures, networks are forced to represent vastly more concepts than they have dimensional bandwidth. To survive this bottleneck, models utilize superposition—packing multiple unrelated concepts into the same active neurons via nearly orthogonal vectors in a high-dimensional continuous space [1].
To untangle this, the AI alignment community relies on Sparse Autoencoders (SAEs) [2]. A standard Euclidean SAE maps the dense residual stream () into a sparse, overcomplete dictionary (, where ), optimizing:
To achieve clean feature separation, labs must train SAEs with billions of features, creating computational costs that rival foundation model pre-training [3]. Throwing dimensions at a flat space is a brute-force software engineering hack, not a mathematical solution.
2. The Geometry of Thought: Trees vs. Flatland
Semantic concepts are strictly hierarchical. The concept of a Golden Retriever sits at the terminal leaf of a deeply nested ontological tree.
graph TD
E["Entity (Root)"]
E --> Ab["Abstract"]
E --> LO["Living Organism"]
LO --> Pl["Plant"]
LO --> An["Animal"]
An --> Ma["Mammal"]
Ma --> Ca["Canine"]
Ca --> Wo["Wolf"]
Ca --> GR["Golden Retriever"]
In a discrete tree with a branching factor , the number of nodes grows exponentially with depth : . However, in flat Euclidean space , the volume of a hypersphere grows only polynomially: .
This creates a fundamental physical impossibility. You cannot pack an exponentially growing number of concepts into a polynomially growing volume without crushing them together [4]. When Euclidean SAEs lack sufficient dimensional bandwidth, the hierarchy shatters. A single coherent high-level concept fractures into redundantly specific, microscopic latents—a pathology known as feature-splitting—rendering the dictionary an unstructured, uninterpretable cloud.
3. The Hyperbolic Manifold: Continuous Trees
To resolve this packing limit, we must abandon flat affine spaces and transition to Riemannian geometry—specifically, manifolds with constant negative curvature.
In Hyperbolic space (often modeled via the Poincaré disk), the geometry warps such that the circumference and volume of a sphere grow exponentially with its radius [5]:
Because hyperbolic space expands exponentially in the transverse (angular) direction, it provides mathematically limitless orthogonal room to pack highly specific leaf nodes without them overlapping.
graph LR
subgraph flat["EUCLIDEAN SAE — Flat Space"]
direction TB
FA["Feature A"] --- FB["Feature B"]
FB --- FC["Feature C"]
FC --- FD["Feature D"]
FA --- FC
FB --- FD
FA -.-|"Feature Splits"| FD
end
subgraph hyper["HYPERBOLIC SAE — Poincaré Disk"]
direction TB
R["Root (Abstract)"]
R --> C1["Category 1"]
R --> C2["Category 2"]
C1 --> L1["Leaf A"]
C1 --> L2["Leaf B"]
C2 --> L3["Leaf C"]
C2 --> L4["Leaf D"]
end
| Property | Euclidean SAE (Flat Space) | Hyperbolic SAE (Poincaré Disk) |
|---|---|---|
| Volume Growth | Polynomial | Exponential |
| Hierarchy Encoding | Flawed / Feature-splitting | Native Radial Depth |
| Concept Mapping | Dense, crowded grids | Native ontological trees |
4. The Physics of Optimization: Escaping the Execution Traps
If hyperbolic space is the perfect topological fit, why hasn’t it been natively adopted for SAEs? Because theoretical geometry routinely shatters against the unyielding bare-metal physics of unsupervised GPU optimization. A naive attempt to map the SAE’s reconstruction objective directly into hyperbolic space triggers mathematically fatal collapses:
- The Category Error (Tokens vs. Concepts): Applying hyperbolic geometry to the transient token activations () organizes random sequence snippets in a batch into a tree. Concepts live in the Weights (Parameters), not the Token Activations [6].
- The RMSNorm Spherical Trap: Modern frontier LLMs utilize RMSNorm. This strict constraint projects every target residual token onto a high-dimensional Euclidean hypersphere with a fixed magnitude. Attempting to extract a radially varying ontology by mapping into a Poincaré cone is mathematically incoherent; the mapping will uniformly crush all tokens against a thin, equidistant spherical shell [7].
- The Algebraic Cancellation (The Ghost Manifold): If one maps the residual stream to the Poincaré ball () and then immediately back to the tangent space () to apply an sparsity penalty, the functions cancel out (). The forward pass remains mathematically identical to a vanilla Euclidean SAE, wasting GPU FLOPs computing a trivial Identity function [8].
- The Collinear Rank-Collapse: If one attempts to build hierarchies by placing parent and leaf concepts on the exact same angular ray to save Euclidean dimensions, one violates the Linear Representation Hypothesis. Collinear dictionary atoms are linearly dependent, and the sparsity penalty will immediately prune one to zero [9]. By Gauss’s Lemma, radial geodesics are isometric to Euclidean lines; true exponential packing capacity exists purely in transverse (angular) deviations [10].
5. The Solution: Weight-Space Regularization
To survive the physical limits of hardware, the topology of the LLM, and linear algebra, we execute a total geometric decoupling.
We abandon mapping the transient token activations () to the manifold. Instead, we apply the Poincaré manifold exclusively to the Dictionary Weights ()—the concepts themselves—acting as an elite structural regularizer.
5.1 The Pure Euclidean Forward Pass (The Interventional Territory)
The forward pass must remain strictly Euclidean. This respects the RMSNorm of the target data , guarantees blindingly fast GPU inference, avoids Algebraic Cancellation, and perfectly preserves the ability to perform causal Manifold Steering downstream [11].
5.2 The Hyperbolic Concept Embeddings (The Unit-Norm Paradox)
The ontology lives natively in the columns of the Euclidean decoder weight matrix . Let be the -dimensional Euclidean dictionary atom for Feature .
To prevent the solver from collapsing the activations to zero while inflating dictionary magnitudes to infinity, standard SAEs strictly enforce a unit-norm constraint on the decoder columns: . However, mapping pure unit vectors directly via the exponential map would force every concept onto a single, equidistant radial shell in hyperbolic space, mathematically crushing the ontological tree onto a 2D sphere.
The Fix (Decoupled Radial Depth): To solve this Unit-Norm Trap, we decouple the interventional direction from the ontological depth. We introduce a learned 1D scalar parameter for each feature to explicitly encode its Hierarchical Depth.
We map these scaled Feature Weights into the -dimensional Poincaré ball via the Exponential Map at the origin [8]:
(Since , the norm cancels out of the denominator, yielding an elegant, numerically stable projection).
Now we are computing geometry in dimensions (which is computationally stable and actively utilizes the low-dimensional packing capacity), and we are evaluating the structural relationships between static concepts. remains a flat Euclidean unit-vector handling interventional direction, while independently governs ontological depth (Origin = Root, Boundary = Leaf). The structural topology is successfully decoupled from the Euclidean solver.
5.3 Asymmetric Entailment via Momentum Contrast (MoCo)
To construct the structural prior without hitting the Petabyte memory wall of tracking a dense feature co-activation graph, we implement a Momentum Contrast (MoCo) Queue [12].
When a batch of tokens causes Feature (Parent) and Feature (Child) to heavily co-activate across the Zipfian token distribution, we sample this relationship from the MoCo queue and pull their hyperbolic weight embeddings ( and ) into a Hyperbolic Entailment Cone [13]:
This asymmetric loss mathematically enforces that highly specific child concepts fall within the geometric “cone” cast outward from their broad parents.
(Where the ConePenalty utilizes the standard isometric invariant geodesic distance scaled by ):
5.4 The Tri-Partite HypSAE Loss
The final optimization objective seamlessly merges flat interventional mechanics with curved topological sorting:
6. The Architectural Payoff: Monosemanticity at Scale
By decoupling the geometry into the Weight-Space, the HypSAE functions as an elite Structural Regularizer.
Because dictates how features superimpose in the LLM’s flat residual stream, the optimizer is forced to organize the Euclidean dictionary weights into an ontological tree to minimize the structural penalty. Root concepts remain near the origin (); distinct child concepts fan out transversely toward the boundary.
- Drastic Dictionary Reduction: Exponential volume handles packing efficiency natively within the structural prior, requiring exponentially fewer dimensions () to achieve the exact same feature separation.
- Elimination of Feature-Splitting: The network is mathematically incentivized to group related concepts hierarchically rather than fracturing a single concept into redundant Euclidean vectors to escape MSE penalties.
- Flawless Interventional Compatibility: Because the forward pass () and decoder () remain perfectly flat matrix multiplications, identity cancellations are avoided, and causal linear steering is preserved flawlessly.
6.1 HypSAE vs. Prior Art
It is vital to contrast the HypSAE with recent attempts to resolve the ontological bottleneck, namely Hierarchical SAEs (Muchane et al., 2025) [14] and Matryoshka SAEs (M-SAE) [15]. These architectures attempt to simulate depth via discrete structural hacks—nesting smaller Euclidean sub-spaces inside larger ones.
The HypSAE abandons this discrete engineering for a continuous geometric solution. Furthermore, the mathematics of Hyperbolic Dictionary Learning have been rigorously established in fields demanding complex topological mapping—such as charting brain ventricles via hyperbolic Ricci flow in Alzheimer’s biomedical imaging [16]. By importing these mature Riemannian algorithms and restructuring the execution graph via Momentum Contrast, we achieve native hierarchical extraction perfectly aligned with both the LLM’s geometry and the physics of the GPU cluster.
7. Conclusion: Escaping Flatland
In 1884, Edwin A. Abbott published Flatland, a novella detailing the tragic inability of two-dimensional shapes to comprehend the existence of a three-dimensional world above them. The interpretability community is currently trapped in a similar dimensional delusion—attempting to decode the rich, exponentially branching reality of artificial cognition using the rigid, polynomial geometry of a flat Euclidean blackboard.
Euclidean space is structurally incapable of capturing the dense, hierarchical ontology of human knowledge without triggering catastrophic feature-splitting. However, pure differential geometry will mathematically shatter if blindly forced to reconstruct the RMSNorm hypersphere of modern foundation models or applied directly to transient token activations.
By executing Weight-Space Regularization and embracing Transverse Angular Squeezing via the learned scalar , we respect the immutable laws of linear representation while unleashing the infinite packing capacity of Riemannian manifolds. We unlock structural, monosemantic interpretability, organizing the chaotic dictionaries of frontier models into mathematically guaranteed ontological maps.
To solve the black box of AI, we do not need more parameters; we need a better geometry.
Next in this Series: In the upcoming article, we will translate this decoupled geometry into a numerically stable, highly-optimized open-source PyTorch library:
hypersae.
References
- [1] Elhage, N., et al. (2022). “Toy Models of Superposition.” Anthropic Transformer Circuits Thread.
- [2] Bricken, T., et al. (2023). “Towards Monosemanticity: Decomposing Language Models With Dictionary Learning.” Anthropic Research.
- [3] Gao, L., et al. (2024). “Scaling and evaluating sparse autoencoders.” OpenAI Research.
- [4] Gromov, M. (1987). “Hyperbolic groups.” Essays in Group Theory, 8, 75–263.
- [5] Nickel, M., & Kiela, D. (2017). “Poincaré Embeddings for Learning Hierarchical Representations.” NeurIPS, 30.
- [6] Olah, C., et al. (2020). “Zoom In: An Introduction to Circuits.” Distill.
- [7] Zhang, B., & Sennrich, R. (2019). “Root Mean Square Layer Normalization.” NeurIPS.
- [8] Ganea, O., Bécigneul, G., & Hofmann, T. (2018). “Hyperbolic Neural Networks.” NeurIPS, 31.
- [9] Tropp, J. A. (2004). “Greed is good: Algorithmic results for sparse approximation.” IEEE Transactions on Information Theory.
- [10] Lee, J. M. (2018). Introduction to Riemannian Manifolds. Springer.
- [11] Lin, Z., et al. (2026). “Manifold Steering Reveals the Shared Geometry of Neural Network Representation and Behavior.” arXiv preprint.
- [12] He, K., et al. (2020). “Momentum Contrast for Unsupervised Visual Representation Learning.” CVPR.
- [13] Ganea, O. E., Bécigneul, G., & Hofmann, T. (2018). “Hyperbolic Entailment Cones for Learning Hierarchical Embeddings.” ICML.
- [14] Muchane, R., et al. (2025). “Incorporating Hierarchical Semantics in Sparse Autoencoder Architectures.” ICML.
- [15] Matryoshka SAE Authors (2025). “Matryoshka Sparse Autoencoders.” Preprint.
- [16] Shi, J., et al. (2020). “Alzheimer’s Disease progression prediction using hyperbolic Ricci flow and sparse coding.” Medical Image Analysis.