Skip to main content

Mechanistic interpretability for hybrid Mamba-Attention LLMs (Jamba, Zamba2): cross-component SAEs, ablation attribution, and a Hub registry.

Project description

hybridlens

Alpha License: Apache-2.0 PyPI

Mechanistic interpretability for hybrid Mamba-Attention LLMsJamba, Zamba2 — with cross-component SAEs, ablation attribution, and a Hub registry.

hybridlens is the recurrentlens sibling for mixed-architecture models, where attention layers, SSM layers, and MoE routers coexist in a single forward pass. Pure-transformer tools (SAELens, OpenMOSS) and pure-SSM tools (recurrentlens) both assume a single component family; hybrid models need per-component analysis.

Why hybrid-specific tooling

2026's mainline open-weight LLMs are no longer pure transformers: AI21's Jamba 1.5, Zyphra's Zamba2, NVIDIA's Nemotron Nano 2, and Microsoft's Samba all ship Mamba+Attention+MoE in a single stack. The existing mechanistic interpretability toolchain (SAELens, OpenMOSS, circuit-tracer) was written against decoder-only transformers and silently mis-attributes features when an SSM layer or MoE router is present. hybridlens fills that gap with per-component SAEs trained on the actual hidden stream each layer carries, plus attribution that respects the component boundary.

Solo-maintained OSS. Best-effort response. No production SLA. Released 2026 by @hinanohart under Apache-2.0.

v0.1.0 scope

Eight public APIs covering Jamba + Zamba2:

  • load_hybrid_model — load and tag attention / SSM / MoE components
  • HybridLensModel — wrapper that composes (does not subclass) the recurrentlens primitives (BaseSAE, HookManager, ActivationCache)
  • LayerLayout — pydantic schema for per-layer component map
  • train_component_sae — train an SAE on one component's hidden stream
  • ComponentSAEBundle — save/load a per-component SAE collection
  • attribute(method='ablation') — zero-ablation attribution to target metric
  • steer_component — context-managed activation steering
  • hub_push_bundle / hub_load_bundle — HF Hub I/O for bundles

Deferred to v0.1.1+: captum integrated-gradients attribution, MoE expert routing trace, Sankey cross-component flow viz, Nemotron / Samba adapters, pretrained Hub bundles, HF Space deploy.

v0.1.0a0 alpha — what runs vs what raises

The alpha freezes the eight public signatures and ships the bookkeeping half of the library — anything that does not need a live forward pass through a hybrid LLM. Concretely:

API v0.1.0a0 status
LayerLayout / LayerSlot runs
JambaAdapter.discover_layout / Zamba2Adapter.discover_layout runs (config-driven)
hooks.resolve_site runs
ComponentSAEBundle (in-memory + save / load round-trip) runs
viz.render_hybrid_explorer runs
hub_push_bundle / hub_load_bundle runs (lazy HfApi import)
load_hybrid_model NotImplementedError until v0.1.0 final
train_component_sae NotImplementedError until v0.1.0 final
attribute(method='ablation') NotImplementedError until v0.1.0 final
steer_component (scale ≠ 0 branch) NotImplementedError until v0.1.0 final
hooks.install_capture_hook NotImplementedError until v0.1.0 final

The five remaining APIs require torch + transformers + recurrentlens to be installed in the same environment; they land in v0.1.0 (final) after the Phase 5 release-verify venv exercises them end-to-end. The signatures are frozen — you can write code against them today.

Install

pip install hybridlens
# Optional Mamba kernels (CUDA only):
pip install "hybridlens[mamba]"

Quick look

import hybridlens as hl

model = hl.load_hybrid_model("ai21labs/Jamba-tiny-dev", family="jamba")
bundle = hl.train_component_sae(
    model,
    component="attn",
    layer=4,
    n_tokens=200_000,
)
print(bundle.summary())

⚠️ HF Hub anonymous rate limit (100 req/h) — set HF_TOKEN for steady-state usage.

Known limitations (v0.1.0)

  • Only Jamba and Zamba2 adapters. Nemotron-Nano / Samba arrive in v0.1.1.
  • Attribution is zero-ablation only; captum IG is v0.1.1.
  • MoE expert tracing and cross-component Sankey viz are v0.2.
  • Default training dataset is Pile-uncopyrighted; users supplying their own corpora are responsible for license compliance.
  • Bundle Hub assets are user-uploaded; no pretrained release ships with v0.1.0.

Citation

@software{hinanohart_hybridlens_2026,
  author  = {hinanohart},
  title   = {hybridlens: cross-component SAE attribution for hybrid LLMs},
  year    = {2026},
  url     = {https://github.com/hinanohart/hybridlens},
  version = {0.1.0a0},
}

Related

  • recurrentlens — sibling project for pure-SSM models. hybridlens reuses its BaseSAE, HookManager, ActivationCache, and hub.io primitives.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

hybridlens-0.1.0a1.tar.gz (70.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

hybridlens-0.1.0a1-py3-none-any.whl (27.7 kB view details)

Uploaded Python 3

File details

Details for the file hybridlens-0.1.0a1.tar.gz.

File metadata

  • Download URL: hybridlens-0.1.0a1.tar.gz
  • Upload date:
  • Size: 70.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.15

File hashes

Hashes for hybridlens-0.1.0a1.tar.gz
Algorithm Hash digest
SHA256 b43106857996498bf3c756aef932a0449161f4bb2877ee2212de93a9f3cc093f
MD5 954979c1813d22f6fab6360a2bb0e437
BLAKE2b-256 be5f78d4df18dded851be3d1f27038a97dbe2cce7478c269f237ac00cb9ee9cb

See more details on using hashes here.

File details

Details for the file hybridlens-0.1.0a1-py3-none-any.whl.

File metadata

  • Download URL: hybridlens-0.1.0a1-py3-none-any.whl
  • Upload date:
  • Size: 27.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.15

File hashes

Hashes for hybridlens-0.1.0a1-py3-none-any.whl
Algorithm Hash digest
SHA256 e3dbee5cc2e0a769a7965b8cb6f9515a152f5a5aaac4f158fb9dad3d0eebb579
MD5 ce54c7215afaa7966c641e6762a0452a
BLAKE2b-256 e2321583c751c63c529a4b1e828725e0f63db3bf46b39a4d538db60fb99cebe7

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page