HyperSAE: High-Performance Hyperbolic Sparse Autoencoders

hypersae is a high-performance mechanistic interpretability engine designed to extract hierarchical concept ontologies from Large Language Models (LLMs). By decoupling hyperbolic geometry from the forward pass, it provides the zero-latency execution of standard Euclidean Sparse Autoencoders alongside the semantic mapping power of Riemannian negative curvature.

Installation

Install directly via PyPI:

pip install hypersae

Or install locally from source:

git clone https://github.com/vishal-dehurdle/hypersae.git
cd hypersae
pip install -e .

1. Core Architecture: Decoupled Weight-Space Regularization

To preserve high GPU throughput and model compatibility, hypersae separates execution into two computational speeds:

1. The Fast-Path (Euclidean Forward Pass): The massive volume of active token data remains entirely in flat, high-speed Euclidean space ($\mathbb{R}^d$). This avoids the latency of Riemannian manifolds, respects base model normalizations (e.g., RMSNorm), and maintains direct causal steering compatibility.
2. The Slow-Path (Hyperbolic Weight Regularization): The structural, hierarchical relationships of concepts are enforced exclusively in the dictionary parameter space during optimization via Poincaré ball projections $(\mathcal{B}^d, g_{\mathbf{x}})$.

graph TD
    subgraph Fast_Path["Fast-Path: Euclidean Forward Pass (bfloat16)"]
        X["Normalized Token Activations x"] --> ENC["Euclidean Encoder"]
        ENC --> F["Sparse Activations f"]
        F --> DEC["Euclidean Decoder (W_dec)"]
        DEC --> X_HAT["Reconstructed Activations x̂"]
    end

    subgraph Slow_Path["Slow-Path: Hyperbolic Weight Optimization (Upcast to float32)"]
        W_dec["Decoder Weights (W_dec)"] & R_depth["Depth Scalars (r_i)"] --> MAP["Poincaré Manifold Projection"]
        MAP --> H_coords["Hyperbolic Coordinates (h_i)"]
        H_coords --> MOCO["CoActivation Queue"]
        MOCO --> LOSS["Asymmetric Poincaré Entailment Loss"]
    end
    
    LOSS -.->|"Dual-Optimizer Update (AdamW / RiemannianAdam)"| W_dec

2. Empirical Benchmark Results (Gemma-2-2B Layer 13)

Evaluated at scale on Google Gemma-2-2B Layer 13 residual stream activations ($d=2304$, dict size $M=16384$) streaming over 20M tokens of FineWeb-Edu on an NVIDIA L4 GPU cluster:

Downstream Reasoning Retention (Single-Token Substitution)

Pareto Reconstruction & Sparsity Performance

Key Takeaway: HyperSAE achieves a 9.8% reduction in reconstruction MSE and a +3.4% boost in Cross-Entropy Loss Recovery over flat SAEs at matching sparsity ($L_0 \approx 53$).

3. Quickstart Example

import torch
from hypersae import HyperSAE, CoActivationQueue, TriPartiteLoss, HyperSAETrainer

device = "cuda" if torch.cuda.is_available() else "cpu"

# 1. Instantiate HyperSAE model, CoActivationQueue, and TriPartiteLoss
sae = HyperSAE(d_model=2304, dict_size=16384).to(device)
queue = CoActivationQueue(dict_size=16384).to(device)
loss_fn = TriPartiteLoss(l1_coeff=0.005, entail_coeff=0.01)

# 2. Instantiate HyperSAETrainer
trainer = HyperSAETrainer(model=sae, queue=queue, loss_fn=loss_fn, lr=1e-3)

# 3. Train step on residual stream activation batch
x = torch.randn(64, 2304, device=device)
metrics = trainer.train_step(x)

print(f"Total Loss: {metrics['loss_total']:.4f}")
print(f"Reconstruction MSE: {metrics['loss_recon']:.4f}")
print(f"Poincaré Entailment Penalty: {metrics['loss_entail']:.4f}")

4. Software Architecture

• hypersae.HyperSAE: Core model module implementing linear forward pass and learnable radial depths $r_i \in [0, 1)$.
• hypersae.FlatSAE: Standard Euclidean baseline for benchmark comparison.
• hypersae.TriPartiteLoss: Loss orchestrator combining MSE, $L_1$ sparsity, and Poincaré cone entailment penalties.
• hypersae.CoActivationQueue: Asynchronous GPU memory queue tracking feature co-occurrences without $\mathcal{O}(M^2)$ memory growth.
• hypersae.hooks: PyTorch and TransformerLens forward hook utilities for steering and intervention.

5. Research Papers & Publications

• Theoretical Paper: Escaping Flatland: Weight-Space Regularization and Hyperbolic Geometry in Mechanistic Interpretability
• Empirical Paper: Hyperbolic Sparse Autoencoders: Empirical Validation of Poincaré Manifold Geometry on LLM Activations

License

This project is licensed under the MIT License — see the LICENSE file for details.