kiji-inspector-google-gemma-4-E2B-it

kiji-inspector — SAE interpretability run: google/gemma-4-E2B-it

JumpReLU Sparse Autoencoders trained on residual-stream activations of google/gemma-4-E2B-it (35-layer decoder, d_model = 1536), to probe how the model selects tools across network depth.

• Subject model: google/gemma-4-E2B-it — 35 layers, d_model = 1536.
• Layer sweep: 13, 14, 15, 18, 19, 20, 23, 24, 25 (depths ≈ 37 % – 71 %, mid-to-deep).
• SAE width: d_sae = 16384 (≈ 10.7× expansion over d_model).
• Sparsity target: target_l0 = 75; 10 epochs; float32 activations.
• Generation / labelling / judge model: Qwen/Qwen3-VL-235B-A22B-Instruct-FP8.

SAE health

L0 vs target. No layer reached target_l0 = 75; L13 is closest (95). L14 and L15 are under-converged — L0 overshoots to 655 / 483 and reconstruction MSE blows up to 20.6 / 3.4 (vs ~0.05–0.12 elsewhere). Treat those two SAEs as not-yet-trained rather than as faithful decompositions. Deeper layers (18–25) reconstruct cleanly (MSE 0.05–0.12) but still carry L0 ≈ 115–194, i.e. 1.5–2.6× the sparsity target. Alive % is low (1–5 %) and dead % high (59–75 %), consistent with the aggressive 10.7× dictionary expansion.

Token-level fuzzing (autointerp)

Combined score = 0.7·token-accuracy + 0.3·prompt-accuracy. Quality tiers: Excellent > 0.8, Good 0.6–0.8, Poor < 0.6.

Best layer: L19 — highest combined (0.739) and token-acc (0.655), and the lowest poor rate (14.9 %). L18 has the most excellent features (34.4 %) but a higher poor rate (25 %). Fuzzing quality is fairly flat across the mid-to-deep sweep (combined 0.69–0.74); the two under-converged layers (14, 15) still fuzz reasonably because the evaluation scores feature descriptions, not reconstruction fidelity.

Per-layer artefacts

Each layer_<N>/ contains:

• sae_checkpoints/sae_final.pt — trained JumpReLU SAE weights.
• sae_checkpoints/config.json — training hyper-parameters.
• sae_checkpoints/feature_health.json — alive/dead, L0, reconstruction MSE (with 95 % CIs).
• sae_checkpoints/metrics.jsonl — per-step training log.
• activations/contrastive_features.json — Cohen's-d tool-selection features (post-hoc probe).
• activations/feature_descriptions.json — autointerp labels.
• activations/fuzzing_summary.json / fuzzing_results.json — token-level fuzzing eval.
• activations/decision_report.json — per-layer decision summary.
• activations/metadata.json — model, layer, d_model, token/shard counts.

Citation

If you use these artefacts, please cite the kiji-inspector framework (Opening the Black Box: Mechanistic Interpretability of Agent Tool Selection with Sparse Autoencoders).