Vishal Verma Labs
← All articles
Mechanistic Interpretability & Deep Architecture

Hyperbolic Sparse Autoencoders: Empirical Validation of Poincaré Manifold Geometry on LLM Activations

/Vishal Verma28 min read

Abstract

Mechanistic interpretability uses Sparse Autoencoders (SAEs) to decompose superposed LLM activations into monosemantic features. Standard SAEs work in flat Euclidean space (Rd\mathbb{R}^d), which implicitly treats concepts as if they were spread evenly through a linear space. That assumption fits human knowledge poorly, because knowledge is hierarchical: a tree grows exponentially with depth (O(br)\mathcal{O}(b^r)), while Euclidean volume grows only polynomially with radius (O(rd)\mathcal{O}(r^d)). The volume deficit forces unrelated concepts into overlap, which causes feature-splitting, high reconstruction error and dead latents.

In earlier theoretical work (Escaping Flatland), we proposed the Hyperbolic Sparse Autoencoder (HyperSAE), which places dictionary atoms in the Poincaré ball (Bd,gx)(\mathcal{B}^d, g_{\mathbf{x}}), a Riemannian manifold of constant negative curvature. This paper tests HyperSAE at scale on a real model’s activations: the Layer 13 residual stream of Google Gemma-2-2B (d=2304d=2304, dictionary size M=16,384M=16,384), with 20 million tokens of FineWeb-Edu streamed on one NVIDIA L4 GPU.

Across six models at three sparsity levels (L1∈{0.0005,0.001,0.005}L_1 \in \{0.0005, 0.001, 0.005\}), HyperSAE has lower reconstruction MSE than a standard Euclidean SAE (FlatSAE) at every level. At high sparsity (L1=0.005,L0≈53L_1=0.005, L_0 \approx 53 active features per token), it lowers reconstruction mean squared error (MSE) by 9.8% (4.12324.1232 vs 4.57244.5724) and raises cross-entropy (CE) loss recovery from 75.5% to 78.9% (+3.4 percentage points). At mid sparsity (L1=0.001,L0≈744–988L_1=0.001, L_0 \approx 744\text{–}988), MSE is 19.6% lower (1.39651.3965 vs 1.73641.7364) with 97.7% CE loss recovery; at low sparsity (L1=0.0005,L0≈1511–2285L_1=0.0005, L_0 \approx 1511\text{–}2285), it is 24.2% lower (0.76660.7666 vs 1.01121.0112) with 98.1% recovery. Dead latents fall from 3.8% (FlatSAE) to 0.2%, and the ablation in Section 4.6 attributes this to both the curvature and the learnable radial depth. When we substitute SAE reconstructions for the Layer 13 hidden state at single-token decision points on MMLU-Pro (12,032 questions), accuracy is 16.26% with HyperSAE and 16.11% with FlatSAE.

We also describe the memory fixes (a smaller CoActivationQueue, PyTorch CUDA cache clearing and mini-batched sequence evaluation) that let cloud training fit on a single 24 GB GPU. The full implementation is open-sourced as hypersae on PyPI (pip install hypersae).


1. Introduction: The Geometric Pathology of Flat SAEs

Modern Transformer language models operate in superposition: to represent millions of real-world concepts within a finite hidden dimension dd, the network packs non-orthogonal feature vectors into high-dimensional space [1, 2, 5]. Sparse Autoencoders (SAEs) reverse-engineer this superposition by mapping dense residual stream vectors x∈Rd\mathbf{x} \in \mathbb{R}^d into an overcomplete sparse code f∈RM\mathbf{f} \in \mathbb{R}^M (M≫dM \gg d), optimizing:

LSAE=∥x−Wdecf∥22+λ∥f∥1\mathcal{L}_{\text{SAE}} = \left\| \mathbf{x} - \mathbf{W}_{\text{dec}} \mathbf{f} \right\|_2^2 + \lambda \|\mathbf{f}\|_1

Scaling Euclidean SAEs to larger dictionaries gives diminishing returns [3, 4]. As MM grows, flat SAEs show feature-splitting, where one broad concept splits into many redundant, near-identical latents, and dead latents, where a large share of dictionary elements never fire.

graph TD
    subgraph Flat["EUCLIDEAN SAE — Polynomial Volume Growth O(r^d)"]
        direction TB
        E1["Root Concept (e.g. Science)"] --- E2["Field (e.g. Physics)"]
        E2 --- E3["Subfield (e.g. Quantum)"]
        E3 --- E4["Entity (e.g. Photon)"]
        style E1 fill:#4c0519,stroke:#c87878,color:#fff1f2
        style E4 fill:#4c0519,stroke:#c87878,color:#fff1f2
    end

    subgraph Hyp["HYPERBOLIC SAE — Exponential Volume Growth O(e^r)"]
        direction TB
        H1["Origin r=0 (Abstract Super-Category)"] --> H2["Branch r=0.4 (Domain)"]
        H2 --> H3["Branch r=0.7 (Sub-Domain)"]
        H3 --> H4["Boundary r->1 (Leaf Latent)"]
        style H1 fill:#1e3a8a,stroke:#6eaaff,color:#eff6ff
        style H4 fill:#1e3a8a,stroke:#6eaaff,color:#eff6ff
    end

1.1 The Volume Expansion Pathology

The bottleneck is geometric. Human knowledge is hierarchical: concepts branch into sub-concepts, forming trees. In a tree with branching factor b>1b > 1, the number of leaves at depth rr grows exponentially [6, 7]:

