Kiji Inspector — SAEs for tool selection in google/gemma-4-E4B-it
JumpReLU sparse autoencoders trained on the residual stream of
google/gemma-4-E4B-it at the
point where an agent commits to a tool, plus the feature labels, automated
validation, and causal ablation results that go with them.
What is in here
Five layers, each with a trained SAE and the full analysis chain:
| Path | Contents |
|---|
layer_N/sae_checkpoints/sae_final.pt | Trained JumpReLU SAE (d_model 2560 → d_sae 10240) |
layer_N/sae_checkpoints/feature_health.json | Alive/dead counts, L0 distribution, reconstruction MSE |
layer_N/sae_checkpoints/metrics.jsonl | Per-step training metrics |
layer_N/activations/contrastive_features.json | Per-contrast-type discriminative features with Cohen's d |
layer_N/activations/feature_descriptions.json | LLM-generated label + description per feature |
layer_N/activations/decision_report.json | Per-contrast-type natural-language explanation |
layer_N/activations/fuzzing_summary.json | Label-validation scores |
layer_N/activations/fuzzing_results.json | Per-example judgments |
layer_N/ablation/ablation_report.json | Causal intervention results |
Setup
| |
|---|
| Subject model | google/gemma-4-E4B-it (42 layers, d_model 2560, dense — no MoE) |
| Layers captured | 12, 18, 24, 30, 36 |
| Capture point | Residual stream entering block N, i.e. the output of block N−1 |
| Decision token | Assistant prefill "I'll use the", activation taken at the next-token position |
| Contrastive pairs | 1,574,523 pairs → 3,149,046 prompts, 37 contrast types |
| SAE | JumpReLU, d_sae 10240 (4× d_model), 10 epochs, adaptive L1 targeting L0 75 |
| Normalization | Per-dimension mean-centering, then scaling by RMS about the mean |
| Labeling + judging | Qwen/Qwen3.6-35B-A3B (never co-resident with the subject model) |
SAE health
| Layer | Depth | d_sae | Alive | Dead | L0 (median) | Firing rate | FVU |
|---|
| 12 | 29% | 10240 | 3,207 | 5.8% | 32.5 (23.9) | 1.01% | 0.181 |
| 18 | 43% | 10240 | 2,493 | 0.4% | 26.2 (25.1) | 1.05% | 0.188 |
| 24 | 57% | 10240 | 3,889 | 8.0% | 281.9 (37.5) | 7.25% | 0.157 |
| 30 | 71% | 10240 | 4,029 | 36.5% | 232.4 (219.1) | 5.77% | 0.243 |
| 36 | 86% | 10240 | 3,778 | 8.8% | 305.7 (157.9) | 8.09% | 0.144 |
Feature interpretation and validation
| Layer | Contrastive features | Labeled | Fuzzing score | Token acc | Excellent (>0.8) | Poor (<0.6) |
|---|
| 12 | 1,394 | 3,700 | 0.940 ± 0.002 | 0.943 | 90.2% | 3.3% |
| 18 | 1,822 | 2,494 | 0.922 ± 0.003 | 0.926 | 85.9% | 5.1% |
| 24 | 4,700 | 4,021 | 0.952 ± 0.002 | 0.951 | 92.6% | 2.9% |
| 30 | 6,890 | 4,087 | 0.963 ± 0.002 | 0.964 | 93.4% | 1.7% |
| 36 | 5,705 | 3,768 | 0.930 ± 0.002 | 0.930 | 87.1% | 4.5% |
Causal ablation
Top-10 contrastive features per contrast type are zeroed during the forward
pass; the measure is whether the predicted tool changes, against a
random-feature control and a reconstruction-only baseline.
Unweighted mean across contrast types, as stored in ablation_report.json:
| Layer | Contrastive flip | Random / reconstruction | Directed flip | CATE | Wilcoxon p<0.05 |
|---|
| 12 | 0.0236 | 0.0111 | 0.0043 | +0.0000 | 9/35 |
| 18 | 0.0214 | 0.0170 | 0.0093 | +0.0002 | 13/34 |
| 24 | 0.0259 | 0.0129 | 0.0046 | +0.0002 | 23/36 |
| 30 | 0.0935 | 0.0329 | 0.0363 | +0.0009 | 26/37 |
| 36 | 0.2552 | 0.0428 | 0.0837 | -0.0005 | 25/37 |
Those are means over contrast types, which weights a type with 3 tested prompts
the same as one with 100. Pooling over all tested prompts instead, with exact
(Clopper–Pearson) intervals on flip rates and a 2000-sample bootstrap on CATE:
| Layer | n tested | Contrastive flip [95% CI] | Random / recon flip [95% CI] | Directed flip [95% CI] | CATE [95% CI] |
|---|
| 12 | 1966 | 0.0183 [0.0129, 0.0253] | 0.0097 [0.0058, 0.0151] | 0.0036 [0.0014, 0.0073] | +0.00003 [-0.00007, +0.00014] |
| 18 | 1923 | 0.0198 [0.0140, 0.0270] | 0.0114 [0.0072, 0.0173] | 0.0094 [0.0056, 0.0148] | +0.00027 [+0.00014, +0.00045] |
| 24 | 2042 | 0.0269 [0.0204, 0.0349] | 0.0088 [0.0052, 0.0139] | 0.0044 [0.0020, 0.0084] | +0.00018 [-0.00000, +0.00040] |
| 30 | 2094 | 0.0855 [0.0739, 0.0983] | 0.0224 [0.0165, 0.0297] | 0.0325 [0.0253, 0.0410] | +0.00086 [+0.00034, +0.00153] |
| 36 | 2086 | 0.2704 [0.2514, 0.2900] | 0.0364 [0.0288, 0.0454] | 0.0786 [0.0674, 0.0910] | -0.00046 [-0.00078, -0.00020] |
Tool selection is decided late in the network.
- Layers 12, 18 and 24 — weak or no effect. The contrastive and random
flip-rate intervals overlap at layers 12 and 18, and the CATE interval
includes zero at layers 12 and 24. Layer 24's flip intervals are disjoint
(0.0204 > 0.0139) but its probability shift is not distinguishable from zero.
- Layers 30 and 36 — clear effect. Flip intervals are disjoint by a wide
margin (0.0739 > 0.0297 and 0.2514 > 0.0454). Layer 36 flips 27% of
decisions against a 3.6% baseline.
Directed flips — where the model switches to the specific paired alternative —
rise ~22× from layer 12 to layer 36 (0.0036 → 0.0786), the most consistent
gradient in the data.
Read the intervals, not the p-values. Layer 24 has 23/36 contrast types
at Wilcoxon p<0.05 while its pooled CATE interval touches zero. The test
detects a consistent sign among shifts too small to matter.
Note that layer 36's pooled CATE is negative ([−0.00078, −0.00020]) while
it flips the most decisions: ablation moves probability away from the paired
contrast tool, so most of its flips land on some third tool. Directed flips
(7.9%) are a small fraction of all flips (27%).
The two tables disagree in magnitude (e.g. layer 36, 0.2552 vs 0.2704) because
of the weighting; the pooled figures are the ones to quote.
How to read these results
The features are interpretable. Mean fuzzing score 0.941 across 18,070
labeled features. In this evaluation an independent judge sees only a feature's
label and a set of token spans, and must predict where the feature fires; it
succeeds ~94% of the time, with 86–93% of features scoring above 0.8. This is
the best-supported claim in the release.
The features are strongly diagnostic of tool choice. Contrastive analysis
finds 1,394–6,890 discriminating features per layer, with effect sizes up to
Cohen's d ≈ 48 on well-populated contrast types.
Causal necessity is confined to deep layers. Zeroing the 10 most
discriminative features changes the tool choice on 1.8–2.7% of prompts at
layers 12–24 — with control intervals overlapping at 12 and 18, i.e. no
demonstrable effect — but on 8.6% at layer 30 and 27.0% at layer 36, where
the intervals separate decisively.
A plausible explanation for the shallow-layer nulls is redundant encoding:
tool-selection information spread across far more than 10 features, with deeper
layers more committed and less redundant. This has not been tested — every
ablation here uses --n-features 10. A sweep over that parameter is the
decisive experiment, and it is cheap (no judge required).
Diagnostic strength and causal necessity are therefore not the same axis, and
they do not peak at the same layer: label quality peaks at layer 30 (0.963),
causal effect at layer 36 (27.0% flip rate).
Caveats
- The random control is degenerate.
random_flip_rate equals
reconstruction_flip_rate in 34–37 of 34–37 contrast types at every layer.
Random indices are drawn from all 10240 features, but only L0 of them are
active on any given prompt, so a random 10 touches an active feature just
3% of the time at layers 12/18 and 21–26% at layers 24/30/36. The
"random control" therefore measures SAE reconstruction error, not a matched
intervention. The layer-30/36 effect survives this qualitatively — its
interval clears the baseline by a wide margin — but the magnitude is
uncalibrated: it is measured against reconstruction noise rather than against
ablating 10 comparably-active features. A frequency-matched control is the
right fix before quoting effect sizes.
- Baseline pass rate is ~58%. About 42% of prompts do not predict the
dataset's own
anchor_tool before any intervention; all ablation metrics are
conditioned on those that do.
- Exclude
cost_vs_speed_optimization. Of its 66,203 generated pairs, only
3 have different tools on the two sides — the contrast is about tool
parameters, not tool choice. It yields degenerate statistics (Cohen's d up
to 604 from n=3) that sort to the top of their group. spot_buy_vs_contract
is similarly thin at ~2%.
- Layer 36's headline is not fully verified. Several contrast types flip
75–100% of prompts. That could be genuine strong causality or a structural
collapse in which removing a high-norm feature breaks the layer's computation
wholesale. Distinguishing them requires checking whether post-ablation
predictions are uniform, which the current reports do not record.
- L0 is bimodal at layers 24, 30 and 36 — report medians, not means. Layer
24's mean of 281.9 against a median of 37.5 describes neither population.
- Layer 30 has 36.5% dead features, by far the highest here.
- Ablation reliability varies by contrast type. A few types have high
reconstruction-only flip rates, where the SAE round-trip alone disturbs the
prediction and the intervention is uninformative.
reconstruction_baseline
is recorded per type so these can be filtered; a threshold around 10% is
reasonable.
- CATE is ≈0 everywhere: ablation perturbs decisions without cleanly
redirecting probability mass toward the paired tool.
Reproducing
1docker pull 575lab/kiji-inspector:dev
2
3python -m kiji_inspector.pipeline \
4 --subject-model google/gemma-4-E4B-it \
5 --judging-model Qwen/Qwen3.6-35B-A3B \
6 --layers 12 18 24 30 36 \
7 --target-l0 75 \
8 --step all
9
10python -m kiji_inspector.experiments.ablation \
11 --sae-checkpoint layer_36/sae_checkpoints/sae_final.pt \
12 --contrastive-features layer_36/activations/contrastive_features.json \
13 --pairs-dir output/pairs \
14 --model google/gemma-4-E4B-it \
15 --layer 36 --n-features 10
Note that step 5 caches per-token activations sized layers × features ×
examples; running all five layers in one pass needs several hundred GB. Run it
one layer at a time.
Environment: vLLM fork Davidnet/vllm @ b6455d43b, torch 2.11.0+cu129,
transformers 5.14.1, CUDA 12.9. Wall clock on a single RTX PRO 6000 Blackwell
(96 GB): ~3h12m extraction, ~28m SAE training, ~7m contrastive analysis, ~40m
labeling, ~64m fuzzing, ~30m ablation.
Loading an SAE
[!IMPORTANT]
These checkpoints are mean-centered: activations must be normalized as
(x - mean_vec) / rms_scale. The kiji-inspector release currently on PyPI
predates centering and applies x / rms_scale only. Using it on these
checkpoints is not a small error — measured on 20k real layer-30
activations, FVU degrades from 0.46 to 26.0 (56× worse, and 26× worse
than predicting the mean) and L0 inflates from 265 to 3,181. The offset
is ~168× the signal scale.
Use the snippet below, which reads both constants from the checkpoint and
works with any version, or install from
source where
SAE.normalize_input() does this for you.
1import torch
2from huggingface_hub import hf_hub_download
3
4REPO, LAYER = "575-lab/kiji-inspector-google-gemma-4-E4B-it", 30
5
6ckpt = torch.load(
7 hf_hub_download(REPO, f"layer_{LAYER}/sae_checkpoints/sae_final.pt"),
8 map_location="cpu", weights_only=False,
9)
10cfg = ckpt["config"]
11mean_vec = torch.tensor(cfg["mean_vec"]) # (d_model,)
12rms_scale = cfg["rms_scale"] # scalar
13
14# x: raw residual-stream activations, shape (..., 2560)
15x_norm = (x - mean_vec) / rms_scale # <- both terms are required
With a source install the same thing is one call:
1from kiji_inspector import SAE
2
3sae, descriptions = SAE.from_pretrained(
4 repo_id="575-lab/kiji-inspector-google-gemma-4-E4B-it", layer=30,
5)
6features = sae.encode(sae.normalize_input(x))
7reconstruction = sae.denormalize_output(sae.decode(features))
Verified on CPU against layer 30: d_model 2560, d_sae 10240, mean_vec
present, normalize/denormalize round trip accurate to 9e-6, post-normalization
variance 0.94, and measured L0 and MSE within 35% of the values reported in
feature_health.json (differences are sampling — health metrics use 1.6M
vectors, this check used 20k).
License
Apache 2.0. Copyright (c) 2026 Dataiku SAS.
