Yusser/Soofi-S-Instruct-Preview-SAEs

JumpReLU SAEs for Soofi-S-Instruct-Preview

Sparse autoencoders trained on Soofi-Project/Soofi-S-Instruct-Preview — the Soofi Project's multilingual (EN/DE/ES/FR/IT) instruction-tuned reasoning model, built on the hybrid Mamba-2 + GQA-Attention + MoE Nemotron-H architecture (Nemotron 3 Nano 30B-A3B class). The base model is gated (closed beta) — request access on its model page to run the extraction examples below. Companion release: JumpReLU SAEs for NVIDIA Nemotron 3 Nano 30B-A3B.

Trained with the Gemma Scope 2 JumpReLU recipe (quadratic L0 penalty, rectangular-kernel STE, unit-norm decoder) by the nemotron-sae pipeline.

The activation-normalization factor and pre-encoder bias are folded into the weights, so every SAE here is a standard JumpReLU on raw model activations and loads with stock SAELens — no custom code:

from sae_lens import SAE

sae = SAE.from_pretrained("Yusser/Soofi-S-Instruct-Preview-SAEs", "L2_resid_post/w16384_l0_10", device="cuda")
feats = sae.encode(acts)      # acts: raw activations at sae.cfg.metadata.hook_name
recon = sae.decode(feats)

Equivalently with nemotron-lens (pip install nemotron-lens), which also hooks the model itself — activation capture and SAE splicing on this hybrid architecture, which TransformerLens does not support:

from nemotron_lens import HookedNemotron, SAE

sae = SAE.from_pretrained("Yusser/Soofi-S-Instruct-Preview-SAEs", "L2_resid_post/w16384_l0_10", device="cuda")
model = HookedNemotron.from_pretrained("Soofi-Project/Soofi-S-Instruct-Preview")
logits, cache = model.run_with_cache("The Eiffel Tower is in",
                                    names=[sae.metadata["hook_name"]])
feats = sae.encode(cache[sae.metadata["hook_name"]].float().flatten(0, 1))

Getting the activations needs only HuggingFace transformers (the hook point is a real module path, stored in the SAE's metadata):

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Soofi-Project/Soofi-S-Instruct-Preview"
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True,
                                             torch_dtype="bfloat16", device_map="auto")
tok = AutoTokenizer.from_pretrained(model_id)

acts = {}
mod = model.get_submodule(sae.cfg.metadata.hook_name)
h = mod.register_forward_hook(lambda m, i, o: acts.__setitem__('x', o[0] if isinstance(o, tuple) else o))
model(**tok("The Eiffel tower is in", return_tensors="pt").to(model.device))
h.remove()
feats = sae.encode(acts['x'].float().flatten(0, 1))

Available SAEs

All are width 16384, trained 1,000,000 steps (batch 4,096 -> ~4.1B token-activations) on nvidia/Nemotron-CC-v2.1 (High-Quality), seq_len 2048, L0 target 10.

sae_id	layer	component	hook (module path)	L0	FVU	dead %
L2_resid_post/w16384_l0_10	2	resid_post	backbone.layers.2	9.9	0.401	12.0
L11_mamba_out/w16384_l0_10	11	mamba_out	backbone.layers.11.mixer	10.0	0.284	84.1
L12_attn_out_prelinear/w16384_l0_10	12	attn_out_prelinear	backbone.layers.12.mixer	9.8	0.410	89.6
L12_resid_post/w16384_l0_10	12	resid_post	backbone.layers.12	10.1	0.382	58.7
L13_moe_out/w16384_l0_10	13	moe_out	backbone.layers.13.mixer	10.0	0.471	30.6
L25_mamba_out/w16384_l0_10	25	mamba_out	backbone.layers.25.mixer	9.9	0.525	84.6
L26_attn_out_prelinear/w16384_l0_10	26	attn_out_prelinear	backbone.layers.26.mixer	10.0	0.445	84.3
L26_resid_post/w16384_l0_10	26	resid_post	backbone.layers.26	10.2	0.574	86.2
L27_moe_out/w16384_l0_10	27	moe_out	backbone.layers.27.mixer	10.1	0.532	59.8
L41_mamba_out/w16384_l0_10	41	mamba_out	backbone.layers.41.mixer	9.9	0.447	89.9
L42_attn_out_prelinear/w16384_l0_10	42	attn_out_prelinear	backbone.layers.42.mixer	9.8	0.292	86.2
L42_resid_post/w16384_l0_10	42	resid_post	backbone.layers.42	10.2	0.604	93.4
L43_moe_out/w16384_l0_10	43	moe_out	backbone.layers.43.mixer	9.9	0.482	23.4
L51_resid_post/w16384_l0_10	51	resid_post	backbone.layers.51	10.0	0.317	75.0

Components: resid_post = residual stream after the block (hook = block module output); mamba_out / attn_out_prelinear / moe_out = mixer-module outputs (Mamba-2, attention pre-out-projection, MoE block).

Notes & caveats

• L0* = 10 is the sparsest anchor of a planned sparsity frontier: expect modest FVU and (on non-residual sites) high dead-latent fractions at this operating point.
• Metrics above were computed by the in-house eval (65k tokens, training distribution); eval fields are also embedded in each SAE's cfg.json metadata.
• The base model requires trust_remote_code=True, mamba_ssm and causal-conv1d.
• Training code, recipe provenance, and full telemetry: nemotron-sae.

Citation

If you use these SAEs, please cite the nemotron-sae repository (https://github.com/Yusser96/nemotron-sae) and Gemma Scope (arXiv:2408.05147), whose recipe this follows.