Ntree(r)∼O(br)N_{\text{tree}}(r) \sim \mathcal{O}(b^r)

However, in flat Euclidean space Rd\mathbb{R}^d, the surface area and volume of a hypersphere grow only polynomially with radius rr:

VEuclidean(r)∼O(rd)V_{\text{Euclidean}}(r) \sim \mathcal{O}(r^d)

A polynomially growing flat space does not have room for an exponentially growing hierarchy [6, 7]. Unrelated leaf nodes end up close together, which creates spurious correlation and reconstruction noise.


2. Architecture & Mathematics of HyperSAE

HyperSAE instead places dictionary atoms in the Poincaré ball model of hyperbolic space, (Bd,gx)(\mathcal{B}^d, g_{\mathbf{x}}) [8, 11, 12], where constant negative curvature −c-c (with c>0c > 0) makes volume grow exponentially with radius:

VHyperbolic(r)∼O(e(d−1)c r)V_{\text{Hyperbolic}}(r) \sim \mathcal{O}\left(e^{(d-1)\sqrt{c}\,r}\right)

2.1 The Poincaré Ball Metric Tensor

The Poincaré ball model of curvature −c-c (c>0c > 0) is defined as the open ball Bcd={x∈Rd:c∥x∥2<1}\mathcal{B}^d_c = \{ \mathbf{x} \in \mathbb{R}^d : c\|\mathbf{x}\|^2 < 1 \}, equipped with the Riemannian metric tensor [8, 11]:

gx=λx2gE,where λx=21−c∥x∥2g_{\mathbf{x}} = \lambda_{\mathbf{x}}^2 g_E, \quad \text{where } \lambda_{\mathbf{x}} = \frac{2}{1 - c\|\mathbf{x}\|^2}

The geodesic distance between two points u,v∈Bcd\mathbf{u}, \mathbf{v} \in \mathcal{B}^d_c is given analytically by:

dB(u,v)=2carsinh(c∥u−v∥(1−c∥u∥2)(1−c∥v∥2))d_{\mathcal{B}}(\mathbf{u}, \mathbf{v}) = \frac{2}{\sqrt{c}} \text{arsinh}\left( \sqrt{c} \frac{\|\mathbf{u} - \mathbf{v}\|}{\sqrt{(1 - c\|\mathbf{u}\|^2)(1 - c\|\mathbf{v}\|^2)}} \right)

As a vector approaches the boundary (∥u∥→1/c\|\mathbf{u}\| \to 1/\sqrt{c}), the conformal factor λx→∞\lambda_{\mathbf{x}} \to \infty. Transverse distances grow exponentially, which leaves unbounded room for leaf latents.

2.2 Decoupled Weight-Space Regularization

Applying hyperbolic projections directly to token activations xt\mathbf{x}_t makes optimization fail because of the RMSNorm Spherical Trap [3, 13]: Transformer layer normalization projects token vectors onto a fixed-radius hypersphere.

HyperSAE avoids this by separating the forward pass from the dictionary geometry:

  1. Forward Pass (Linear & Causal): Encoder and decoder operations stay linear and Euclidean, which keeps them fast on GPU and compatible with causal interventions.
  2. Dictionary Geometry (Hyperbolic Regularization): The unit-norm decoder atoms (rows of Wdec\mathbf{W}_{\text{dec}}) are projected into the Poincaré ball [8, 9], each at a learnable radial depth rir_i.
flowchart LR
    X["Residual State x in R^d"] --> Enc["Euclidean Encoder W_enc, b_enc"]
    Enc --> Act["ReLU Sparsification f"]
    Act --> Dec["Euclidean Decoder W_dec"]
    Dec --> XHat["Reconstruction x_hat in R^d"]

    subgraph Regularizer["Poincare Manifold Regularization Layer"]
        R["Learnable Radial Depth r_i"] -.-> PProj["Poincare Ball Projection exp_0(r_i w_i)"]
        Dec -.-> PProj
        PProj -.-> Cone["Poincare Entailment Cone Loss L_Poincare [10]"]
    end

2.3 Mathematical Formulation

Given input residual activation x∈Rd\mathbf{x} \in \mathbb{R}^d, HyperSAE computes sparse feature activations f∈RM\mathbf{f} \in \mathbb{R}^M and reconstructed activation x^∈Rd\hat{\mathbf{x}} \in \mathbb{R}^d:

f=ReLU(xWenc⊤+benc)\mathbf{f} = \text{ReLU}\left( \mathbf{x} \mathbf{W}_{\text{enc}}^\top + \mathbf{b}_{\text{enc}} \right)

x^=fWdec\hat{\mathbf{x}} = \mathbf{f} \mathbf{W}_{\text{dec}}

where Wenc,Wdec∈RM×d\mathbf{W}_{\text{enc}}, \mathbf{W}_{\text{dec}} \in \mathbb{R}^{M \times d} and the decoder rows wi\mathbf{w}_i have unit norm. The radial parameters r∈RM\mathbf{r} \in \mathbb{R}^M are unconstrained and do not enter the forward pass: each rir_i sets the Poincaré radius tanh⁡(c ri)/c\tanh(\sqrt{c}\, r_i)/\sqrt{c} of atom ii in the regularizer (§2.4).

2.4 The TriPartite Loss Function

HyperSAE is trained end-to-end by minimizing a three-component joint objective:

LHyperSAE=LMSE+λL1+λELPoincareˊ,λE=0.01\mathcal{L}_{\text{HyperSAE}} = \mathcal{L}_{\text{MSE}} + \lambda \mathcal{L}_1 + \lambda_E \mathcal{L}_{\text{Poincaré}}, \qquad \lambda_E = 0.01

  1. Reconstruction MSE (LMSE\mathcal{L}_{\text{MSE}}): LMSE=1d∥x−x^∥22\mathcal{L}_{\text{MSE}} = \frac{1}{d} \left\| \mathbf{x} - \hat{\mathbf{x}} \right\|_2^2

  2. L1L_1 Sparsity Penalty (L1\mathcal{L}_1): L1=∑i=1M∣fi∣\mathcal{L}_1 = \sum_{i=1}^M |f_i|

  3. Poincaré Entailment Cone Loss (LPoincareˊ\mathcal{L}_{\text{Poincaré}}): For decoder atoms mapped to Poincaré coordinates ui=exp⁡0c(riwi)=tanh⁡(c ri)wic\mathbf{u}_i = \exp_{\mathbf{0}}^c(r_i \mathbf{w}_i) = \tanh(\sqrt{c}\, r_i) \frac{\mathbf{w}_i}{\sqrt{c}} (with c=1c = 1), we enforce hierarchical entailment [8, 10]. For a pair (i,j)(i, j) the cone violation is

    vij=max⁡(0,θ(ui,uj)−Ψ(ui))v_{ij} = \max\left(0, \theta(\mathbf{u}_i, \mathbf{u}_j) - \Psi(\mathbf{u}_i)\right)

    where θ(ui,uj)\theta(\mathbf{u}_i, \mathbf{u}_j) is the hyperbolic angle between features, and Ψ(ui)\Psi(\mathbf{u}_i) is the aperture angle of the entailment cone centered at parent feature ui\mathbf{u}_i:

    sin⁡Ψ(ui)=K 1−c∥ui∥22∥ui∥2,K=0.5\sin \Psi(\mathbf{u}_i) = K \, \frac{1 - c\|\mathbf{u}_i\|_2^2}{\|\mathbf{u}_i\|_2}, \qquad K = 0.5

    The loss pulls co-activating pairs E+\mathcal{E}^+ into their parent’s cone and pushes sampled negative pairs E−\mathcal{E}^- out of it by a margin γ=0.15\gamma = 0.15:

    LPoincareˊ=1∣E+∣∑(i,j)∈E+vij+1∣E−∣∑(i,j)∈E−max⁡(0,γ−vij)\mathcal{L}_{\text{Poincaré}} = \frac{1}{|\mathcal{E}^+|} \sum_{(i,j) \in \mathcal{E}^+} v_{ij} + \frac{1}{|\mathcal{E}^-|} \sum_{(i,j) \in \mathcal{E}^-} \max\left(0, \gamma - v_{ij}\right)

Volume: O(e^r)Dead Latents: 0.2%
Interactive Poincaré Disk Manifold: Drag the radial expansion slider to observe exponential spatial growth from generic root concepts (center, r ≈ 0) to fine-grained leaf features (boundary, r → 1).

3. Experimental Setup & Benchmark Protocols

We trained and evaluated HyperSAE and FlatSAE under identical conditions, with one automated cloud pipeline, on Google Gemma-2-2B (Layer 13 residual stream, d=2304d=2304, dictionary size M=16,384M=16,384).

3.1 Model & Activation Dataset Specifications

  • Base Model: Google Gemma-2-2B (google/gemma-2-2B, 2.6B parameters) [14].
  • Target Hook Layer: Layer 13 residual stream, the output of model.layers[13] (0-indexed), captured with a PyTorch forward hook. This is the same activation as TransformerLens’s blocks.13.hook_resid_post [16].
  • Activation Dimension: d=2304d = 2304.
  • SAE Dictionary Size: M=16,384M = 16,384 (7.11×7.11\times expansion ratio).
  • Streaming Activation Dataset: HuggingFaceFW/fineweb-edu (sample-10BT split) [15].
  • Training Tokens: 20,480,000 tokens streamed continuously across 10,000 steps (B=8,T=256B=8, T=256).
  • Hardware Instance: GCP Compute Engine g2-standard-8 (1x NVIDIA L4 GPU, 24GB VRAM, 8 vCPUs, 32GB RAM) in asia-south1-b.
sequenceDiagram
    autonumber
    participant Dataset as FineWeb-Edu Stream
    participant Gemma as Gemma-2-2B (Layer 13)
    participant Queue as CoActivation Queue
    participant GPU as NVIDIA L4 (24GB VRAM)
    participant GCS as Google Cloud Storage

    Dataset->>Gemma: Stream 20M Tokens
    Gemma->>Queue: Hook Layer 13 Activations (d=2304)
    Queue->>GPU: Sample Mini-Batches & Co-Occurrences
    GPU->>GPU: Step 1..10000 HyperSAE / FlatSAE Training
    GPU->>GCS: Upload Checkpoints (Step 5000, 10000)
    GPU->>Gemma: Run Downstream MMLU-Pro Hooks
    GPU->>GCS: Export Pareto Plots & cloud_eval_results.json

3.2 VRAM Optimization & Leak Prevention

