Yusser/nemotron-3-nano-30b-a3b-saes

JumpReLU SAEs for Nemotron 3 Nano 30B-A3B

Sparse autoencoders trained on nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16, NVIDIA's hybrid Mamba-2 + GQA-Attention + MoE language model — to our knowledge the first public SAEs for a hybrid-architecture LLM. 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/nemotron-3-nano-30b-a3b-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)

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 = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
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 1M steps (batch 4096 -> ~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	10.1	0.462	7.6
L11_mamba_out/w16384_l0_10	11	mamba_out	backbone.layers.11.mixer	10.1	0.333	80.7
L12_attn_out_prelinear/w16384_l0_10	12	attn_out_prelinear	backbone.layers.12.mixer	10.0	0.446	82.3
L12_resid_post/w16384_l0_10	12	resid_post	backbone.layers.12	10.1	0.422	51.7
L13_moe_out/w16384_l0_10	13	moe_out	backbone.layers.13.mixer	10.0	0.541	18.4
L25_mamba_out/w16384_l0_10	25	mamba_out	backbone.layers.25.mixer	10.1	0.578	85.7
L26_attn_out_prelinear/w16384_l0_10	26	attn_out_prelinear	backbone.layers.26.mixer	10.0	0.464	83.0
L26_resid_post/w16384_l0_10	26	resid_post	backbone.layers.26	10.1	0.596	75.0
L27_moe_out/w16384_l0_10	27	moe_out	backbone.layers.27.mixer	10.0	0.609	59.9
L41_mamba_out/w16384_l0_10	41	mamba_out	backbone.layers.41.mixer	10.0	0.515	91.3
L42_attn_out_prelinear/w16384_l0_10	42	attn_out_prelinear	backbone.layers.42.mixer	9.9	0.319	82.4
L42_resid_post/w16384_l0_10	42	resid_post	backbone.layers.42	9.9	0.646	92.1
L43_moe_out/w16384_l0_10	43	moe_out	backbone.layers.43.mixer	10.0	0.599	19.7
L51_resid_post/w16384_l0_10	51	resid_post	backbone.layers.51	10.0	0.470	89.2

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.