kiji-inspector-google-gemma-4-E4B-it

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.

Produced with dataiku/kiji-inspector.

What is in here

Five layers, each with a trained SAE and the full analysis chain:

Setup

SAE health

Feature interpretation and validation

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:

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:

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

docker pull 575lab/kiji-inspector:dev

python -m kiji_inspector.pipeline \
  --subject-model google/gemma-4-E4B-it \
  --judging-model Qwen/Qwen3.6-35B-A3B \
  --layers 12 18 24 30 36 \
  --target-l0 75 \
  --step all

python -m kiji_inspector.experiments.ablation \
  --sae-checkpoint layer_36/sae_checkpoints/sae_final.pt \
  --contrastive-features layer_36/activations/contrastive_features.json \
  --pairs-dir output/pairs \
  --model google/gemma-4-E4B-it \
  --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.

import torch
from huggingface_hub import hf_hub_download

REPO, LAYER = "575-lab/kiji-inspector-google-gemma-4-E4B-it", 30

ckpt = torch.load(
    hf_hub_download(REPO, f"layer_{LAYER}/sae_checkpoints/sae_final.pt"),
    map_location="cpu", weights_only=False,
)
cfg = ckpt["config"]
mean_vec = torch.tensor(cfg["mean_vec"])   # (d_model,)
rms_scale = cfg["rms_scale"]               # scalar

# x: raw residual-stream activations, shape (..., 2560)
x_norm = (x - mean_vec) / rms_scale        # <- both terms are required

With a source install the same thing is one call:

from kiji_inspector import SAE

sae, descriptions = SAE.from_pretrained(
    repo_id="575-lab/kiji-inspector-google-gemma-4-E4B-it", layer=30,
)
features = sae.encode(sae.normalize_input(x))
reconstruction = 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.