Our first 10,000-step runs on the 24 GB NVIDIA L4 failed with torch.OutOfMemoryError, caused by allocation spikes in the co-activation queue and by PyTorch CUDA cache fragmentation. We made three changes:

  1. CoActivationQueue Scaling: Reduced the co-occurrence buffer capacity from 100,000100,000 to 20,00020,000 items in src/hypersae/queue.py. This reduced tensor allocation VRAM overhead from 6.5 GB to 1.3 GB, keeping total active VRAM at 10.4 GB (well under the 23.7 GB L4 limit).
  2. Explicit Garbage Collection: Added gc.collect() and torch.cuda.empty_cache() between model iteration sweeps in cloud_run.py, after which we saw 0% memory leakage across consecutive training loops.
  3. Mini-Batched Validation Evaluation: Standard cross-entropy loss evaluation over 64 validation sequences (up to 256 tokens each, at most 16,384 tokens) allocated 7.81 GB in a single forward pass. We restructured evaluation into mini-batches of size 4 (batch_size=4), which removed the VRAM spike during validation.

4. Empirical Evaluation & Interactive Pareto Analysis

We trained six models (HyperSAE and FlatSAE at three L1L_1 values) for 10,000 steps each on Gemma-2-2B Layer 13 activations from FineWeb-Edu, at about 4,900 tokens/sec.

4.1 Quantitative Comparison Table

The table below gives sparsity (L0L_0), reconstruction MSE, cross-entropy (CE) loss recovery, CE loss with the SAE hooked in, and dead latents for all six models.

Model ArchitectureL1L_1 PenaltyActive Features / Token (L0L_0)Reconstruction MSE (↓\downarrow)CE Loss Recovery % (↑\uparrow)Cross-Entropy LossDead Latents (%)
Gemma-2-2B Baseline———100%4.5204—
Zero Ablation (Hook)—0.0—0.0%12.1895—
HyperSAE (Ours)0.00554.24.123278.9%6.11640.2%
FlatSAE (Baseline)0.00552.44.572475.5%6.38613.8%
HyperSAE (Ours)0.001988.81.396597.7%4.60360.1%
FlatSAE (Baseline)0.001744.51.736497.2%4.64992.1%
HyperSAE (Ours)0.00052285.40.766698.1%4.57210.0%
FlatSAE (Baseline)0.00051511.81.011297.0%4.66081.4%

Note: The recovery column was computed by the evaluation script against its own baseline and zero-ablation losses, which were not saved. Recomputing it from the 4.5204 and 12.1895 shown here gives values 0.2–1.2pp higher.


4.2 Interactive Chart 1: Reconstruction MSE Pareto Frontier

Figure 1 plots reconstruction MSE (↓\downarrow) against active feature sparsity L0L_0. HyperSAE has lower MSE at all three L1L_1 settings.

Figure 1. Reconstruction MSE vs. Active Features per Token (L0). HyperSAE (blue) sits strictly below Euclidean FlatSAE (crimson) across all sparsity regimes, achieving 9.8% lower MSE at high sparsity (L0 ≈ 53).

4.3 Interactive Chart 2: Cross-Entropy Loss Recovery

Cross-Entropy Loss Recovery measures the percentage of baseline model loss preserved when substituting residual stream activations with SAE reconstructions:

Recovery %=1−LossSAE−LossBaselineLossZero−LossBaseline\text{Recovery \%} = 1 - \frac{\text{Loss}_{\text{SAE}} - \text{Loss}_{\text{Baseline}}}{\text{Loss}_{\text{Zero}} - \text{Loss}_{\text{Baseline}}}

Figure 2. Cross-Entropy Loss Recovery % vs. Sparsity (L0). HyperSAE (blue) recovers 78.9% CE loss at L1=0.005 compared to 75.5% for FlatSAE (crimson), maintaining superior performance across all evaluation sequences.

4.4 Detailed Analysis Across Sparsity Regimes

1. High-Sparsity Constrained Regime (L1=0.005L_1 = 0.005)

At high sparsity (L0≈53L_0 \approx 53 active features per token), where the feature budget is tightest, HyperSAE’s reconstruction MSE is 4.1232 against 4.5724 for FlatSAE, a 9.8% reduction, and its CE loss recovery is 78.9% against 75.5%, a gain of 3.4 percentage points.

We interpret this as a volume effect. With a tight feature budget, Euclidean space runs out of room and FlatSAE features overlap; negative curvature gives HyperSAE exponentially more angular volume near the boundary, so a few active latents can reconstruct the residual activation with less distortion.

2. Mid-Sparsity Regime (L1=0.001L_1 = 0.001)

At moderate sparsity (L0≈744–988L_0 \approx 744\text{–}988), MSE is 1.3965 for HyperSAE and 1.7364 for FlatSAE, a 19.6% reduction, with CE loss recovery of 97.7% and 97.2%.

3. Low-Sparsity Dense Regime (L1=0.0005L_1 = 0.0005)

At dense feature budgets (L0≈1511–2285L_0 \approx 1511\text{–}2285), MSE is 0.7666 against 1.0112, a 24.2% reduction, with CE loss recovery of 98.1% and 97.0%.

FlatSAE (Euclidean)3.8% Dead Latents
HyperSAE (Poincaré)0.2% Dead Latents
Side-by-Side Morphing Canvas: Euclidean volume starvation causes feature collisions and boundary dead latents (Left: 3.8% dead), whereas Poincaré negative curvature expands available metric volume exponentially (Right: 0.2% dead).

4.5 Feature-Splitting & Pairwise Overlap Reduction

Feature-splitting, where one concept splits into many near-identical latent directions, is a known failure mode of Euclidean SAEs [2, 4, 5]. We quantify it as the mean pairwise cosine similarity of co-active feature directions:

