Skip to main content

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

Reason this release was yanked:

Superseded by 0.1.0a2 (security fix). Use 0.1.0a2 or later.

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.0},
}

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.0a0.tar.gz (63.5 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.0a0-py3-none-any.whl (26.5 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: hybridlens-0.1.0a0.tar.gz
  • Upload date:
  • Size: 63.5 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.0a0.tar.gz
Algorithm Hash digest
SHA256 527fed81285407be646aeecedc63f62b1433c4238520395974d0f1ea76b229eb
MD5 b133144d8203d957bb2963712f51a316
BLAKE2b-256 3482f560cf12762ec3a97535a54accb90df4763bad55cdcf48f45bfd66b4ee24

See more details on using hashes here.

File details

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

File metadata

  • Download URL: hybridlens-0.1.0a0-py3-none-any.whl
  • Upload date:
  • Size: 26.5 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.0a0-py3-none-any.whl
Algorithm Hash digest
SHA256 5549db8f59911c4e069fe9cf08c89b37e306a86d1270708762ced43f0fc952f6
MD5 afd96e664bda0b47be3ec1035013a45e
BLAKE2b-256 b7df0c02de630ebe805497b25d2f1704d61898c152da07fdcc3f04e7c3be8a65

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