juiceb0xc0de/gemma-4-e4b-it-atlas-SAE

juiceb0xc0de/gemma-4-e4b-it-atlas-sae

This joins the SAE half of the E4B atlas to the census half. Built from `juiceb0xc0de/gemma-4-e4b-it-SAE-v2` (42 layers, 32x, d_sae 81,920) against `juiceb0xc0de/gemma-4-e4b-it-atlas`.

Three splits, all queryable in the browser. No download needed to look around.

The join

The atlas and the SAEs describe the same model in two index spaces that could not talk to each other:

atlas features.feature_idx   in [0, 10240)   an MLP hidden unit
SAE   feature_idx            in [0, 81920)   a dictionary element

Both write into the residual stream at d_model = 2560. Finding 8 in the atlas card says the census side of it: a hidden unit writes through one column of down_proj, so its output direction is a single vector. The SAE side is W_dec[:, f]. Cosine between the two is the join.

Before this, merge_sae.py finished by telling you to steer with W_dec[:, feat], a vector the database did not contain. The 70 GB of sae.pt sat outside the atlas entirely.

What is in here

42 files per split, one per layer, l00 through l41. About 1.18 GB total. Per-layer provenance (sae_id, EV, L0, dead %, trainer commit, training job id, gate policy) is in sae_bridge_manifest.json.

channel_idx runs 0 to 10,239 against an intermediate size of 10,240, so all 3,413,504 bridge rows resolve to a real census channel. The 106,672 rows in vectors match the 106,672 materialized rows in geometry exactly.

Query it on the Hub

Hit the SQL Console button. The splits are DuckDB views named after themselves.

The row-preview pane above the console will show a cast error. That is expected: the three tables have different column sets, and the viewer tries to cast every split in a config down to one schema. The console reads the Parquet directly and does not care. Keeping all three in one config is deliberate, because it is what lets you join geometry to bridge_mlp in a single query.

Peak persona alignment per layer:

SELECT layer_id, round(max(abs(persona_cos_all)), 3) AS peak
FROM geometry
GROUP BY layer_id
ORDER BY layer_id;

The features worth steering with, and where each one lands in the census:

SELECT layer_id, feature_idx, persona_cos_all, top_mlp_channel, top_mlp_cos
FROM geometry
WHERE materialized
ORDER BY abs(persona_cos_all) DESC
LIMIT 50;

The tightest decoder-to-channel matches in the model:

SELECT layer_id, feature_idx AS sae_feat, channel_idx AS mlp_chan, cos
FROM bridge_mlp
WHERE "rank" = 0
ORDER BY abs(cos) DESC
LIMIT 50;

Quote "rank". DuckDB reads a bare rank as the window function and throws don't know what type: at you, which is not a helpful error message.

Query it locally, against the census atlas

The point of the bridge is joining back to atlas.sqlite, and the Hub console cannot see that file. Vectors round-trip through Parquet as variable FLOAT[], and DuckDB's array_* functions bind only to fixed-size ARRAY, so cast once in a view:

INSTALL sqlite; LOAD sqlite;
ATTACH 'atlas.sqlite' AS atlas (TYPE sqlite, READ_ONLY);

CREATE VIEW geom   AS SELECT * FROM 'data/geometry_l*.parquet';
CREATE VIEW bridge AS SELECT * FROM 'data/bridge_mlp_l*.parquet';
CREATE VIEW vecs   AS
  SELECT sae_id, layer_id, feature_idx, CAST(w_dec AS FLOAT[2560]) AS w_dec
  FROM 'data/vectors_l*.parquet';

Score and steering vector in one row:

SELECT g.layer_id, g.feature_idx, g.persona_cos_all, v.w_dec
FROM geom g JOIN vecs v USING (sae_id, layer_id, feature_idx)
WHERE g.materialized
ORDER BY abs(g.persona_cos_all) DESC LIMIT 50;

The join that did not exist before:

SELECT b.layer_id, b.feature_idx AS sae_feat, b.channel_idx AS mlp_chan,
       b.cos, f.taxonomy_class, f.activation_rate
FROM bridge b
JOIN atlas.features f
  ON f.layer_id = b.layer_id AND f.component = 'mlp'
 AND f.feature_idx = b.channel_idx
WHERE b."rank" = 0
ORDER BY abs(b.cos) DESC;