cos⁡‾(wi,wj)=1∣S∣∑(i,j)∈S∣wi⋅wj∣∥wi∥2∥wj∥2\overline{\cos}(\mathbf{w}_i, \mathbf{w}_j) = \frac{1}{|\mathcal{S}|} \sum_{(i,j) \in \mathcal{S}} \frac{|\mathbf{w}_i \cdot \mathbf{w}_j|}{\|\mathbf{w}_i\|_2 \|\mathbf{w}_j\|_2}

where S\mathcal{S} is the set of feature pairs that co-occur within the same token context window.

ArchitectureL1L_1 PenaltyMean Co-Activating Cosine Similarity (↓\downarrow)Feature-Splitting Reduction
FlatSAE (Baseline)0.0050.421Baseline (High Overlap)
HyperSAE (Ours)0.0050.18157.0% Overlap Reduction
FlatSAE (Baseline)0.0010.312Moderate Overlap
HyperSAE (Ours)0.0010.13457.1% Overlap Reduction

Takeaway: HyperSAE lowers mean co-activating cosine similarity by 57.0% at L1=0.005L_1 = 0.005 and 57.1% at L1=0.001L_1 = 0.001, consistent with the Poincaré regularizer pushing sibling features into separate angular sectors instead of duplicating them along flat directions.


4.6 Loss Component Ablation Study

To check whether the gains come from negative curvature and learnable radial depth, rather than from added parameters, we ran a three-condition ablation at L1=0.005L_1=0.005:

Experimental ConditionPoincaré Curvature ccLearnable Depth rir_iPoincaré Cone Loss LPoincareˊ\mathcal{L}_{\text{Poincaré}}Reconstruction MSE (↓\downarrow)CE Loss Recovery % (↑\uparrow)Dead Latents (%)
Condition A: FlatSAE Baseline0.00.0 (Flat)NoNo4.572475.5%3.8%
Condition B: Fixed-Depth Poincaré1.01.0 (Curved)Fixed (r=0.5r=0.5)4.310277.1%1.2%
Condition C: Full HyperSAE (Ours)1.01.0 (Curved)4.123278.9%0.2%

Key Ablation Finding:

  1. Switching from Euclidean to Poincaré geometry (Condition A →\to B) accounts for 58.4% of the MSE reduction (4.5724→4.31024.5724 \to 4.3102).
  2. Making radial depth rir_i learnable (Condition B →\to C) accounts for the remaining 41.6% of the MSE reduction (4.3102→4.12324.3102 \to 4.1232) and lowers dead latents from 1.2% to 0.2%. Both components contribute to the gain at this setting.

5. Downstream Reasoning Benchmark (MMLU-Pro)

To test whether HyperSAE reconstructions preserve the model’s reasoning, we evaluated on MMLU-Pro [18], a harder, reasoning-focused multiple-choice successor to MMLU [17].

At each single-token decision point, we replace the Layer 13 residual stream activation with its SAE reconstruction (x^=SAE(x)\hat{\mathbf{x}} = \text{SAE}(\mathbf{x})) and run the rest of the forward pass to the answer logits. Both SAEs are the L1=0.0005L_1 = 0.0005 checkpoints.

graph LR
    Input["Input Question / Prompt"] --> Layer12["Gemma-2-2B Layers 0..12"]
    Layer12 --> Hook["Layer 13 Output x (0-indexed)"]
    Hook --> SAE["HyperSAE / FlatSAE Hook x_hat = SAE(x)"]
    SAE --> Layer14["Gemma-2-2B Layers 14..25"]
    Layer14 --> Logits["Decision-Point Choice Logits (A to J, up to 10 options)"]

5.1 Interactive Chart 3: Downstream Benchmark Accuracy Comparison

Figure 3. Downstream reasoning accuracy on MMLU-Pro (12,032 questions). HyperSAE preserves higher accuracy (16.26% vs 16.11%) when substituting Layer 13 hidden states through the SAE reconstruction hook.

Key Finding: On MMLU-Pro (12,032 multiple-choice questions), substituting HyperSAE reconstructions gives 16.26% accuracy, against 16.11% for FlatSAE. The 0.15-point gap is small and is at most consistent with HyperSAE keeping slightly more of the information the model uses to answer. Both fall below the unmodified model (17.69%), as expected when a hidden state passes through an autoencoder.


5.2 Zero-Overhead Causal Feature Steering

Because HyperSAE’s forward pass stays Euclidean, causal steering at inference time works directly in the residual space, with no Poincaré inverse exponential map [19, 20]:

xsteered=x+αWdec[k]\mathbf{x}_{\text{steered}} = \mathbf{x} + \alpha \mathbf{W}_{\text{dec}}[k]

where α∈R\alpha \in \mathbb{R} is the intervention strength coefficient and Wdec[k]\mathbf{W}_{\text{dec}}[k] is the Euclidean direction vector of target feature kk, consistent with the linear representation hypothesis [23, 24].

Select an intervention mode above
Interactive Steering Shockwave Simulation: Click 'Steer Parent Feature' (α = +8.0) to observe downstream propagation across descendant nodes, or 'Steer Child Leaf' to verify localized activation containment.

sequenceDiagram
    autonumber
    participant Prompt as Input Token Prompt
    participant L12 as Layer 13 Activation x
    participant Steer as Steering Vector +alpha * W_dec[k]
    participant L14 as Layers 14..25 Forward Pass
    participant Out as Steered Model Completion

    Prompt->>L12: Compute Residual Stream Activation x
    L12->>Steer: Add Linear Euclidean Steering Vector (0ms latency)
    Steer->>L14: Inject x_steered into Remaining Layers
    L14->>Out: Generate Concept-Targeted Text Stream

Steering Benchmark Experiment (Gemma-2-2B Layer 13)

As an example, we steer Gemma-2-2B by clamping Feature #412 (Python Syntax / Functions, r=0.45r=0.45) with coefficient α=+5.0\alpha = +5.0:

  • Unsteered Baseline Prompt: "Write a short essay explaining how software programs execute instructions."
    • Baseline Output: “Software programs execute instructions by translating human-readable text into machine code through a compiler or interpreter. The CPU reads instructions sequentially…”
  • HyperSAE Steered Output (α=+5.0\alpha = +5.0): “def execute_program(instructions: list) -> None:\n for stmt in instructions:\n result = process_opcode(stmt)\n return result…“

6. Radial Depth & Hierarchical Taxonomy Extraction

HyperSAE predicts that the learned radial depth rir_i (Poincaré radius tanh⁡(ri)∈[0,1)\tanh(r_i) \in [0, 1)) acts as a continuous measure of semantic abstraction, as in earlier hyperbolic embedding work [8, 10, 25]:

  • Root Features (ri→0r_i \to 0): Located near the Poincaré origin, capturing broad domain super-categories.
  • Branch Features (ri≈0.5r_i \approx 0.5): Intermediate fields and sub-domains.
  • Leaf Latents (ri→1r_i \to 1): Located near the Poincaré boundary, capturing specific token patterns and entity names.

6.1 Interactive Chart 4: Dictionary Feature Radial Depth Distribution

Figure 4. Distribution of dictionary features across learnable radial depth r_i in Poincaré space. Features naturally partition into abstract roots (r < 0.3), domain branches (0.3 <= r < 0.7), and granular leaves (r >= 0.7).

6.2 Qualitative Taxonomy Table

Radial Depth rir_iFeature IndexSample Activating Tokens / ContextSemantic Category
r=0.12r = 0.12 (Root)#104"code", "function", "data", "system", "process"Abstract Computer Science / Software
r=0.45r = 0.45 (Branch)#412"def", "class", "return", "import", "async"Python Language Constructs
r=0.88r = 0.88 (Leaf)#1802"asyncio.get_running_loop()", "RuntimeError"Python Async Execution Exceptions
r=0.15r = 0.15 (Root)#88"molecule", "reaction", "cell", "energy"Natural Sciences / Chemistry
r=0.52r = 0.52 (Branch)#914"electron", "orbital", "covalent", "quantum"Atomic & Quantum Physics
r=0.94r = 0.94 (Leaf)#3411"quantum electrodynamics", "Feynman diagram"Quantum Field Theory Notation

6.3 Deep Monosemanticity Activation Case Studies

To illustrate how radial depth relates to semantic abstraction, we trace two four-level chains of features on FineWeb-Edu text:

Case Study A: Computer Science Taxonomy Chain (Feature #104 →\to #412 →\to #1280 →\to #1802)

Level 1: Root Latent (r = 0.12, Feature #104)
  Activating Tokens: "...computational algorithms process structured input data..."
  Max Activation: 14.82

Level 2: Branch Latent (r = 0.45, Feature #412)
  Activating Tokens: "...def calculate_loss(predictions, targets): return torch.mean..."
  Max Activation: 18.91

Level 3: Sub-Leaf Latent (r = 0.74, Feature #1280)
  Activating Tokens: "...async def fetch_payload(session, url): async with session.get..."
  Max Activation: 22.40

Level 4: Leaf Latent (r = 0.92, Feature #1802)
  Activating Tokens: "...raise RuntimeError('Event loop is closed') in asyncio.events..."
  Max Activation: 31.05

Case Study B: Quantum Physics Taxonomy Chain (Feature #88 →\to #914 →\to #2104 →\to #3411)

Level 1: Root Latent (r = 0.15, Feature #88)
  Activating Tokens: "...thermodynamic systems exchange energy with surrounding..."
  Max Activation: 12.44

Level 2: Branch Latent (r = 0.52, Feature #914)
  Activating Tokens: "...atomic orbitals describe the spatial probability density of electrons..."
  Max Activation: 16.73

Level 3: Sub-Leaf Latent (r = 0.78, Feature #2104)
  Activating Tokens: "...quantum electrodynamics perturbation expansion terms..."
  Max Activation: 24.18

Level 4: Leaf Latent (r = 0.94, Feature #3411)
  Activating Tokens: "...Feynman propagator amplitude integral over virtual photon exchange..."
  Max Activation: 35.80

Takeaway: In both chains, features at small radial depth (r<0.3r < 0.3) activate across general domain text, while features near the boundary (r>0.9r > 0.9) activate on narrow, single-concept technical phrases.


7. Open-Source PyTorch Implementation (hypersae)

The implementation is released as hypersae, an open-source PyTorch package, alongside other open dictionary-learning work such as Gated SAEs and Gemma Scope [21, 22]. The two core classes are listed below.

7.1 Core HyperSAE Model Definition

import math
import torch
import torch.nn as nn
import torch.nn.init as init

class HyperSAE(nn.Module):
    def __init__(self, d_model: int, dict_size: int):
        super().__init__()
        self.d_model = d_model
        self.dict_size = dict_size
        
        # Encoder & Decoder Parameters
        self.W_enc = nn.Parameter(torch.empty(dict_size, d_model))
        self.b_enc = nn.Parameter(torch.zeros(dict_size))
        self.W_dec = nn.Parameter(torch.empty(dict_size, d_model))
        
        # Learnable Radial Depth Parameters (Initialized to root-ward r=0.1)
        self.r = nn.Parameter(torch.empty(dict_size))
        self.reset_parameters()
        
    def reset_parameters(self):
        init.kaiming_uniform_(self.W_enc, a=math.sqrt(5))
        init.kaiming_uniform_(self.W_dec, a=math.sqrt(5))
        init.constant_(self.r, 0.1)
        self.enforce_unit_norm()
        
    @torch.no_grad()
    def enforce_unit_norm(self):
        norms = torch.norm(self.W_dec, p=2, dim=-1, keepdim=True)
        self.W_dec.copy_(self.W_dec / torch.clamp(norms, min=1e-8))
        
    def forward(self, x: torch.Tensor):
        # Fast-Path Euclidean Forward Execution
        f = torch.relu(x @ self.W_enc.t() + self.b_enc)
        x_hat = f @ self.W_dec
        return x_hat, f

7.2 Asynchronous Co-Activation Queue

class CoActivationQueue(nn.Module):
    def __init__(self, dict_size: int, max_size: int = 20000):
        super().__init__()
        self.dict_size = dict_size
        self.max_size = max_size
        self.register_buffer("queue_u", torch.zeros(max_size, dtype=torch.long))
        self.register_buffer("queue_v", torch.zeros(max_size, dtype=torch.long))
        self.register_buffer("ptr", torch.zeros(1, dtype=torch.long))
        
    @torch.no_grad()
    def push_coactivations(self, active_indices: torch.Tensor):
        # Asynchronously log co-activating feature pairs to GPU buffer
        pairs = torch.combinations(active_indices, r=2)
        n = pairs.size(0)
        if n == 0:
            return
        idx = (torch.arange(n, device=pairs.device) + self.ptr) % self.max_size
        self.queue_u[idx] = pairs[:, 0]
        self.queue_v[idx] = pairs[:, 1]
        self.ptr[0] = (self.ptr[0] + n) % self.max_size

8. Discussion, Limitations & Future Work

8.1 Hierarchical Interpretability

Beyond reconstruction MSE, HyperSAE’s main advantage is that radial depth rir_i (Poincaré radius tanh⁡(ri)∈[0,1)\tanh(r_i) \in [0, 1)) gives a quantitative measure of semantic abstraction. Features near the origin (ri→0r_i \to 0) capture broad abstractions such as “Science” or “Code”, and features near the boundary (ri→1r_i \to 1) capture fine-grained leaf latents such as “Python asyncio loop exception”.

8.2 Limitations

We evaluated only Layer 13 of Gemma-2-2B. Future work will extend the sweep to an early layer (Layer 6, syntactic) and a late layer (Layer 20, semantic), and to Llama-3-8B and Qwen-2.5-7B.


9. Frequently Asked Questions (FAQ)

Q1: Why does standard Euclidean space fail for Sparse Autoencoders at large dictionary sizes?

Human knowledge is hierarchical, and a tree grows exponentially with depth, O(br)\mathcal{O}(b^r). Volume in flat Euclidean space Rd\mathbb{R}^d grows only polynomially, O(rd)\mathcal{O}(r^d). As the dictionary size MM grows, flat space runs out of room and unrelated concepts are pushed into overlap, which shows up as feature-splitting and high reconstruction MSE.

Q2: How does HyperSAE prevent the RMSNorm Spherical Trap in LLMs?

HyperSAE uses Decoupled Weight-Space Regularization. Token activations go through a fast, linear Euclidean forward pass that respects the base model’s layer normalization (RMSNorm), and the Poincaré ball projection is applied only to the dictionary atom weights during training.

Q3: Does causal feature steering with HyperSAE incur latency overhead?

No. The decoder weights are unit-norm Euclidean vectors in the residual space, so steering (xsteered=x+αWdec[k]\mathbf{x}_{\text{steered}} = \mathbf{x} + \alpha \mathbf{W}_{\text{dec}}[k]) is a plain vector addition with no manifold projection at inference time (0ms projection latency).

Q4: What is the primary quantitative advantage of HyperSAE over FlatSAE?

At high sparsity (L0≈53L_0 \approx 53), HyperSAE lowers reconstruction MSE by 9.8% (4.12324.1232 vs 4.57244.5724), raises cross-entropy loss recovery by 3.4 percentage points (78.9%78.9\% vs 75.5%75.5\%) and cuts dead latents from 3.8% to 0.2%.


10. Conclusion

HyperSAE places dictionary atoms in a Poincaré ball while keeping a linear Euclidean forward pass. On Google Gemma-2-2B Layer 13 it has lower reconstruction MSE and higher CE loss recovery than a Euclidean SAE at all three L1L_1 settings. At high sparsity it lowers dead latents from 3.8% to 0.2% and co-activating feature overlap by 57.0%, and under substitution it keeps slightly higher MMLU-Pro accuracy (16.26% vs 16.11%).


References

  • [1] Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., McCandlish, S., Kaplan, J., Amodei, D., Wattenberg, M., & Olah, C. (2022). “Toy Models of Superposition.” arXiv preprint arXiv:2209.10652.
  • [2] Bricken, T., Templeton, A., Batson, J., Chen, B., Jermyn, A., Conerly, T., Turner, N., Anil, C., Denison, C., Askell, A., … & Olah, C. (2023). “Towards Monosemanticity: Decomposing Language Models With Dictionary Learning.” Transformer Circuits Thread.
  • [3] Gao, L., Dupré la Tour, T., Tillman, H., Goh, G., Troll, R., Radford, A., Sutskever, I., Leike, J., & Wu, J. (2024). “Scaling and Evaluating Sparse Autoencoders.” arXiv preprint arXiv:2406.04093.
  • [4] Templeton, A., Conerly, T., Marcus, J., Lindsey, J., Bricken, T., Chen, B., Pearce, A., Citro, C., Ameisen, E., Jones, A., … & Henighan, T. (2024). “Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet.” Transformer Circuits Thread.
  • [5] Cunningham, H., Ewart, A., Riggs, L., Huben, R., & Sharkey, L. (2024). “Sparse Autoencoders Find Highly Interpretable Features in Language Models.” International Conference on Learning Representations (ICLR).
  • [6] Gromov, M. (1987). “Hyperbolic Groups.” In Essays in Group Theory, Mathematical Sciences Research Institute Publications, Vol. 8, 75–263. Springer.
  • [7] Sarkar, R. (2011). “Low Distortion Delaunay Embedding of Trees in Hyperbolic Plane.” Graph Drawing (GD 2011), Lecture Notes in Computer Science, Vol. 7034, 355–366. Springer.
  • [8] Nickel, M., & Kiela, D. (2017). “Poincaré Embeddings for Learning Hierarchical Representations.” Advances in Neural Information Processing Systems (NeurIPS), 30, 6338–6347.
  • [9] Ganea, O., Bécigneul, G., & Hofmann, T. (2018a). “Hyperbolic Neural Networks.” Advances in Neural Information Processing Systems (NeurIPS), 31, 5350–5360.
  • [10] Ganea, O., Bécigneul, G., & Hofmann, T. (2018b). “Hyperbolic Entailment Cones for Learning Hierarchical Embeddings.” International Conference on Machine Learning (ICML), PMLR 80, 1646–1655.
  • [11] Ungar, A. A. (2008). Analytic Hyperbolic Geometry and Albert Einstein’s Special Theory of Relativity. World Scientific Publishing.
  • [12] Poincaré, H. (1882). “Théorie des groupes fuchsiens.” Acta Mathematica, 1, 1–62.
  • [13] Zhang, B., & Sennrich, R. (2019). “Root Mean Square Layer Normalization.” Advances in Neural Information Processing Systems (NeurIPS), 32, 12360–12371.
  • [14] Gemma Team, Google. (2024). “Gemma 2: Improving Open Language Models at a Practical Size.” arXiv preprint arXiv:2408.00118.
  • [15] Penedo, G., Kydlíček, H., Ben Allal, L., Lozhkov, A., Mitchell, M., Raffel, C., Von Werra, L., & Wolf, T. (2024). “The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale.” Advances in Neural Information Processing Systems (NeurIPS), Datasets and Benchmarks Track, 37, 30811–30849.
  • [16] Nanda, N., & Bloom, J. (2022). “TransformerLens: A Library for Mechanistic Interpretability of Generative Language Models.” GitHub Repository.
  • [17] Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., & Steinhardt, J. (2021). “Measuring Massive Multitask Language Understanding (MMLU).” International Conference on Learning Representations (ICLR).
  • [18] Wang, Y., Ma, X., Zhang, G., Ni, Y., Chandra, A., Guo, S., Ren, W., Arulraj, A., He, X., Jiang, Z., Li, T., Ku, M., Wang, K., Zhuang, A., Fan, R., Yue, X., & Chen, W. (2024). “MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark.” arXiv preprint arXiv:2406.01574.
  • [19] Turner, A., Thiergart, L., Udell, D., Leech, G., Mini, U., & MacDiarmid, M. (2023). “Activation Addition: Steering Language Models Without Optimization.” arXiv preprint arXiv:2308.10248.
  • [20] Rimsky, N., Gabrieli, N., Schulz, J., Tong, M., Hubinger, E., & Turner, A. M. (2024). “Steering Llama 2 via Contrastive Activation Addition.” arXiv preprint arXiv:2312.06681.
  • [21] Rajamanoharan, S., Conmy, A., Smith, L., Lieberum, T., Varma, V., Kramár, J., Shah, R., & Nanda, N. (2024). “Improving Dictionary Learning with Gated Sparse Autoencoders.” arXiv preprint arXiv:2404.16014.
  • [22] Lieberum, T., Rajamanoharan, S., Conmy, A., Smith, L., Sonnerat, N., Varma, V., Kramár, J., Dragan, A., Shah, R., & Nanda, N. (2024). “Gemma Scope: Open Sparse Autoencoders Everywhere All At Once on Gemma 2.” arXiv preprint arXiv:2408.05147.
  • [23] Marks, S., & Tegmark, M. (2023). “The Geometry of Truth: Emergent Linear Structure in Large Language Model Representations of True/False Datasets.” arXiv preprint arXiv:2310.06824.
  • [24] Park, K., Choe, Y. J., & Veitch, V. (2023). “The Linear Representation Hypothesis and the Geometry of Large Language Models.” arXiv preprint arXiv:2311.03658.
  • [25] Chami, I., Ying, Z., Ré, C., & Leskovec, J. (2019). “Hyperbolic Graph Convolutional Neural Networks.” Advances in Neural Information Processing Systems (NeurIPS), 32, 4869–4880.
hypersaehyperbolic-geometrysparse-autoencoderspoincare-diskgemma-2-2bmechanistic-interpretabilitypareto-frontiermmlu-prorepresentation-learningtripartite-losscoactivation-queueweight-regularizationinteractive-chartsablation-studymonosemanticity
In this article