Per finding 8 the same hidden unit backs mlp, gate, and up, so one bridge row serves all three component views.

What the geometry shows

max |cos(W_dec[:,f], persona_direction[layer])| climbs with depth:

The handoff reports E2B topping out at 0.28 to 0.43 across tested layers. E4B's late layers go well past that. Single decoder features carry substantially more of the Bella direction than the E2B result suggested. The dip at layers 18 to 20 (0.216 to 0.325) sits inside the middle-depth EV trough the handoff flagged.

The bridge tightens with depth too. Mean |cos| for the best-matching MLP channel per feature:

Early decoder columns are spread across many channels. Late ones increasingly sit on one. The single tightest match in the model is at layer 24, |cos| 0.865.

What is gated, and one thing wrong with how

Vectors are materialized for the top 2,048 features per layer by |persona_cos|, unioned with the top 512 by dec_norm. That lands between 2,520 and 2,552 features per layer after overlap. All 81,920 features keep their scores and a top_mlp_channel pointer either way, so a missing vector means "not materialized", never "dead". The materialized column says which. Storing every vector would be 17.6 GB. 2,048 per layer is about 10x the largest k the handoff's P2 experiment calls for.

The dec_norm half of that policy did nothing useful. Decoder columns are unit-normalized, so across all 3,440,640 rows dec_norm takes exactly 7 distinct float32 values, running from 0.99999976 to 1.00000024. That is a spread of 4.8e-07, which is float32 rounding and not signal. 2,187,021 of those rows sit at literal 1.0. So "top 512 by dec_norm" sorted a column that is constant to within rounding, and the extra features per layer arrived on whatever order the sort happened to produce rather than on norm. They are still real features with real scores. They were just not picked for the reason the policy states. If you want a principled second criterion, choose one and rebuild the gate. enc_norm does carry a little more range (0.998085 to 1.0) if you need a tiebreak that is at least not noise.

Checked against the causal table

The atlas card's finding 2 reports that the static composite score separated causally-kept from causally-dropped directions by 0.000075, which is no signal, while the causal test flipped sign. This build reproduces that exactly: 0.483327 against 0.483252. A new static score deserves the same suspicion, so here is persona_cos against axis_causal, all 5,160 causally tested directions:

Causal keep-rate and persona alignment both rise into the back half, which reproduces finding 3 ("separation front-loaded, editability back-loaded") from a completely independent measurement.

That is correlational and layer-level, not per-direction. axis_causal scores singular vectors of projections, not SAE features, so the two only join by layer. Four aggregate points. It is a consistent picture, not evidence that persona_cos predicts causal effect. Getting that needs a causal test run directly on SAE features. Do not use persona_cos as a targeting signal until that test exists.

Still open

The SAE hook point is assumed to pair layer L's decoder with layer L's down_proj (bridge_offset = 0, recorded in the manifest). If the trainer captured resid_pre rather than post-block, the correct pairing is L-1 and the bridge needs rebuilding with --bridge-offset -1. Nothing here settles that.

Build

hf jobs run --flavor l4x1 --timeout 2h --secrets HF_TOKEN \
  -v hf://datasets/juiceb0xc0de/gemma-4-e4b-it-SAE-v2:/saes:ro \
  -v hf://models/google/gemma-4-E4B-it:/model:ro \
  -v hf://buckets/juiceb0xc0de/atlas-runs:/out \
  pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime \
  /bin/sh -c '... build_sae_bridge_e4b.py --layers 0-41 --gate 2048 --stage-model'

About 9 s per layer on one L4 plus 134 s staging the base model, so 42 layers in about 12 minutes. The 70 GB of SAE weights never leaves the datacenter.

Scripts in this repo: build_sae_bridge_e4b.py, verify_sae_bridge.py, freeze_axis_subspace.py, build_axis_surgery_models.py. Also sae_atlas_views.sql in JuiceB0xC0de/GWIQ-atlas.

License

Gemma license, matching the source model.

More

• —Model: https://huggingface.co/google/gemma-4-E4B-it
• —Census atlas: https://huggingface.co/datasets/juiceb0xc0de/gemma-4-e4b-it-atlas
• —SAEs: https://huggingface.co/datasets/juiceb0xc0de/gemma-4-e4b-it-SAE-v2
• —Follow: https://huggingface.co/juiceb0xc0de